#!/usr/bin/env python3
# Windows 侧 TCP stub server（本机联调用）
# 用途：模拟 Windows 车端代理，收到任意帧后回复对应的 RSP 帧（含最小合法 JSON payload）。
# 启动方式：python3 windows_stub_server.py
# 默认监听：
#   9001 端口 —— 底盘域
#   9002 端口 —— 运控域
#   9003 端口 —— 传感器域
#   9004 端口 —— 外部真值位姿输入

import json
import socket
import struct
import threading
import time
import logging

logging.basicConfig(level=logging.INFO, format="[%(threadName)s] %(message)s")

# 帧类型（与 proto_frame.hpp 和 proto/transport_contract.proto 保持一致）
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
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
EXTERNAL_POSE_PUSH_REQ      = 31
EXTERNAL_POSE_PUSH_RSP      = 32

# REQ → RSP 类型映射
REQ_TO_RSP = {
    CHASSIS_GET_READINESS_REQ:   CHASSIS_GET_READINESS_RSP,
    CHASSIS_MOTION_PRIMITIVE_REQ: CHASSIS_MOTION_PRIMITIVE_RSP,
    CHASSIS_EMERGENCY_BRAKE_REQ:  CHASSIS_EMERGENCY_BRAKE_RSP,
    CONTROL_GET_READINESS_REQ:   CONTROL_GET_READINESS_RSP,
    CONTROL_EVALUATION_REQ:      CONTROL_EVALUATION_RSP,
    SENSOR_GET_READINESS_REQ:    SENSOR_GET_READINESS_RSP,
    SENSOR_GET_LATEST_FRAME_REQ: SENSOR_GET_LATEST_FRAME_RSP,
    SENSOR_LIST_SENSORS_REQ:     SENSOR_LIST_SENSORS_RSP,
    EXTERNAL_POSE_PUSH_REQ:      EXTERNAL_POSE_PUSH_RSP,
}


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


def _make_payload(req_type: int, req_payload: dict) -> dict:
    """根据请求类型生成最小合法的成功响应 payload。"""
    if req_type == CHASSIS_GET_READINESS_REQ:
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 底盘就绪。",
            "agent_ready": True,
            "chassis_driver_online": True,
            "motion_control_ready": True,
            "estop_released": True,
            "vehicle_safe_to_move": True,
            "telemetry_available": True,
            "checked_timestamp_us": _ts_us(),
        }
    elif req_type == CHASSIS_MOTION_PRIMITIVE_REQ:
        job_id = req_payload.get("header", {}).get("request_id", "stub_job")
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 动作原语执行完成。",
            "job_id": job_id,
            "data_quality_passed": True,
            "suitable_for_commit": True,
            "recommended_parameter_version": "stub_v1",
            "estimated_straight_line_bias": 0.001,
            "validation_summary": {
                "max_lateral_error_m": 0.010,
                "max_yaw_error_rad": 0.005,
                "rms_lateral_error_m": 0.005,
                "rms_yaw_error_rad": 0.003,
                "repeatability_error_m": 0.001,
                "curvature_error": 0.001,
                "module_consistency_error": 0.0,
                "auto_acceptance_passed": True,
            },
            "artifacts": [],
        }
    elif req_type == CHASSIS_EMERGENCY_BRAKE_REQ:
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 紧急制动已执行。",
        }
    elif req_type == CONTROL_GET_READINESS_REQ:
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 运控就绪。",
            "agent_ready": True,
            "trajectory_executor_ready": True,
            "vehicle_feedback_ready": True,
            "control_output_ready": True,
            "estop_released": True,
            "vehicle_safe_to_move": True,
            "checked_timestamp_us": _ts_us(),
            "external_pose_feedback_ready": True,
            "external_pose_source_name": "stub_external_truth",
            "external_pose_age_ms": 10.0,
            "external_pose_quality_score": 1.0,
            "external_pose_transport": "wifi6_stub_tcp",
        }
    elif req_type == CONTROL_EVALUATION_REQ:
        job_id = req_payload.get("header", {}).get("request_id", "stub_job")
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 控制评估完成。",
            "job_id": job_id,
            "data_quality_passed": True,
            "suitable_for_commit": True,
            "recommended_parameter_version": "stub_control_v1",
            "validation_summary": {
                "rms_lateral_error_m": 0.015,
                "rms_heading_error_rad": 0.010,
                "rms_speed_error_ms": 0.05,
                "overshoot_ratio": 0.08,
                "settle_time_sec": 1.2,
                "stop_position_error_m": 0.02,
                "max_jerk": 0.3,
                "saturation_ratio": 0.05,
                "auto_acceptance_passed": True,
            },
            "artifacts": [],
        }
    elif req_type == SENSOR_GET_READINESS_REQ:
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 传感器数据流就绪。",
            "agent_ready": True,
            "capture_pipeline_ready": True,
            "storage_ready": True,
            "telemetry_ready": True,
            "vehicle_safe_to_move": True,
            "arm_ready": True,
            "ready_sensor_ids": ["demo_front_camera", "demo_down_camera", "demo_lidar_3d", "demo_lidar_2d", "demo_imu"],
            "issues": [],
            "checked_timestamp_us": _ts_us(),
            "sensor_data_stream_ready": True,
            "streaming_sensor_ids": ["demo_front_camera", "demo_down_camera", "demo_lidar_3d", "demo_lidar_2d", "demo_imu"],
            "sensor_data_transport": "wifi6_stub_tcp",
        }
    elif req_type == SENSOR_LIST_SENSORS_REQ:
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 传感器列表。",
            "sensors": [
                {"sensor_id": "demo_front_camera", "sensor_type": 2, "payload_type": 1, "topic": "/sensor/front_camera/image_raw", "frame_id": "front_camera_link"},
                {"sensor_id": "demo_down_camera", "sensor_type": 1, "payload_type": 1, "topic": "/sensor/down_camera/image_raw", "frame_id": "down_camera_link"},
                {"sensor_id": "demo_lidar_3d", "sensor_type": 4, "payload_type": 2, "topic": "/sensor/lidar_3d/pointcloud", "frame_id": "lidar_3d_link"},
                {"sensor_id": "demo_lidar_2d", "sensor_type": 5, "payload_type": 3, "topic": "/sensor/lidar_2d/scan", "frame_id": "lidar_2d_link"},
                {"sensor_id": "demo_imu", "sensor_type": 6, "payload_type": 4, "topic": "/sensor/imu/data", "frame_id": "imu_link"},
            ],
        }
    elif req_type == SENSOR_GET_LATEST_FRAME_REQ:
        sensor_id = req_payload.get("sensor_id", "demo_front_camera")
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 最新传感器帧。",
            "sensor_frame": {
                "header": {
                    "sensor_id": sensor_id,
                    "sensor_type": 2,
                    "frame_id": "front_camera_link",
                    "hardware_timestamp_us": _ts_us(),
                    "sequence_id": 1,
                    "active_job_id": "",
                    "capture_intent": 0,
                },
                "payload_type": 1,
                "stream_id": f"{sensor_id}:1",
                "fragment_index": 0,
                "total_fragments": 1,
                "is_final_fragment": True,
                "uncompressed_size_bytes": 0,
                "payload_digest": "",
                "image": {"width": 0, "height": 0, "encoding": "stub", "is_bigendian": False, "step": 0, "data": "", "compression": "none"},
            },
        }
    elif req_type == EXTERNAL_POSE_PUSH_REQ:
        return {
            "success": True,
            "error_code": {"code": 1},
            "message": "Stub: 外部真值位姿已接收。",
            "accepted_timestamp_us": _ts_us(),
            "reference_source_name": req_payload.get("reference_source_name", "stub_external_truth"),
        }
    return {}


