#!/usr/bin/env python3
"""四角 LiDAR 外部真值定位的配置、诊断和落盘公共工具。"""

from __future__ import annotations

import csv
import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import yaml

from four_lidar_ball_algorithm import (
    BallTargetConfig,
    FourLidarBallLocalizationOutput,
    LidarFrame,
    LidarMountConfig,
    Pose3DValue,
)


POINTCLOUD_TYPES = {"pointcloud2", "point_cloud2", "pointcloud"}
LASER_SCAN_TYPES = {"laserscan", "laser_scan", "scan"}
SUPPORTED_MESSAGE_TYPES = POINTCLOUD_TYPES | LASER_SCAN_TYPES


def now_us() -> int:
    return int(time.time() * 1_000_000)


def load_config(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as stream:
        config = yaml.safe_load(stream) or {}
    if not isinstance(config, dict):
        raise ValueError(f"配置文件必须是 YAML 字典: {path}")
    return config


def pose_from_dict(data: dict[str, Any] | None) -> Pose3DValue:
    data = data or {}
    return Pose3DValue(
        x_m=float(data.get("x_m", 0.0) or 0.0),
        y_m=float(data.get("y_m", 0.0) or 0.0),
        z_m=float(data.get("z_m", 0.0) or 0.0),
        roll_rad=float(data.get("roll_rad", 0.0) or 0.0),
        pitch_rad=float(data.get("pitch_rad", 0.0) or 0.0),
        yaw_rad=float(data.get("yaw_rad", 0.0) or 0.0),
    )


def parse_lidar_mounts(config: dict[str, Any]) -> list[LidarMountConfig]:
    mounts: list[LidarMountConfig] = []
    for item in config.get("lidars", []):
        mounts.append(
            LidarMountConfig(
                lidar_id=str(item["lidar_id"]),
                topic=str(item["topic"]),
                frame_id=str(item.get("frame_id", item["lidar_id"])),
                workshop_pose=pose_from_dict(item.get("workshop_pose")),
                enabled=bool(item.get("enabled", True)),
            )
        )
    return mounts


def parse_ball_targets(config: dict[str, Any]) -> list[BallTargetConfig]:
    targets: list[BallTargetConfig] = []
    for item in config.get("ball_targets", []):
        targets.append(
            BallTargetConfig(
                ball_id=str(item["ball_id"]),
                radius_m=float(item["radius_m"]),
                center_in_base_link=pose_from_dict(item.get("center_in_base_link")),
                required=bool(item.get("required", True)),
            )
        )
    return targets


@dataclass
class ConfigValidationResult:
    errors: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)

    @property
    def ok(self) -> bool:
        return not self.errors


def _is_nonempty_string(value: Any) -> bool:
    return isinstance(value, str) and bool(value.strip())


def _has_placeholder(value: Any) -> bool:
    if not isinstance(value, str):
        return False
    markers = ("replace_with", "measured_on_site", "session_xxx")
    return any(marker in value for marker in markers)


def _require_pose_map(value: Any, field_name: str, result: ConfigValidationResult) -> None:
    if not isinstance(value, dict):
        result.errors.append(f"{field_name} 必须是 YAML 字典。")
        return
    for key in ("x_m", "y_m", "z_m", "roll_rad", "pitch_rad", "yaw_rad"):
        try:
            float(value.get(key, 0.0) or 0.0)
        except (TypeError, ValueError):
            result.errors.append(f"{field_name}.{key} 必须是数字。")


