#!/usr/bin/env python3
"""TCP smoke test for the simulated vehicle-side agent."""

from __future__ import annotations

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

from sim_profile import as_float, as_int, get_nested, load_profile, validate_sim_profile


CHASSIS_GET_READINESS_REQ = 1
CHASSIS_GET_READINESS_RSP = 2
CHASSIS_MOTION_PRIMITIVE_REQ = 3
CHASSIS_MOTION_PRIMITIVE_RSP = 4
CHASSIS_EMERGENCY_BRAKE_REQ = 5
CHASSIS_EMERGENCY_BRAKE_RSP = 6
CONTROL_GET_READINESS_REQ = 11
CONTROL_GET_READINESS_RSP = 12
CONTROL_EVALUATION_REQ = 13
CONTROL_EVALUATION_RSP = 14
EXTERNAL_POSE_PUSH_REQ = 31
EXTERNAL_POSE_PUSH_RSP = 32


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


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 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 = recv_exactly(conn, 8)
        rsp_type, payload_len = struct.unpack("<II", header)
        raw_payload = recv_exactly(conn, payload_len) if payload_len else b"{}"
    return rsp_type, json.loads(raw_payload.decode("utf-8"))


def assert_success(name: str, expected_type: int, actual_type: int, payload: dict[str, Any]) -> None:
    if actual_type != expected_type:
        raise AssertionError(f"{name}: 响应类型错误，期望 {expected_type}，实际 {actual_type}")
    if not payload.get("success", False):
        raise AssertionError(f"{name}: success=false, payload={payload}")
    print(f"[PASS] {name}: {payload.get('message', '')}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="仿真车端 TCP 接口快速验收")
    parser.add_argument("--profile", default="", help="仿真 profile 路径")
    parser.add_argument("--host", default="", help="覆盖 profile 中的车端 host")
    parser.add_argument("--timeout-sec", type=float, default=5.0)
    parser.add_argument("--motion", action="store_true", help="发送一个小距离直线动作，默认不动车")
    parser.add_argument("--control-eval", action="store_true", help="发送一个短时速度阶跃运控评估任务")
    parser.add_argument("--trajectory", action="store_true", help="发送一个短轨迹跟踪任务")
    parser.add_argument("--split-trajectory", action="store_true", help="把短轨迹拆成两段下发，验证分段缓存和最终拼接执行")
    parser.add_argument("--fake-telemetry", action="store_true", help="发布一段假的外部真值和底盘遥测，用于无 Isaac 时测试轨迹接口")
    parser.add_argument("--distance-m", type=float, default=0.20)
    parser.add_argument("--speed-mps", type=float, default=0.10)
    parser.add_argument("--skip-brake", action="store_true", help="跳过急停接口测试")
    return parser.parse_args()


def start_fake_vehicle_feedback(
    external_pose_topic: str,
    chassis_topic: str,
    x_m: float,
    y_m: float,
    yaw_rad: float,
    external_pose_tcp_host: str = "",
    external_pose_tcp_port: int = 0,
    external_pose_tcp_timeout_sec: float = 0.2,
) -> tuple[threading.Event, threading.Thread]:
    import rclpy
    from calibration_chassis_interfaces.msg import ChassisTelemetry
    from calibration_external_localization_interfaces.msg import ExternalLocalizationTelemetry

    stop_event = threading.Event()

    def run() -> None:
        if not rclpy.ok():
            rclpy.init(args=None)
        node = rclpy.create_node("vehicle_agent_smoke_fake_vehicle_feedback")
        external_publisher = node.create_publisher(ExternalLocalizationTelemetry, external_pose_topic, 10)
        chassis_publisher = node.create_publisher(ChassisTelemetry, chassis_topic, 10)
        try:
            while not stop_event.is_set():
                timestamp_us = now_us()

                pose_msg = ExternalLocalizationTelemetry()
                pose_msg.hardware_timestamp_us = timestamp_us
                pose_msg.pose_valid = True
                pose_msg.workshop_pose.x_m = x_m
                pose_msg.workshop_pose.y_m = y_m
                pose_msg.workshop_pose.z_m = 0.0
                pose_msg.workshop_pose.yaw_rad = yaw_rad
                pose_msg.position_stddev_m = 0.001
                pose_msg.yaw_stddev_rad = 0.001
                pose_msg.tracking_loss_ratio = 0.0
                pose_msg.time_sync_offset_ms = 1.0
                pose_msg.quality_score = 1.0
                pose_msg.observed_target_count = 4
                pose_msg.reference_source_name = "smoke_fake_external_truth"
                pose_msg.active_job_id = "vehicle_agent_smoke_fake_pose"
                external_publisher.publish(pose_msg)
                if external_pose_tcp_host and external_pose_tcp_port > 0:
                    try:
                        send_request(
                            external_pose_tcp_host,
                            external_pose_tcp_port,
                            EXTERNAL_POSE_PUSH_REQ,
                            {
                                "hardware_timestamp_us": timestamp_us,
                                "pose_valid": True,
                                "workshop_pose": {
                                    "x_m": x_m,
                                    "y_m": y_m,
                                    "z_m": 0.0,
                                    "roll_rad": 0.0,
                                    "pitch_rad": 0.0,
                                    "yaw_rad": yaw_rad,
                                },
                                "position_stddev_m": 0.001,
                                "yaw_stddev_rad": 0.001,
                                "tracking_loss_ratio": 0.0,
                                "time_sync_offset_ms": 1.0,
                                "quality_score": 1.0,
                                "observed_target_count": 4,
                                "reference_source_name": "smoke_fake_external_truth",
                                "active_job_id": "vehicle_agent_smoke_fake_pose",
                            },
                            external_pose_tcp_timeout_sec,
                        )
                    except Exception:
                        pass

                chassis_msg = ChassisTelemetry()
                chassis_msg.hardware_timestamp_us = timestamp_us
                chassis_msg.odom_x_m = x_m
                chassis_msg.odom_y_m = y_m
                chassis_msg.odom_yaw_rad = yaw_rad
                chassis_msg.linear_velocity_ms = 0.0
                chassis_msg.angular_velocity_rads = 0.0
                chassis_msg.estop_engaged = False
                chassis_msg.driver_error_code = 0
                chassis_msg.active_job_id = "vehicle_agent_smoke_fake_velocity"
                chassis_publisher.publish(chassis_msg)

                rclpy.spin_once(node, timeout_sec=0.0)
                time.sleep(0.05)
        finally:
            node.destroy_node()
            if rclpy.ok():
                rclpy.shutdown()

    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    return stop_event, thread


def main() -> int:
    args = parse_args()
    profile = load_profile(args.profile or None)
    validate_sim_profile(profile)

    gateway = profile["workshop_pc"]["gateway"]
    host = args.host or str(gateway.get("vehicle_host", gateway["chassis_host"]))
    vehicle_port = as_int(
        gateway.get("vehicle_port", gateway.get("chassis_port", 9000)),
        "workshop_pc.gateway.vehicle_port",
    )
    vehicle_id = str(profile["vehicle_id"])
    max_speed = as_float(get_nested(profile, "safety.max_speed_mps", 0.3), "safety.max_speed_mps")
    speed = min(abs(args.speed_mps), max_speed)
    external_pose_topic = str(get_nested(profile, "vehicle_agent.external_pose_topic", "/isaac/external_localization/telemetry"))
    external_pose_transport = str(get_nested(profile, "vehicle_agent.external_pose_transport", "both"))
    external_pose_tcp_host = str(gateway.get("vehicle_host", gateway.get("control_host", "127.0.0.1")))
    external_pose_tcp_port = as_int(gateway.get("vehicle_port", 9000), "workshop_pc.gateway.vehicle_port")
    chassis_topic = str(get_nested(profile, "vehicle_agent.chassis_telemetry_topic", "/chassis/telemetry"))
    fake_stop = None
    fake_thread = None

    if args.trajectory and args.fake_telemetry:
        fake_stop, fake_thread = start_fake_vehicle_feedback(
            external_pose_topic,
            chassis_topic,
            x_m=abs(args.distance_m),
            y_m=0.0,
            yaw_rad=0.0,
            external_pose_tcp_host=external_pose_tcp_host if external_pose_transport in ("wifi6_tcp", "both") else "",
            external_pose_tcp_port=external_pose_tcp_port if external_pose_transport in ("wifi6_tcp", "both") else 0,
        )
        time.sleep(0.2)

    try:
        chassis_req = {
            "agent_name": "chassis",
            "include_details": True,
            "target_resource_ids": [],
        }
        rsp_type, payload = send_request(host, vehicle_port, CHASSIS_GET_READINESS_REQ, chassis_req, args.timeout_sec)
        assert_success("chassis readiness", CHASSIS_GET_READINESS_RSP, rsp_type, payload)

        control_req = {
            "agent_name": "control",
            "include_details": True,
            "target_resource_ids": [],
        }
        rsp_type, payload = send_request(host, vehicle_port, CONTROL_GET_READINESS_REQ, control_req, args.timeout_sec)
        assert_success("control readiness", CONTROL_GET_READINESS_RSP, rsp_type, payload)

        if args.motion:
            request_id = f"sim_smoke_straight_{now_us()}"
            motion_req = {
                "header": {
                    "session_id": f"sim_smoke_{now_us()}",
                    "task_id": "sim_smoke_straight",
                    "vehicle_id": vehicle_id,
                    "request_id": request_id,
                    "client_send_timestamp_us": now_us(),
                    "operator_id": "sim_smoke",
                    "workshop_host": "simulation",
                },
                "test_case_id": "smoke_straight_line",
                "task_purpose": 1,
                "selected_primitive": 1,
                "brake_when_finished": True,
                "timeout_sec": max(3.0, args.distance_m / max(speed, 0.01) + 2.0),
                "source_iteration_id": "smoke_iter_001",
                "straight_line": {
                    "target_speed_ms": speed,
                    "target_distance_m": abs(args.distance_m),
                    "reverse": False,
                },
            }
            rsp_type, payload = send_request(host, vehicle_port, CHASSIS_MOTION_PRIMITIVE_REQ, motion_req, args.timeout_sec + 10.0)
            assert_success("chassis straight motion", CHASSIS_MOTION_PRIMITIVE_RSP, rsp_type, payload)

        if args.control_eval:
            request_id = f"sim_smoke_control_{now_us()}"
            control_req = {
                "header": {
                    "session_id": f"sim_smoke_{now_us()}",
                    "task_id": "sim_smoke_velocity_step",
                    "vehicle_id": vehicle_id,
                    "request_id": request_id,
                    "client_send_timestamp_us": now_us(),
                    "operator_id": "sim_smoke",
                    "workshop_host": "simulation",
                },
                "test_case_id": "smoke_velocity_step",
                "task_purpose": 1,
                "selected_task": 2,
                "source_iteration_id": "smoke_iter_001",
                "velocity_step": {
                    "target_velocity_ms": speed,
                    "hold_time_sec": 0.5,
                    "settle_before_step_sec": 0.1,
                },
            }
            rsp_type, payload = send_request(host, vehicle_port, CONTROL_EVALUATION_REQ, control_req, args.timeout_sec + 10.0)
            assert_success("control velocity step", CONTROL_EVALUATION_RSP, rsp_type, payload)

        if args.trajectory:
            request_id = f"sim_smoke_trajectory_{now_us()}"
            path = [
                {"x_m": 0.0, "y_m": 0.0, "yaw_rad": 0.0, "target_speed_ms": speed},
                {"x_m": abs(args.distance_m), "y_m": 0.0, "yaw_rad": 0.0, "target_speed_ms": speed},
            ]

            def make_trajectory_request(segment_path: list[dict[str, float]], segment_index: int, total_segments: int, is_final: bool) -> dict[str, Any]:
                return {
                    "header": {
                        "session_id": f"sim_smoke_{now_us()}",
                        "task_id": "sim_smoke_trajectory",
                        "vehicle_id": vehicle_id,
                        "request_id": f"{request_id}_seg{segment_index}" if total_segments > 1 else request_id,
                        "client_send_timestamp_us": now_us(),
                        "operator_id": "sim_smoke",
                        "workshop_host": "simulation",
                    },
                    "test_case_id": "smoke_trajectory_tracking",
                    "task_purpose": 1,
                    "selected_task": 1,
                    "source_iteration_id": "smoke_iter_001",
                    "trajectory_tracking": {
                        "path": segment_path,
                        "stop_at_end": True,
                        "timeout_sec": max(3.0, args.distance_m / max(speed, 0.01) + 2.0),
                        "required_external_pose_source_id": "smoke_fake_external_truth" if args.fake_telemetry else "",
                        "max_external_pose_age_ms": 500.0 if args.fake_telemetry else 0.0,
                        "min_external_pose_quality_score": 0.1,
                        "trajectory_id": request_id,
                        "segment_index": segment_index,
                        "total_segments": total_segments,
                        "is_final_segment": is_final,
                    },
                }

            if args.split_trajectory:
                segments = [path[:1], path[1:]]
                for segment_index, segment_path in enumerate(segments):
                    is_final = segment_index == len(segments) - 1
                    control_req = make_trajectory_request(segment_path, segment_index, len(segments), is_final)
                    rsp_type, payload = send_request(
                        host,
                        vehicle_port,
                        CONTROL_EVALUATION_REQ,
                        control_req,
                        args.timeout_sec + 10.0,
                    )
                    name = "control trajectory tracking" if is_final else f"control trajectory segment {segment_index} upload"
                    assert_success(name, CONTROL_EVALUATION_RSP, rsp_type, payload)
            else:
                control_req = make_trajectory_request(path, 0, 1, True)
                rsp_type, payload = send_request(host, vehicle_port, CONTROL_EVALUATION_REQ, control_req, args.timeout_sec + 10.0)
                assert_success("control trajectory tracking", CONTROL_EVALUATION_RSP, rsp_type, payload)

        if not args.skip_brake:
            rsp_type, payload = send_request(host, vehicle_port, CHASSIS_EMERGENCY_BRAKE_REQ, {}, args.timeout_sec)
            assert_success("chassis emergency brake", CHASSIS_EMERGENCY_BRAKE_RSP, rsp_type, payload)
    finally:
        if fake_stop is not None:
            fake_stop.set()
        if fake_thread is not None:
            fake_thread.join(timeout=2.0)

    print("[PASS] vehicle agent TCP smoke test completed.")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"[FAIL] {exc}", file=sys.stderr)
        raise SystemExit(1)
