#!/usr/bin/env python3
"""TCP smoke test for the simulated vehicle-side sensor 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_int, get_nested, load_profile, validate_sim_profile


SENSOR_GET_READINESS_REQ = 21
SENSOR_GET_READINESS_RSP = 22
SENSOR_GET_LATEST_FRAME_REQ = 23
SENSOR_GET_LATEST_FRAME_RSP = 24
SENSOR_LIST_SENSORS_REQ = 25
SENSOR_LIST_SENSORS_RSP = 26


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("--fake-sensors", action="store_true", help="发布假的图像、点云、LaserScan、IMU 数据")
    parser.add_argument("--hold-fake-sensors-sec", type=float, default=0.0, help="验收完成后继续发布假传感器数据的秒数")
    parser.add_argument("--max-payload-bytes", type=int, default=256)
    return parser.parse_args()


def start_fake_sensors(topics: dict[str, str]) -> tuple[threading.Event, threading.Thread]:
    import rclpy
    from sensor_msgs.msg import Image, Imu, LaserScan, PointCloud2, PointField

    stop_event = threading.Event()

    def stamp_msg(msg: Any, frame_id: str) -> None:
        msg.header.frame_id = frame_id
        now_ns = time.time_ns()
        msg.header.stamp.sec = now_ns // 1_000_000_000
        msg.header.stamp.nanosec = now_ns % 1_000_000_000

    def run() -> None:
        if not rclpy.ok():
            rclpy.init(args=None)
        node = rclpy.create_node("vehicle_sensor_agent_smoke_fake_sensors")
        front_pub = node.create_publisher(Image, topics["front_camera"], 10)
        down_pub = node.create_publisher(Image, topics["down_camera"], 10)
        cloud_pub = node.create_publisher(PointCloud2, topics["lidar_3d"], 10)
        scan_pub = node.create_publisher(LaserScan, topics["lidar_2d"], 10)
        imu_pub = node.create_publisher(Imu, topics["imu"], 10)

        image_data = bytes([0, 32, 64, 96, 128, 160, 192, 255] * 6)
        point_data = bytes([0, 0, 0, 0, 0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64] * 4)
        try:
            while not stop_event.is_set():
                front = Image()
                stamp_msg(front, "front_camera_link")
                front.width = 4
                front.height = 4
                front.encoding = "rgb8"
                front.is_bigendian = 0
                front.step = 12
                front.data = image_data
                front_pub.publish(front)

                down = Image()
                stamp_msg(down, "down_camera_link")
                down.width = 4
                down.height = 4
                down.encoding = "mono8"
                down.is_bigendian = 0
                down.step = 4
                down.data = image_data[:16]
                down_pub.publish(down)

                cloud = PointCloud2()
                stamp_msg(cloud, "lidar_3d_link")
                cloud.height = 1
                cloud.width = 4
                cloud.fields = [
                    PointField(name="x", offset=0, datatype=PointField.FLOAT32, count=1),
                    PointField(name="y", offset=4, datatype=PointField.FLOAT32, count=1),
                    PointField(name="z", offset=8, datatype=PointField.FLOAT32, count=1),
                    PointField(name="intensity", offset=12, datatype=PointField.FLOAT32, count=1),
                ]
                cloud.is_bigendian = False
                cloud.point_step = 16
                cloud.row_step = 64
                cloud.is_dense = True
                cloud.data = point_data
                cloud_pub.publish(cloud)

                scan = LaserScan()
                stamp_msg(scan, "lidar_2d_link")
                scan.angle_min = -1.57
                scan.angle_max = 1.57
                scan.angle_increment = 0.785
                scan.time_increment = 0.0
                scan.scan_time = 0.1
                scan.range_min = 0.05
                scan.range_max = 20.0
                scan.ranges = [1.0, 1.2, 1.5, 1.2, 1.0]
                scan.intensities = [10.0, 20.0, 30.0, 20.0, 10.0]
                scan_pub.publish(scan)

                imu = Imu()
                stamp_msg(imu, "imu_link")
                imu.orientation.w = 1.0
                imu.angular_velocity.z = 0.01
                imu.linear_acceleration.z = 9.81
                imu_pub.publish(imu)

                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"]
    sensor_agent = profile.get("vehicle_sensor_agent", {})
    host = args.host or str(gateway.get("sensor_host", gateway.get("control_host", "127.0.0.1")))
    port = as_int(
        gateway.get("vehicle_port", gateway.get("sensor_port", 9003)),
        "workshop_pc.gateway.vehicle_port",
    )
    topics = {
        "front_camera": str(get_nested(profile, "vehicle_sensor_agent.front_camera_topic", "/sensor/front_camera/image_raw")),
        "down_camera": str(get_nested(profile, "vehicle_sensor_agent.down_camera_topic", "/sensor/down_camera/image_raw")),
        "lidar_3d": str(get_nested(profile, "vehicle_sensor_agent.lidar_3d_topic", "/sensor/lidar_3d/pointcloud")),
        "lidar_2d": str(get_nested(profile, "vehicle_sensor_agent.lidar_2d_topic", "/sensor/lidar_2d/scan")),
        "imu": str(get_nested(profile, "vehicle_sensor_agent.imu_topic", "/sensor/imu/data")),
    }
    sensor_ids = [
        str(sensor_agent.get("front_camera_sensor_id", "demo_front_camera")),
        str(sensor_agent.get("down_camera_sensor_id", "demo_down_camera")),
        str(sensor_agent.get("lidar_3d_sensor_id", "demo_lidar_3d")),
        str(sensor_agent.get("lidar_2d_sensor_id", "demo_lidar_2d")),
        str(sensor_agent.get("imu_sensor_id", "demo_imu")),
    ]

    fake_stop = None
    fake_thread = None
    if args.fake_sensors:
        fake_stop, fake_thread = start_fake_sensors(topics)
        time.sleep(0.5)

    try:
        rsp_type, payload = send_request(host, port, SENSOR_LIST_SENSORS_REQ, {}, args.timeout_sec)
        assert_success("sensor list", SENSOR_LIST_SENSORS_RSP, rsp_type, payload)

        rsp_type, payload = send_request(host, port, SENSOR_GET_READINESS_REQ, {}, args.timeout_sec)
        assert_success("sensor readiness", SENSOR_GET_READINESS_RSP, rsp_type, payload)

        for sensor_id in sensor_ids:
            rsp_type, payload = send_request(
                host,
                port,
                SENSOR_GET_LATEST_FRAME_REQ,
                {
                    "sensor_id": sensor_id,
                    "include_payload": True,
                    "max_payload_bytes": args.max_payload_bytes,
                    "fragment_index": 0,
                },
                args.timeout_sec,
            )
            assert_success(f"latest frame {sensor_id}", SENSOR_GET_LATEST_FRAME_RSP, rsp_type, payload)
            frame = payload.get("sensor_frame", {})
            if not frame.get("header", {}).get("sequence_id", 0):
                raise AssertionError(f"{sensor_id}: sequence_id missing in payload={payload}")
        if args.fake_sensors and args.hold_fake_sensors_sec > 0.0:
            time.sleep(args.hold_fake_sensors_sec)
    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 sensor 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)