def validate_config(config: dict[str, Any], config_path: Path | None = None) -> ConfigValidationResult:
    """校验现场配置是否满足四角 LiDAR 外部定位的最小落地要求。"""

    result = ConfigValidationResult()
    validation = config.get("config_validation", {}) or {}
    if not isinstance(validation, dict):
        result.errors.append("config_validation 必须是 YAML 字典。")
        validation = {}

    min_enabled_lidar_count = int(validation.get("min_enabled_lidar_count", 4) or 4)
    min_required_ball_count = int(validation.get("min_required_ball_count", 2) or 2)

    output_topic = config.get("output_topic", "")
    if not _is_nonempty_string(output_topic):
        result.errors.append("output_topic 不能为空。")
    elif not str(output_topic).startswith("/workshop/"):
        result.warnings.append("output_topic 建议使用 /workshop/ 前缀，避免和仿真 topic 混淆。")

    reference_source_name = config.get("reference_source_name", "")
    if not _is_nonempty_string(reference_source_name) or _has_placeholder(reference_source_name):
        result.errors.append("reference_source_name 必须填写真实外部定位源名称。")

    workcell_zone_id = config.get("workcell_zone_id", "")
    if not _is_nonempty_string(workcell_zone_id) or _has_placeholder(workcell_zone_id):
        result.errors.append("workcell_zone_id 必须填写真实工位区域 ID。")

    publish_hz = float(config.get("publish_hz", 20.0) or 0.0)
    if publish_hz <= 0.0:
        result.errors.append("publish_hz 必须大于 0。")
    elif publish_hz < 5.0:
        result.warnings.append("publish_hz 低于 5Hz，运控使用外部位姿时可能不够平滑。")

    max_frame_age_ms = float(config.get("max_frame_age_ms", 100.0) or 0.0)
    if max_frame_age_ms <= 0.0:
        result.errors.append("max_frame_age_ms 必须大于 0。")

    lidars = config.get("lidars", [])
    if not isinstance(lidars, list):
        result.errors.append("lidars 必须是列表。")
        lidars = []
    lidar_ids: set[str] = set()
    enabled_lidar_count = 0
    for index, item in enumerate(lidars):
        field = f"lidars[{index}]"
        if not isinstance(item, dict):
            result.errors.append(f"{field} 必须是 YAML 字典。")
            continue
        lidar_id = str(item.get("lidar_id", "")).strip()
        if not lidar_id:
            result.errors.append(f"{field}.lidar_id 不能为空。")
        elif lidar_id in lidar_ids:
            result.errors.append(f"{field}.lidar_id 重复: {lidar_id}")
        lidar_ids.add(lidar_id)
        if bool(item.get("enabled", True)):
            enabled_lidar_count += 1
        topic = str(item.get("topic", "")).strip()
        if not topic or _has_placeholder(topic):
            result.errors.append(f"{field}.topic 必须填写真实 LiDAR topic。")
        message_type = str(item.get("message_type", "pointcloud2")).strip().lower()
        if message_type not in SUPPORTED_MESSAGE_TYPES:
            result.errors.append(f"{field}.message_type 不支持: {message_type}")
        frame_id = str(item.get("frame_id", "")).strip()
        if not frame_id or _has_placeholder(frame_id):
            result.errors.append(f"{field}.frame_id 必须填写真实 frame_id。")
        _require_pose_map(item.get("workshop_pose"), f"{field}.workshop_pose", result)

    if enabled_lidar_count < min_enabled_lidar_count:
        result.errors.append(
            f"启用 LiDAR 数量不足，当前 {enabled_lidar_count}，要求至少 {min_enabled_lidar_count}。"
        )

    ball_targets = config.get("ball_targets", [])
    if not isinstance(ball_targets, list):
        result.errors.append("ball_targets 必须是列表。")
        ball_targets = []
    ball_ids: set[str] = set()
    required_ball_count = 0
    for index, item in enumerate(ball_targets):
        field = f"ball_targets[{index}]"
        if not isinstance(item, dict):
            result.errors.append(f"{field} 必须是 YAML 字典。")
            continue
        ball_id = str(item.get("ball_id", "")).strip()
        if not ball_id:
            result.errors.append(f"{field}.ball_id 不能为空。")
        elif ball_id in ball_ids:
            result.errors.append(f"{field}.ball_id 重复: {ball_id}")
        ball_ids.add(ball_id)
        if bool(item.get("required", True)):
            required_ball_count += 1
        try:
            radius_m = float(item.get("radius_m", 0.0) or 0.0)
        except (TypeError, ValueError):
            radius_m = 0.0
        if radius_m <= 0.0:
            result.errors.append(f"{field}.radius_m 必须大于 0。")
        _require_pose_map(item.get("center_in_base_link"), f"{field}.center_in_base_link", result)

    if required_ball_count < min_required_ball_count:
        result.errors.append(
            f"required 标定球数量不足，当前 {required_ball_count}，要求至少 {min_required_ball_count}。"
        )

    recording = config.get("recording", {}) or {}
    if not isinstance(recording, dict):
        result.errors.append("recording 必须是 YAML 字典。")
        recording = {}
    if bool(recording.get("enabled", False)):
        for field_name in ("session_id", "site_id", "vehicle_id"):
            value = str(recording.get(field_name, "")).strip()
            if not value or _has_placeholder(value):
                result.errors.append(f"recording.{field_name} 必须填写真实值。")
        session_dir = str(recording.get("session_dir", "")).strip()
        if not session_dir or _has_placeholder(session_dir):
            result.errors.append("recording.session_dir 必须填写真实会话目录。")
        dataset_index_path = str(recording.get("dataset_index_path", "")).strip()
        if dataset_index_path and _has_placeholder(dataset_index_path):
            result.errors.append("recording.dataset_index_path 不能保留模板占位符。")

    if config_path is not None and not config_path.exists():
        result.errors.append(f"配置文件不存在: {config_path}")
    return result


