#!/usr/bin/env python3
"""统一车端 agent 仿真。

该进程模拟真实 Windows 小车车端电脑：对车间工控机只暴露一个 WiFi6/TCP
入口端口，所有底盘、运控、外部真值位姿和传感器请求都通过 frame_type 分发。
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import socket
import socketserver
import struct
import sys
import threading
from pathlib import Path
from types import ModuleType
from typing import Any

import rclpy
from rclpy.executors import MultiThreadedExecutor


PROJECT_ROOT = Path(__file__).resolve().parents[3]


def load_module(name: str, path: Path) -> ModuleType:
    spec = importlib.util.spec_from_file_location(name, path)
    if spec is None or spec.loader is None:
        raise ImportError(f"无法加载模块: {path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


vehicle_mod = load_module(
    "isaac_vehicle_agent_sim_module",
    PROJECT_ROOT / "src" / "simulation" / "vehicle_agent_sim" / "scripts" / "isaac_vehicle_agent_sim.py",
)
sensor_mod = load_module(
    "isaac_vehicle_sensor_agent_sim_module",
    PROJECT_ROOT / "src" / "simulation" / "vehicle_sensor_agent_sim" / "scripts" / "isaac_vehicle_sensor_agent_sim.py",
)


VEHICLE_REQUEST_TYPES = {
    vehicle_mod.CHASSIS_GET_READINESS_REQ,
    vehicle_mod.CHASSIS_MOTION_PRIMITIVE_REQ,
    vehicle_mod.CHASSIS_EMERGENCY_BRAKE_REQ,
    vehicle_mod.CONTROL_GET_READINESS_REQ,
    vehicle_mod.CONTROL_EVALUATION_REQ,
    vehicle_mod.EXTERNAL_POSE_PUSH_REQ,
}

SENSOR_REQUEST_TYPES = {
    sensor_mod.SENSOR_GET_READINESS_REQ,
    sensor_mod.SENSOR_GET_LATEST_FRAME_REQ,
    sensor_mod.SENSOR_LIST_SENSORS_REQ,
}


def recv_exactly(conn: socket.socket, size: int) -> bytes:
    data = b""
    while len(data) < size:
        chunk = conn.recv(size - len(data))
        if not chunk:
            raise ConnectionError("连接提前关闭")
        data += chunk
    return data


def read_frame(conn: socket.socket) -> tuple[int, dict[str, Any]]:
    header = recv_exactly(conn, 8)
    msg_type, payload_len = struct.unpack("<II", header)
    raw = recv_exactly(conn, payload_len).decode("utf-8") if payload_len else "{}"
    return msg_type, json.loads(raw)


def send_frame(conn: socket.socket, msg_type: int, payload: dict[str, Any]) -> None:
    encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    conn.sendall(struct.pack("<II", msg_type, len(encoded)) + encoded)


def error_payload(message: str) -> dict[str, Any]:
    return {
        "success": False,
        "error_code": {"code": vehicle_mod.ERROR_INVALID_ARGUMENT},
        "message": message,
    }


class UnifiedVehicleAgent:
    def __init__(self, args: argparse.Namespace) -> None:
        self.args = args
        self.vehicle_agent = vehicle_mod.IsaacVehicleAgentSim(args)
        self.sensor_agent = sensor_mod.IsaacVehicleSensorAgentSim(args)

    def handle_vehicle_request(self, msg_type: int, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]:
        if msg_type == vehicle_mod.CHASSIS_GET_READINESS_REQ:
            return vehicle_mod.CHASSIS_GET_READINESS_RSP, self.vehicle_agent.readiness_payload("chassis")
        if msg_type == vehicle_mod.CHASSIS_MOTION_PRIMITIVE_REQ:
            return vehicle_mod.CHASSIS_MOTION_PRIMITIVE_RSP, self.vehicle_agent.handle_motion_primitive(payload)
        if msg_type == vehicle_mod.CHASSIS_EMERGENCY_BRAKE_REQ:
            self.vehicle_agent.stop()
            return vehicle_mod.CHASSIS_EMERGENCY_BRAKE_RSP, {
                "success": True,
                "error_code": {"code": vehicle_mod.ERROR_OK},
                "message": "Isaac simulated brake executed.",
            }
        if msg_type == vehicle_mod.CONTROL_GET_READINESS_REQ:
            return vehicle_mod.CONTROL_GET_READINESS_RSP, self.vehicle_agent.readiness_payload("control")
        if msg_type == vehicle_mod.CONTROL_EVALUATION_REQ:
            return vehicle_mod.CONTROL_EVALUATION_RSP, self.vehicle_agent.handle_control_evaluation(payload)
        if msg_type == vehicle_mod.EXTERNAL_POSE_PUSH_REQ:
            return vehicle_mod.EXTERNAL_POSE_PUSH_RSP, self.vehicle_agent.handle_external_pose_push(payload)
        return 0, error_payload(f"unsupported vehicle msg_type={msg_type}")

    def handle_request(self, msg_type: int, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]:
        if msg_type in VEHICLE_REQUEST_TYPES:
            return self.handle_vehicle_request(msg_type, payload)
        if msg_type in SENSOR_REQUEST_TYPES:
            return self.sensor_agent.handle_request(msg_type, payload)
        return 0, error_payload(f"unsupported unified vehicle frame_type={msg_type}")

    def stop(self) -> None:
        self.vehicle_agent.stop()


class UnifiedTcpHandler(socketserver.BaseRequestHandler):
    def handle(self) -> None:
        agent: UnifiedVehicleAgent = self.server.agent
        try:
            msg_type, payload = read_frame(self.request)
            response_type, response_payload = agent.handle_request(msg_type, payload)
        except Exception as exc:
            response_type = 0
            response_payload = error_payload(str(exc))
        send_frame(self.request, response_type, response_payload)


class ThreadingTcpServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True
    daemon_threads = True


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Isaac 统一车端 agent 仿真")
    parser.add_argument("--vehicle-id", default="demo_agv_001")
    parser.add_argument("--bind-host", default="0.0.0.0")
    parser.add_argument("--vehicle-port", type=int, default=9000)
    parser.add_argument("--cmd-vel-topic", default="/vehicle/demo_agv_001/actuator/cmd_vel")
    parser.add_argument("--external-pose-topic", default="/isaac/external_localization/telemetry")
    parser.add_argument(
        "--external-pose-transport",
        choices=["ros_topic", "wifi6_tcp", "both"],
        default="wifi6_tcp",
        help="外部真值位姿输入方式；统一车端仿真默认只通过 WiFi6/TCP",
    )
    parser.add_argument("--external-pose-port", type=int, default=9000)
    parser.add_argument("--external-pose-min-quality", type=float, default=0.1)
    parser.add_argument("--chassis-telemetry-topic", default="/chassis/telemetry")
    parser.add_argument("--allow-chassis-pose-fallback", action="store_true")
    parser.add_argument("--publish-hz", type=float, default=20.0)
    parser.add_argument("--max-speed-mps", type=float, default=0.3)
    parser.add_argument("--wheel-base-m", type=float, default=0.80)
    parser.add_argument("--max-steering-angle-rad", type=float, default=0.60)
    parser.add_argument("--internal-command-timeout-sec", type=float, default=0.5)
    parser.add_argument("--internal-ackermann-command-topic", default="/vehicle/demo_agv_001/internal/ackermann_cmd")
    parser.add_argument("--internal-state-topic", default="/vehicle/demo_agv_001/internal/state")
    parser.add_argument("--internal-health-topic", default="/vehicle/demo_agv_001/internal/health")
    parser.add_argument("--internal-time-sync-topic", default="/vehicle/demo_agv_001/internal/time_sync_status")
    parser.add_argument("--max-angular-speed-rps", type=float, default=1.2)
    parser.add_argument("--max-motion-duration-sec", type=float, default=60.0)
    parser.add_argument("--pose-wait-timeout-sec", type=float, default=2.0)
    parser.add_argument("--pose-stale-timeout-sec", type=float, default=0.5)
    parser.add_argument("--trajectory-lookahead-m", type=float, default=0.45)
    parser.add_argument("--trajectory-goal-tolerance-m", type=float, default=0.08)
    parser.add_argument("--trajectory-slowdown-radius-m", type=float, default=0.5)
    parser.add_argument("--trajectory-slowdown-gain", type=float, default=0.8)
    parser.add_argument("--default-tracking-speed-mps", type=float, default=0.12)
    parser.add_argument("--min-tracking-speed-mps", type=float, default=0.03)
    parser.add_argument("--trajectory-max-rms-lateral-error-m", type=float, default=0.20)
    parser.add_argument("--trajectory-max-rms-heading-error-rad", type=float, default=0.60)

    parser.add_argument("--max-frame-payload-bytes", type=int, default=262144)
    parser.add_argument("--frame-stale-timeout-sec", type=float, default=1.0)
    parser.add_argument("--enabled-sensor-ids", nargs="*", default=[], help="只启用指定 sensor_id；默认启用全部")
    parser.add_argument("--sensor-link-status-topic", default="/vehicle/demo_agv_001/internal/sensor_link_status")
    parser.add_argument("--front-camera-sensor-id", default="demo_front_camera")
    parser.add_argument("--down-camera-sensor-id", default="demo_down_camera")
    parser.add_argument("--lidar-3d-sensor-id", default="demo_lidar_3d")
    parser.add_argument("--lidar-2d-sensor-id", default="demo_lidar_2d")
    parser.add_argument("--imu-sensor-id", default="demo_imu")
    parser.add_argument("--front-camera-topic", default="/sensor/front_camera/image_raw")
    parser.add_argument("--down-camera-topic", default="/sensor/down_camera/image_raw")
    parser.add_argument("--lidar-3d-topic", default="/sensor/lidar_3d/pointcloud")
    parser.add_argument("--lidar-2d-topic", default="/sensor/lidar_2d/scan")
    parser.add_argument("--imu-topic", default="/sensor/imu/data")
    return parser


def main() -> int:
    args = build_parser().parse_args()
    rclpy.init(args=None)
    agent = UnifiedVehicleAgent(args)
    server = None
    executor = MultiThreadedExecutor()
    executor.add_node(agent.vehicle_agent.node)
    executor.add_node(agent.sensor_agent.node)
    try:
        server = ThreadingTcpServer((args.bind_host, args.vehicle_port), UnifiedTcpHandler)
        server.agent = agent
        threading.Thread(target=server.serve_forever, daemon=True).start()
        print(f"[*] unified vehicle agent sim listening on {args.bind_host}:{args.vehicle_port}")
        print(f"[*] publishing private actuator commands to {args.cmd_vel_topic}")
        print("[*] subscribed vehicle sensor topics:")
        for sensor in agent.sensor_agent.sensors.values():
            print(f"  - {sensor.sensor_id}: {sensor.topic}")
        while rclpy.ok():
            executor.spin_once(timeout_sec=0.1)
    except OSError as exc:
        print(f"[ERROR] unified vehicle agent network startup failed: {exc}")
        return 1
    except KeyboardInterrupt:
        print("[*] unified vehicle agent shutdown requested.")
    finally:
        if server is not None:
            server.shutdown()
            server.server_close()
        try:
            if rclpy.ok():
                agent.stop()
        except Exception as exc:
            print(f"[WARN] unified vehicle agent stop during shutdown failed: {exc}")
        executor.remove_node(agent.vehicle_agent.node)
        executor.remove_node(agent.sensor_agent.node)
        agent.vehicle_agent.node.destroy_node()
        agent.sensor_agent.node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()
    return 0


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