#!/usr/bin/env python3
"""Bridge external truth pose from ROS to the simulated WiFi6 TCP vehicle link."""

from __future__ import annotations

import argparse
import json
import socket
import struct
import time
from typing import Any

import rclpy
from rclpy.executors import ExternalShutdownException

try:
    from calibration_external_localization_interfaces.msg import ExternalLocalizationTelemetry
except ImportError:
    ExternalLocalizationTelemetry = None


EXTERNAL_POSE_PUSH_REQ = 31
EXTERNAL_POSE_PUSH_RSP = 32


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


def read_exactly(conn: socket.socket, size: int) -> bytes:
    chunks: list[bytes] = []
    remaining = size
    while remaining > 0:
        chunk = conn.recv(remaining)
        if not chunk:
            raise ConnectionError("connection closed")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def send_request(
    host: str,
    port: int,
    msg_type: int,
    payload: dict[str, Any],
    timeout_sec: float,
) -> tuple[int, dict[str, Any]]:
    encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    with socket.create_connection((host, port), timeout=timeout_sec) as conn:
        conn.settimeout(timeout_sec)
        conn.sendall(struct.pack("<II", msg_type, len(encoded)) + encoded)
        header = read_exactly(conn, 8)
        rsp_type, payload_len = struct.unpack("<II", header)
        raw_payload = read_exactly(conn, payload_len) if payload_len else b"{}"
    return rsp_type, json.loads(raw_payload.decode("utf-8"))


class ExternalPoseWifi6Bridge:
    def __init__(self, args: argparse.Namespace):
        if ExternalLocalizationTelemetry is None:
            raise RuntimeError("calibration_external_localization_interfaces is not available")

        self.args = args
        self.node = rclpy.create_node("external_pose_wifi6_bridge")
        self.last_send_monotonic = 0.0
        self.last_error_log_monotonic = 0.0
        self.sent_count = 0
        self.fail_count = 0
        self.subscription = self.node.create_subscription(
            ExternalLocalizationTelemetry,
            args.source_topic,
            self.on_pose,
            50,
        )
        self.node.get_logger().info(
            f"external pose bridge: {args.source_topic} -> {args.vehicle_host}:{args.vehicle_port}"
        )

    def rate_limited(self) -> bool:
        if self.args.max_hz <= 0.0:
            return False
        now = time.monotonic()
        period = 1.0 / self.args.max_hz
        if now - self.last_send_monotonic < period:
            return True
        self.last_send_monotonic = now
        return False

    def on_pose(self, msg: Any) -> None:
        if self.rate_limited():
            return

        pose = msg.workshop_pose
        source_name = self.args.reference_source_name or str(msg.reference_source_name or "external_truth_wifi6")
        payload = {
            "hardware_timestamp_us": int(msg.hardware_timestamp_us or now_us()),
            "pose_valid": bool(msg.pose_valid),
            "workshop_pose": {
                "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),
            },
            "position_stddev_m": float(msg.position_stddev_m),
            "yaw_stddev_rad": float(msg.yaw_stddev_rad),
            "tracking_loss_ratio": float(msg.tracking_loss_ratio),
            "time_sync_offset_ms": float(msg.time_sync_offset_ms),
            "quality_score": float(msg.quality_score),
            "observed_target_count": int(msg.observed_target_count),
            "reference_source_name": source_name,
            "active_job_id": str(msg.active_job_id or ""),
        }

        try:
            rsp_type, rsp = send_request(
                self.args.vehicle_host,
                self.args.vehicle_port,
                EXTERNAL_POSE_PUSH_REQ,
                payload,
                self.args.timeout_sec,
            )
            if rsp_type != EXTERNAL_POSE_PUSH_RSP or not rsp.get("success", False):
                raise RuntimeError(f"unexpected response type={rsp_type}, payload={rsp}")
            self.sent_count += 1
        except Exception as exc:
            self.fail_count += 1
            now = time.monotonic()
            if rclpy.ok() and now - self.last_error_log_monotonic >= self.args.error_log_interval_sec:
                self.last_error_log_monotonic = now
                self.node.get_logger().warning(
                    f"failed to push external pose over wifi6_sim_tcp: {exc}; fail_count={self.fail_count}"
                )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="外部真值位姿 ROS -> WiFi6/TCP 车端桥")
    parser.add_argument("--source-topic", default="/isaac/external_localization/telemetry")
    parser.add_argument("--vehicle-host", default="127.0.0.1")
    parser.add_argument("--vehicle-port", type=int, default=9000)
    parser.add_argument("--reference-source-name", default="")
    parser.add_argument("--max-hz", type=float, default=30.0)
    parser.add_argument("--timeout-sec", type=float, default=2.0)
    parser.add_argument("--error-log-interval-sec", type=float, default=2.0)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    rclpy.init(args=None)
    bridge = None
    try:
        bridge = ExternalPoseWifi6Bridge(args)
        rclpy.spin(bridge.node)
    except (KeyboardInterrupt, ExternalShutdownException):
        pass
    finally:
        if bridge is not None:
            bridge.node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()
    return 0


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