def message_type_for_lidar(config: dict[str, Any], lidar_id: str) -> str:
    for item in config.get("lidars", []):
        if str(item.get("lidar_id", "")) == lidar_id:
            return str(item.get("message_type", "pointcloud2")).lower()
    return "pointcloud2"


def build_frame_status(
    timestamp_us: int,
    latest_frames: dict[str, LidarFrame],
    lidar_mounts: list[LidarMountConfig],
    max_frame_age_ms: float,
) -> dict[str, Any]:
    max_age_us = int(max_frame_age_ms * 1000.0)
    lidar_status: list[dict[str, Any]] = []
    missing_lidar_ids: list[str] = []
    stale_lidar_ids: list[str] = []
    recent_lidar_ids: list[str] = []
    for mount in lidar_mounts:
        if not mount.enabled:
            continue
        frame = latest_frames.get(mount.lidar_id)
        if frame is None:
            missing_lidar_ids.append(mount.lidar_id)
            lidar_status.append({
                "lidar_id": mount.lidar_id,
                "frame_id": mount.frame_id,
                "present": False,
                "fresh": False,
                "age_ms": None,
                "point_count": 0,
            })
            continue
        age_us = timestamp_us - int(frame.stamp_us)
        fresh = age_us <= max_age_us
        if fresh:
            recent_lidar_ids.append(mount.lidar_id)
        else:
            stale_lidar_ids.append(mount.lidar_id)
        lidar_status.append({
            "lidar_id": mount.lidar_id,
            "frame_id": frame.frame_id,
            "present": True,
            "fresh": fresh,
            "age_ms": max(0.0, float(age_us) / 1000.0),
            "point_count": int(frame.point_count),
        })
    return {
        "lidars": lidar_status,
        "enabled_lidar_count": sum(1 for mount in lidar_mounts if mount.enabled),
        "recent_lidar_count": len(recent_lidar_ids),
        "recent_lidar_ids": recent_lidar_ids,
        "missing_lidar_ids": missing_lidar_ids,
        "stale_lidar_ids": stale_lidar_ids,
    }


def pose_to_dict(pose: Pose3DValue) -> dict[str, float]:
    return {
        "x_m": float(pose.x_m),
        "y_m": float(pose.y_m),
        "z_m": float(pose.z_m),
        "roll_rad": float(pose.roll_rad),
        "pitch_rad": float(pose.pitch_rad),
        "yaw_rad": float(pose.yaw_rad),
    }


def output_to_dict(output: FourLidarBallLocalizationOutput) -> dict[str, Any]:
    return {
        "pose_valid": bool(output.pose_valid),
        "workshop_pose": pose_to_dict(output.workshop_pose),
        "position_stddev_m": float(output.position_stddev_m),
        "yaw_stddev_rad": float(output.yaw_stddev_rad),
        "tracking_loss_ratio": float(output.tracking_loss_ratio),
        "time_sync_offset_ms": float(output.time_sync_offset_ms),
        "quality_score": float(output.quality_score),
        "observed_target_count": int(output.observed_target_count),
        "hardware_timestamp_us": int(output.hardware_timestamp_us),
        "diagnostic_messages": list(output.diagnostic_messages),
    }


