#!/usr/bin/env python3
"""车间电脑侧四角 LiDAR 外部真值定位节点模板。"""

from __future__ import annotations

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

import rclpy
from rclpy.node import Node

from calibration_external_localization_interfaces.msg import ExternalLocalizationTelemetry
from sensor_msgs.msg import LaserScan, PointCloud2
from sensor_msgs_py import point_cloud2
from std_msgs.msg import String

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,
    message_type_for_lidar,
    now_us,
    parse_ball_targets,
    parse_lidar_mounts,
    validate_config,
)


def msg_stamp_us(msg: Any) -> int:
    stamp = getattr(getattr(msg, "header", None), "stamp", None)
    if stamp is None:
        return now_us()
    return int(stamp.sec) * 1_000_000 + int(stamp.nanosec) // 1_000


class WorkshopExternalLidarLocalizationNode(Node):
    """固定 ROS 通信和配置加载，真实定位算法只在算法类中实现。"""

    def __init__(self, args: argparse.Namespace) -> None:
        super().__init__("workshop_external_lidar_localization_real")
        self.config_path = Path(args.config).expanduser().resolve()
        self.config = load_config(self.config_path)
        self.apply_cli_overrides(args)
        validation = validate_config(self.config, self.config_path)
        for warning in validation.warnings:
            self.get_logger().warning(f"外部真值定位配置警告: {warning}")
        if not validation.ok:
            raise ValueError("外部真值定位配置无效:\n  - " + "\n  - ".join(validation.errors))

        self.reference_source_name = str(
            self.config.get("reference_source_name", "workshop_four_lidar_ball_truth")
        )
        self.workcell_zone_id = str(self.config.get("workcell_zone_id", "workshop_external_lidar_zone"))
        self.max_frame_age_ms = float(self.config.get("max_frame_age_ms", 100.0))
        self.max_points_per_frame = int(self.config.get("max_points_per_frame", 20000))
        self.lidar_mounts = parse_lidar_mounts(self.config)
        self.ball_targets = parse_ball_targets(self.config)
        self.latest_frames: dict[str, LidarFrame] = {}
        self.algorithm = FourLidarBallLocalizationAlgorithm()
        self.recorder = ExternalLocalizationRecorder(self.config, self.config_path)
        self.publisher = self.create_publisher(
            ExternalLocalizationTelemetry,
            str(self.config.get("output_topic", "/workshop/external_localization/vehicle/pose")),
            10,
        )
        self.diagnostics_publisher = self.create_publisher(
            String,
            str(self.config.get("diagnostics_topic", "/workshop/external_localization/diagnostics")),
            10,
        )

        for mount in self.lidar_mounts:
            if not mount.enabled:
                continue
            message_type = message_type_for_lidar(self.config, mount.lidar_id)
            if message_type in ("pointcloud2", "point_cloud2", "pointcloud"):
                self.create_subscription(
                    PointCloud2,
                    mount.topic,
                    lambda msg, lidar_id=mount.lidar_id: self.on_point_cloud(lidar_id, msg),
                    10,
                )
            elif message_type in ("laserscan", "laser_scan", "scan"):
                self.create_subscription(
                    LaserScan,
                    mount.topic,
                    lambda msg, lidar_id=mount.lidar_id: self.on_laser_scan(lidar_id, msg),
                    10,
                )
            else:
                raise ValueError(f"不支持的 LiDAR 消息类型: {mount.lidar_id} {message_type}")

        publish_hz = float(self.config.get("publish_hz", 20.0))
        period_sec = 1.0 / publish_hz if publish_hz > 0.0 else 0.05
        self.create_timer(period_sec, self.on_timer)
        self.get_logger().info(
            f"四角 LiDAR 外部真值定位模板已启动: config={self.config_path}, source={self.reference_source_name}"
        )
        if self.recorder.enabled:
            self.get_logger().info(
                f"外部真值定位数据落盘已启用: dataset_index={self.recorder.dataset_index_path}"
            )

    def apply_cli_overrides(self, args: argparse.Namespace) -> None:
        recording = self.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 on_point_cloud(self, lidar_id: str, msg: PointCloud2) -> None:
        points: list[tuple[float, float, float]] = []
        for point in point_cloud2.read_points(msg, field_names=("x", "y", "z"), skip_nans=True):
            if isinstance(point, dict):
                x, y, z = point["x"], point["y"], point["z"]
            else:
                x, y, z = point[0], point[1], point[2]
            points.append((float(x), float(y), float(z)))
            if len(points) >= self.max_points_per_frame:
                break

        frame_id = str(getattr(msg.header, "frame_id", "") or lidar_id)
        self.latest_frames[lidar_id] = LidarFrame(
            lidar_id=lidar_id,
            frame_id=frame_id,
            stamp_us=msg_stamp_us(msg),
            points_xyz=points,
            point_count=len(points),
        )

    def on_laser_scan(self, lidar_id: str, msg: LaserScan) -> None:
        points: list[tuple[float, float, float]] = []
        angle = float(msg.angle_min)
        for raw_range in msg.ranges:
            distance = float(raw_range)
            if math.isfinite(distance) and msg.range_min <= distance <= msg.range_max:
                points.append((distance * math.cos(angle), distance * math.sin(angle), 0.0))
                if len(points) >= self.max_points_per_frame:
                    break
            angle += float(msg.angle_increment)

        frame_id = str(getattr(msg.header, "frame_id", "") or lidar_id)
        self.latest_frames[lidar_id] = LidarFrame(
            lidar_id=lidar_id,
            frame_id=frame_id,
            stamp_us=msg_stamp_us(msg),
            points_xyz=points,
            point_count=len(points),
        )

    def recent_frames(self, timestamp_us: int) -> list[LidarFrame]:
        max_age_us = int(self.max_frame_age_ms * 1000.0)
        frames = []
        for frame in self.latest_frames.values():
            if timestamp_us - frame.stamp_us <= max_age_us:
                frames.append(frame)
        return frames

    def on_timer(self) -> None:
        publish_timestamp_us = now_us()
        frames = self.recent_frames(publish_timestamp_us)
        frame_status = build_frame_status(
            publish_timestamp_us,
            self.latest_frames,
            self.lidar_mounts,
            self.max_frame_age_ms,
        )
        algorithm_input = FourLidarBallLocalizationInput(
            frames=frames,
            lidar_mounts=self.lidar_mounts,
            ball_targets=self.ball_targets,
            max_frame_age_ms=self.max_frame_age_ms,
            publish_timestamp_us=publish_timestamp_us,
            reference_source_name=self.reference_source_name,
            workcell_zone_id=self.workcell_zone_id,
        )
        try:
            output = self.algorithm.solve(algorithm_input)
        except Exception as exc:
            output = FourLidarBallLocalizationOutput()
            output.hardware_timestamp_us = publish_timestamp_us
            output.diagnostic_messages.append(f"外部真值定位算法异常: {exc}")
            self.get_logger().warning(output.diagnostic_messages[-1])
        if not output.hardware_timestamp_us:
            output.hardware_timestamp_us = publish_timestamp_us
        diagnostics_payload = build_diagnostics_payload(
            publish_timestamp_us,
            output,
            frame_status,
            self.reference_source_name,
            self.workcell_zone_id,
            self.recorder.recording_files(),
        )
        self.publish_telemetry(output, publish_timestamp_us)
        self.publish_diagnostics(diagnostics_payload)
        self.recorder.record(output, frame_status, diagnostics_payload)

    def publish_telemetry(self, output: FourLidarBallLocalizationOutput, publish_timestamp_us: int) -> None:
        msg = ExternalLocalizationTelemetry()
        msg.hardware_timestamp_us = int(output.hardware_timestamp_us or publish_timestamp_us)
        msg.pose_valid = bool(output.pose_valid)
        msg.workshop_pose.x_m = float(output.workshop_pose.x_m)
        msg.workshop_pose.y_m = float(output.workshop_pose.y_m)
        msg.workshop_pose.z_m = float(output.workshop_pose.z_m)
        msg.workshop_pose.roll_rad = float(output.workshop_pose.roll_rad)
        msg.workshop_pose.pitch_rad = float(output.workshop_pose.pitch_rad)
        msg.workshop_pose.yaw_rad = float(output.workshop_pose.yaw_rad)
        msg.position_stddev_m = float(output.position_stddev_m)
        msg.yaw_stddev_rad = float(output.yaw_stddev_rad)
        msg.tracking_loss_ratio = float(output.tracking_loss_ratio)
        msg.time_sync_offset_ms = float(output.time_sync_offset_ms)
        msg.quality_score = float(output.quality_score)
        msg.observed_target_count = int(output.observed_target_count)
        msg.reference_source_name = self.reference_source_name
        msg.active_job_id = f"{self.workcell_zone_id}:four_lidar_ball_localization"
        self.publisher.publish(msg)

    def publish_diagnostics(self, diagnostics_payload: dict[str, Any]) -> None:
        msg = String()
        msg.data = json.dumps(diagnostics_payload, ensure_ascii=False, separators=(",", ":"))
        self.diagnostics_publisher.publish(msg)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="四角 LiDAR + 标定球外部真值定位节点模板")
    parser.add_argument(
        "--config",
        required=True,
        help="四角 LiDAR + 标定球定位配置文件路径",
    )
    parser.add_argument(
        "--session-dir",
        default="",
        help="覆盖 recording.session_dir，并启用本次会话数据落盘。",
    )
    parser.add_argument(
        "--dataset-index",
        default="",
        help="覆盖 recording.dataset_index_path，并启用数据集索引写入。",
    )
    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="关闭外部定位观测和诊断落盘。",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    rclpy.init(args=None)
    node = None
    try:
        node = WorkshopExternalLidarLocalizationNode(args)
        rclpy.spin(node)
    finally:
        if node is not None:
            node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()
    return 0


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