#!/usr/bin/env python3
"""四角 LiDAR 外部真值定位离线回放和落盘自检入口。"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

from four_lidar_ball_algorithm import (
    FourLidarBallLocalizationAlgorithm,
    FourLidarBallLocalizationInput,
    FourLidarBallLocalizationOutput,
    LidarFrame,
)
from workshop_external_lidar_common import (
    ExternalLocalizationRecorder,
    build_diagnostics_payload,
    build_frame_status,
    load_config,
    now_us,
    parse_ball_targets,
    parse_lidar_mounts,
    validate_config,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="离线验证四角 LiDAR 外部真值定位配置、算法入口、诊断输出和数据落盘。"
    )
    parser.add_argument("--config", required=True, help="四角 LiDAR + 标定球配置文件。")
    parser.add_argument(
        "--frames-jsonl",
        default="",
        help="可选回放数据，每行一个 JSON；不提供时使用空点云合成帧做链路自检。",
    )
    parser.add_argument(
        "--session-dir",
        default="",
        help="覆盖 recording.session_dir，并启用落盘。",
    )
    parser.add_argument(
        "--dataset-index",
        default="",
        help="覆盖 recording.dataset_index_path，并启用 dataset_index 写入。",
    )
    parser.add_argument("--session-id", default="", help="覆盖 recording.session_id。")
    parser.add_argument("--site-id", default="", help="覆盖 recording.site_id。")
    parser.add_argument("--vehicle-id", default="", help="覆盖 recording.vehicle_id。")
    parser.add_argument(
        "--disable-recording",
        action="store_true",
        help="关闭落盘，只检查配置和算法调用链路。",
    )
    parser.add_argument(
        "--require-valid-pose",
        action="store_true",
        help="要求至少输出一帧 pose_valid=true；当前算法模板未实现时不要开启。",
    )
    return parser.parse_args()


def apply_overrides(config: dict[str, Any], args: argparse.Namespace) -> None:
    recording = config.setdefault("recording", {})
    if args.session_dir:
        recording["session_dir"] = args.session_dir
        recording["enabled"] = True
        if not args.dataset_index:
            recording["dataset_index_path"] = "dataset_index.yaml"
    if args.dataset_index:
        recording["dataset_index_path"] = args.dataset_index
        recording["enabled"] = True
    if args.session_id:
        recording["session_id"] = args.session_id
    if args.site_id:
        recording["site_id"] = args.site_id
    if args.vehicle_id:
        recording["vehicle_id"] = args.vehicle_id
    if args.disable_recording:
        recording["enabled"] = False


def points_from_value(value: Any) -> list[tuple[float, float, float]]:
    result: list[tuple[float, float, float]] = []
    if not isinstance(value, list):
        return result
    for item in value:
        if isinstance(item, dict):
            result.append((
                float(item.get("x", item.get("x_m", 0.0)) or 0.0),
                float(item.get("y", item.get("y_m", 0.0)) or 0.0),
                float(item.get("z", item.get("z_m", 0.0)) or 0.0),
            ))
        elif isinstance(item, (list, tuple)) and len(item) >= 2:
            z = item[2] if len(item) >= 3 else 0.0
            result.append((float(item[0]), float(item[1]), float(z)))
    return result


def frame_from_dict(raw: dict[str, Any], fallback_timestamp_us: int) -> LidarFrame:
    lidar_id = str(raw.get("lidar_id", "")).strip()
    if not lidar_id:
        raise ValueError(f"回放帧缺少 lidar_id: {raw!r}")
    points = points_from_value(raw.get("points_xyz", raw.get("points", [])))
    return LidarFrame(
        lidar_id=lidar_id,
        frame_id=str(raw.get("frame_id", lidar_id)),
        stamp_us=int(raw.get("stamp_us", raw.get("timestamp_us", fallback_timestamp_us)) or fallback_timestamp_us),
        points_xyz=points,
        point_count=len(points),
    )


def load_replay_sets(path: Path) -> list[dict[str, LidarFrame]]:
    frame_sets: list[dict[str, LidarFrame]] = []
    with path.open("r", encoding="utf-8") as stream:
        for line_no, line in enumerate(stream, start=1):
            line = line.strip()
            if not line:
                continue
            raw = json.loads(line)
            if not isinstance(raw, dict):
                raise ValueError(f"{path}:{line_no} 每行必须是 JSON 对象。")
            timestamp_us = int(raw.get("timestamp_us", now_us()) or now_us())
            frames: dict[str, LidarFrame] = {}
            if isinstance(raw.get("frames"), list):
                for item in raw["frames"]:
                    frame = frame_from_dict(item, timestamp_us)
                    frames[frame.lidar_id] = frame
            elif isinstance(raw.get("lidars"), dict):
                for lidar_id, item in raw["lidars"].items():
                    item = item if isinstance(item, dict) else {}
                    item.setdefault("lidar_id", lidar_id)
                    frame = frame_from_dict(item, timestamp_us)
                    frames[frame.lidar_id] = frame
            else:
                frame = frame_from_dict(raw, timestamp_us)
                frames[frame.lidar_id] = frame
            frame_sets.append(frames)
    return frame_sets


def synthetic_frame_set(config: dict[str, Any]) -> dict[str, LidarFrame]:
    timestamp_us = now_us()
    result: dict[str, LidarFrame] = {}
    for item in config.get("lidars", []):
        if not bool(item.get("enabled", True)):
            continue
        lidar_id = str(item["lidar_id"])
        result[lidar_id] = LidarFrame(
            lidar_id=lidar_id,
            frame_id=str(item.get("frame_id", lidar_id)),
            stamp_us=timestamp_us,
            points_xyz=[],
            point_count=0,
        )
    return result


def run_once(
    config: dict[str, Any],
    recorder: ExternalLocalizationRecorder,
    latest_frames: dict[str, LidarFrame],
) -> FourLidarBallLocalizationOutput:
    timestamp_us = now_us()
    lidar_mounts = parse_lidar_mounts(config)
    ball_targets = parse_ball_targets(config)
    max_frame_age_ms = float(config.get("max_frame_age_ms", 100.0))
    frame_status = build_frame_status(timestamp_us, latest_frames, lidar_mounts, max_frame_age_ms)
    max_age_us = int(max_frame_age_ms * 1000.0)
    frames = [
        frame
        for frame in latest_frames.values()
        if timestamp_us - int(frame.stamp_us) <= max_age_us
    ]
    algorithm_input = FourLidarBallLocalizationInput(
        frames=frames,
        lidar_mounts=lidar_mounts,
        ball_targets=ball_targets,
        max_frame_age_ms=max_frame_age_ms,
        publish_timestamp_us=timestamp_us,
        reference_source_name=str(config.get("reference_source_name", "")),
        workcell_zone_id=str(config.get("workcell_zone_id", "")),
    )
    output = FourLidarBallLocalizationAlgorithm().solve(algorithm_input)
    if not output.hardware_timestamp_us:
        output.hardware_timestamp_us = timestamp_us
    diagnostics = build_diagnostics_payload(
        timestamp_us,
        output,
        frame_status,
        str(config.get("reference_source_name", "")),
        str(config.get("workcell_zone_id", "")),
        recorder.recording_files(),
    )
    recorder.record(output, frame_status, diagnostics)
    print(json.dumps(diagnostics, ensure_ascii=False, separators=(",", ":")))
    return output


def main() -> int:
    args = parse_args()
    config_path = Path(args.config).expanduser().resolve(strict=False)
    config = load_config(config_path)
    apply_overrides(config, args)
    validation = validate_config(config, config_path)
    for warning in validation.warnings:
        print(f"[WARN] {warning}", file=sys.stderr)
    if not validation.ok:
        for error in validation.errors:
            print(f"[错误] {error}", file=sys.stderr)
        return 1

    recorder = ExternalLocalizationRecorder(config, config_path)
    if args.frames_jsonl:
        frame_sets = load_replay_sets(Path(args.frames_jsonl).expanduser().resolve(strict=False))
    else:
        frame_sets = [synthetic_frame_set(config)]

    latest_frames: dict[str, LidarFrame] = {}
    valid_count = 0
    for frame_set in frame_sets:
        latest_frames.update(frame_set)
        output = run_once(config, recorder, latest_frames)
        if output.pose_valid:
            valid_count += 1
    recorder.update_dataset_index()

    print(f"[OK] 回放完成: frames={len(frame_sets)}, valid_pose={valid_count}", file=sys.stderr)
    if args.require_valid_pose and valid_count <= 0:
        print("[错误] 没有输出有效外部定位位姿。", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