def build_diagnostics_payload(
    timestamp_us: int,
    output: FourLidarBallLocalizationOutput,
    frame_status: dict[str, Any],
    reference_source_name: str,
    workcell_zone_id: str,
    recording_files: dict[str, str] | None = None,
) -> dict[str, Any]:
    payload = {
        "timestamp_us": int(timestamp_us),
        "reference_source_name": reference_source_name,
        "workcell_zone_id": workcell_zone_id,
        "pose": output_to_dict(output),
        "frame_status": frame_status,
        "recording_files": recording_files or {},
    }
    payload["healthy"] = (
        bool(output.pose_valid)
        and frame_status.get("recent_lidar_count", 0) > 0
        and float(output.quality_score) > 0.0
    )
    return payload


def _resolve_recording_path(session_dir: Path, value: str, default_value: str) -> Path:
    raw = value or default_value
    path = Path(raw).expanduser()
    if path.is_absolute():
        return path
    return session_dir / path


def _relative_to_root(path: Path, root: Path) -> str:
    try:
        return str(path.resolve(strict=False).relative_to(root.resolve(strict=False)))
    except ValueError:
        return str(path.resolve(strict=False))


def _load_yaml_if_exists(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    with path.open("r", encoding="utf-8") as stream:
        data = yaml.safe_load(stream) or {}
    if not isinstance(data, dict):
        raise ValueError(f"{path} 顶层结构必须是 YAML 字典。")
    return data


class ExternalLocalizationRecorder:
    """把外部定位实时输出写入本次会话数据集索引。"""

    def __init__(self, config: dict[str, Any], config_path: Path) -> None:
        recording = config.get("recording", {}) or {}
        self.enabled = bool(recording.get("enabled", False))
        self.sample_count = 0
        self.update_every_n_samples = int(recording.get("dataset_index_update_every_n_samples", 20) or 20)
        self.start_us = now_us()
        self.end_us = self.start_us
        self.site_id = str(recording.get("site_id", ""))
        self.vehicle_id = str(recording.get("vehicle_id", ""))
        self.session_id = str(recording.get("session_id", ""))
        self.config_path = config_path
        self._paths: dict[str, Path] = {}
        if not self.enabled:
            self.session_dir = Path("")
            self.dataset_index_path = Path("")
            return

        session_dir_raw = str(recording.get("session_dir", "")).strip()
        self.session_dir = Path(session_dir_raw).expanduser().resolve(strict=False)
        self.dataset_index_path = _resolve_recording_path(
            self.session_dir,
            str(recording.get("dataset_index_path", "")),
            "dataset_index.yaml",
        )
        self._paths = {
            "external_observation_files": _resolve_recording_path(
                self.session_dir,
                str(recording.get("external_observation_file", "")),
                "external/external_observations.csv",
            ),
            "marker_observation_files": _resolve_recording_path(
                self.session_dir,
                str(recording.get("marker_observation_file", "")),
                "external/marker_observations.jsonl",
            ),
            "time_sync_diagnostic_files": _resolve_recording_path(
                self.session_dir,
                str(recording.get("time_sync_diagnostic_file", "")),
                "external/time_sync_diagnostics.jsonl",
            ),
        }
        for path in self._paths.values():
            path.parent.mkdir(parents=True, exist_ok=True)
        self.dataset_index_path.parent.mkdir(parents=True, exist_ok=True)
        self._ensure_observation_header()
        self.update_dataset_index()

    def recording_files(self) -> dict[str, str]:
        if not self.enabled:
            return {}
        return {key: str(path) for key, path in self._paths.items()}

    def _ensure_observation_header(self) -> None:
        path = self._paths["external_observation_files"]
        if path.exists() and path.stat().st_size > 0:
            return
        with path.open("w", newline="", encoding="utf-8") as stream:
            writer = csv.writer(stream)
            writer.writerow([
                "hardware_timestamp_us",
                "pose_valid",
                "x_m",
                "y_m",
                "z_m",
                "roll_rad",
                "pitch_rad",
                "yaw_rad",
                "position_stddev_m",
                "yaw_stddev_rad",
                "tracking_loss_ratio",
                "time_sync_offset_ms",
                "quality_score",
                "observed_target_count",
                "recent_lidar_count",
                "enabled_lidar_count",
                "diagnostic_messages",
            ])

    def record(
        self,
        output: FourLidarBallLocalizationOutput,
        frame_status: dict[str, Any],
        diagnostics_payload: dict[str, Any],
    ) -> None:
        if not self.enabled:
            return
        self.sample_count += 1
        self.end_us = int(diagnostics_payload["timestamp_us"])
        self._append_external_observation(output, frame_status)
        self._append_marker_observations(output, diagnostics_payload)
        self._append_time_sync_diagnostic(diagnostics_payload)
        if self.sample_count % max(1, self.update_every_n_samples) == 0:
            self.update_dataset_index()

    def _append_external_observation(
        self,
        output: FourLidarBallLocalizationOutput,
        frame_status: dict[str, Any],
    ) -> None:
        path = self._paths["external_observation_files"]
        pose = output.workshop_pose
        with path.open("a", newline="", encoding="utf-8") as stream:
            writer = csv.writer(stream)
            writer.writerow([
                int(output.hardware_timestamp_us),
                int(bool(output.pose_valid)),
                float(pose.x_m),
                float(pose.y_m),
                float(pose.z_m),
                float(pose.roll_rad),
                float(pose.pitch_rad),
                float(pose.yaw_rad),
                float(output.position_stddev_m),
                float(output.yaw_stddev_rad),
                float(output.tracking_loss_ratio),
                float(output.time_sync_offset_ms),
                float(output.quality_score),
                int(output.observed_target_count),
                int(frame_status.get("recent_lidar_count", 0)),
                int(frame_status.get("enabled_lidar_count", 0)),
                "|".join(output.diagnostic_messages),
            ])

    def _append_marker_observations(
        self,
        output: FourLidarBallLocalizationOutput,
        diagnostics_payload: dict[str, Any],
    ) -> None:
        path = self._paths["marker_observation_files"]
        entry = {
            "timestamp_us": int(diagnostics_payload["timestamp_us"]),
            "observed_target_count": int(output.observed_target_count),
            "observed_ball_targets": [
                {
                    "ball_id": item.ball_id,
                    "center_in_workshop": pose_to_dict(item.center_in_workshop),
                    "residual_m": float(item.residual_m),
                    "quality_score": float(item.quality_score),
                    "source_lidar_ids": list(item.source_lidar_ids),
                }
                for item in output.observed_ball_targets
            ],
        }
        with path.open("a", encoding="utf-8") as stream:
            stream.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")) + "\n")

    def _append_time_sync_diagnostic(self, diagnostics_payload: dict[str, Any]) -> None:
        path = self._paths["time_sync_diagnostic_files"]
        entry = {
            "timestamp_us": int(diagnostics_payload["timestamp_us"]),
            "time_sync_offset_ms": diagnostics_payload["pose"]["time_sync_offset_ms"],
            "frame_status": diagnostics_payload["frame_status"],
            "healthy": diagnostics_payload["healthy"],
        }
        with path.open("a", encoding="utf-8") as stream:
            stream.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")) + "\n")

    def update_dataset_index(self) -> None:
        if not self.enabled:
            return
        index = _load_yaml_if_exists(self.dataset_index_path)
        index.setdefault("schema_version", 1)
        session = index.setdefault("session", {})
        if not isinstance(session, dict):
            raise ValueError("dataset_index.yaml 中的 session 必须是 YAML 字典。")
        session.setdefault("session_id", self.session_id)
        session.setdefault("site_id", self.site_id)
        session.setdefault("vehicle_id", self.vehicle_id)
        session["dataset_root"] = str(self.session_dir)
        session["data_window_start_timestamp_us"] = int(session.get("data_window_start_timestamp_us") or self.start_us)
        session["data_window_end_timestamp_us"] = int(self.end_us)

        data_inputs = index.setdefault("data_inputs", {})
        external = data_inputs.setdefault("external", {})
        for field_name, path in self._paths.items():
            value = _relative_to_root(path, self.session_dir)
            files = external.setdefault(field_name, [])
            if isinstance(files, str):
                files = [files]
                external[field_name] = files
            if value not in files:
                files.append(value)

        metadata = index.setdefault("metadata", {})
        metadata["external_lidar_config_file"] = str(self.config_path)
        metadata["external_lidar_sample_count"] = int(self.sample_count)
        if self.session_id:
            session["session_id"] = self.session_id
        if self.site_id:
            session["site_id"] = self.site_id
        if self.vehicle_id:
            session["vehicle_id"] = self.vehicle_id
        self.dataset_index_path.write_text(
            yaml.safe_dump(index, sort_keys=False, allow_unicode=True),
            encoding="utf-8",
        )