def _read_exactly(conn, n: int) -> bytes:
    buf = b""
    while len(buf) < n:
        chunk = conn.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("连接已关闭")
        buf += chunk
    return buf


def handle_connection(conn: socket.socket, addr):
    logging.info(f"新连接来自 {addr}")
    try:
        header = _read_exactly(conn, 8)
        msg_type, payload_len = struct.unpack_from("<II", header)
        raw = _read_exactly(conn, payload_len) if payload_len > 0 else b""
        logging.info(f"收到帧 type={msg_type} payload_len={payload_len}")

        req_payload = json.loads(raw.decode("utf-8")) if raw else {}

        rsp_type = REQ_TO_RSP.get(msg_type)
        if rsp_type is None:
            logging.warning(f"未知帧类型 {msg_type}，忽略")
            return

        rsp_dict = _make_payload(msg_type, req_payload)
        rsp_data = json.dumps(rsp_dict, ensure_ascii=False).encode("utf-8")
        rsp_header = struct.pack("<II", rsp_type, len(rsp_data))
        conn.sendall(rsp_header + rsp_data)
        logging.info(f"已回复帧 type={rsp_type} payload_len={len(rsp_data)}")
    except (ConnectionError, json.JSONDecodeError) as e:
        logging.warning(f"连接异常：{e}")
    finally:
        conn.close()


def start_server(port: int, name: str):
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("0.0.0.0", port))
    srv.listen(8)
    logging.info(f"{name} stub server 监听端口 {port}")
    while True:
        conn, addr = srv.accept()
        t = threading.Thread(
            target=handle_connection, args=(conn, addr),
            name=f"{name}:{addr[1]}", daemon=True
        )
        t.start()


if __name__ == "__main__":
    t_chassis = threading.Thread(
        target=start_server, args=(9001, "chassis"), name="chassis-srv", daemon=True
    )
    t_control = threading.Thread(
        target=start_server, args=(9002, "control"), name="control-srv", daemon=True
    )
    t_sensor = threading.Thread(
        target=start_server, args=(9003, "sensor"), name="sensor-srv", daemon=True
    )
    t_external_pose = threading.Thread(
        target=start_server, args=(9004, "external-pose"), name="external-pose-srv", daemon=True
    )
    t_chassis.start()
    t_control.start()
    t_sensor.start()
    t_external_pose.start()
    logging.info("Windows stub server 已启动，按 Ctrl+C 退出。")
    try:
        t_chassis.join()
    except KeyboardInterrupt:
        logging.info("退出。")
