#!/usr/bin/env python3
"""仿真车端 WiFi6 gateway。

车间工控机只连接本脚本暴露的统一端口。脚本根据 frame_type 把请求转发到
仿真车端内部的底盘、运控、传感器、外部真值后端端口。
"""

from __future__ import annotations

import argparse
import json
import logging
import socket
import struct
import threading
from dataclasses import dataclass
from typing import Any


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


@dataclass(frozen=True)
class BackendRoute:
    name: str
    host: str
    port: int
    response_type: int


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, bytes]:
    header = recv_exactly(conn, 8)
    msg_type, payload_len = struct.unpack("<II", header)
    payload = recv_exactly(conn, payload_len) if payload_len else b""
    return msg_type, payload


def send_frame(conn: socket.socket, msg_type: int, payload: bytes) -> None:
    conn.sendall(struct.pack("<II", msg_type, len(payload)) + payload)


def json_error(message: str) -> bytes:
    payload: dict[str, Any] = {
        "success": False,
        "error_code": {"code": 3},
        "message": message,
    }
    return json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")


class Wifi6GatewaySim:
    def __init__(self, args: argparse.Namespace) -> None:
        self.bind_host = args.bind_host
        self.vehicle_port = args.vehicle_port
        self.timeout_sec = args.timeout_sec
        self.routes = self._build_routes(args)

    def _build_routes(self, args: argparse.Namespace) -> dict[int, BackendRoute]:
        chassis = (args.chassis_host, args.chassis_port)
        control = (args.control_host, args.control_port)
        sensor = (args.sensor_host, args.sensor_port)
        external_pose = (args.external_pose_host, args.external_pose_port)
        return {
            CHASSIS_GET_READINESS_REQ: BackendRoute("chassis", *chassis, CHASSIS_GET_READINESS_RSP),
            CHASSIS_MOTION_PRIMITIVE_REQ: BackendRoute("chassis", *chassis, CHASSIS_MOTION_PRIMITIVE_RSP),
            CHASSIS_EMERGENCY_BRAKE_REQ: BackendRoute("chassis", *chassis, CHASSIS_EMERGENCY_BRAKE_RSP),
            CONTROL_GET_READINESS_REQ: BackendRoute("control", *control, CONTROL_GET_READINESS_RSP),
            CONTROL_EVALUATION_REQ: BackendRoute("control", *control, CONTROL_EVALUATION_RSP),
            SENSOR_GET_READINESS_REQ: BackendRoute("sensor", *sensor, SENSOR_GET_READINESS_RSP),
            SENSOR_GET_LATEST_FRAME_REQ: BackendRoute("sensor", *sensor, SENSOR_GET_LATEST_FRAME_RSP),
            SENSOR_LIST_SENSORS_REQ: BackendRoute("sensor", *sensor, SENSOR_LIST_SENSORS_RSP),
            EXTERNAL_POSE_PUSH_REQ: BackendRoute("external_pose", *external_pose, EXTERNAL_POSE_PUSH_RSP),
        }

    def serve_forever(self) -> None:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
            server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            server.bind((self.bind_host, self.vehicle_port))
            server.listen(32)
            logging.info("仿真 WiFi6 gateway 监听 %s:%d", self.bind_host, self.vehicle_port)
            while True:
                conn, addr = server.accept()
                thread = threading.Thread(
                    target=self._handle_client,
                    args=(conn, addr),
                    name=f"wifi6-gateway:{addr[0]}:{addr[1]}",
                    daemon=True,
                )
                thread.start()

    def _handle_client(self, conn: socket.socket, addr: tuple[str, int]) -> None:
        with conn:
            conn.settimeout(self.timeout_sec)
            try:
                msg_type, payload = read_frame(conn)
                route = self.routes.get(msg_type)
                if route is None:
                    logging.warning("未知 WiFi6 frame_type=%d from %s:%d", msg_type, addr[0], addr[1])
                    send_frame(conn, 0, json_error(f"unsupported wifi6 frame_type={msg_type}"))
                    return
                rsp_type, rsp_payload = self._forward(route, msg_type, payload)
                send_frame(conn, rsp_type, rsp_payload)
            except Exception as exc:
                logging.warning("处理 WiFi6 gateway 请求失败: %s", exc)

    def _forward(self, route: BackendRoute, msg_type: int, payload: bytes) -> tuple[int, bytes]:
        try:
            with socket.create_connection((route.host, route.port), timeout=self.timeout_sec) as backend:
                backend.settimeout(self.timeout_sec)
                send_frame(backend, msg_type, payload)
                rsp_type, rsp_payload = read_frame(backend)
                if rsp_type != route.response_type:
                    return route.response_type, json_error(
                        f"{route.name} backend response type mismatch: expected {route.response_type}, got {rsp_type}"
                    )
                return rsp_type, rsp_payload
        except Exception as exc:
            return route.response_type, json_error(
                f"{route.name} backend unavailable at {route.host}:{route.port}: {exc}"
            )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="仿真车端统一 WiFi6 gateway")
    parser.add_argument("--bind-host", default="0.0.0.0")
    parser.add_argument("--vehicle-port", type=int, default=9000)
    parser.add_argument("--timeout-sec", type=float, default=5.0)
    parser.add_argument("--chassis-host", default="127.0.0.1")
    parser.add_argument("--chassis-port", type=int, default=9001)
    parser.add_argument("--control-host", default="127.0.0.1")
    parser.add_argument("--control-port", type=int, default=9002)
    parser.add_argument("--sensor-host", default="127.0.0.1")
    parser.add_argument("--sensor-port", type=int, default=9003)
    parser.add_argument("--external-pose-host", default="127.0.0.1")
    parser.add_argument("--external-pose-port", type=int, default=9004)
    parser.add_argument("--log-level", default="INFO")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    logging.basicConfig(level=getattr(logging, args.log_level.upper(), logging.INFO), format="[%(levelname)s] %(message)s")
    Wifi6GatewaySim(args).serve_forever()
    return 0


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