#!/usr/bin/env python3
"""Vehicle-side sensor agent simulator for Isaac validation.

The process represents the sensor-facing part of the vehicle computer. It
subscribes to Isaac ROS sensor topics and exposes the latest frames through the
same simple TCP frame envelope used by the chassis/control simulator.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import math
import socket
import socketserver
import struct
import threading
import time
from dataclasses import dataclass, field
from typing import Any

import rclpy
from sensor_msgs.msg import Image, Imu, LaserScan, PointCloud2

try:
    from vehicle_internal_interfaces.msg import SensorLinkStatus
except ImportError:
    SensorLinkStatus = None


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

ERROR_OK = 1
ERROR_INVALID_ARGUMENT = 2
ERROR_INVALID_STATE = 3

SENSOR_TYPE_DOWNWARD_CAMERA = 1
SENSOR_TYPE_FRONT_CAMERA = 2
SENSOR_TYPE_LIDAR_3D = 4
SENSOR_TYPE_LIDAR_2D = 5
SENSOR_TYPE_IMU = 6

PAYLOAD_IMAGE = 1
PAYLOAD_POINT_CLOUD = 2
PAYLOAD_LASER_SCAN = 3
PAYLOAD_IMU = 4


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


def ros_stamp_to_us(stamp: Any) -> int:
    sec = int(getattr(stamp, "sec", 0))
    nanosec = int(getattr(stamp, "nanosec", 0))
    if sec == 0 and nanosec == 0:
        return now_us()
    return sec * 1_000_000 + nanosec // 1000


def read_exactly(conn: socket.socket, count: int) -> bytes:
    chunks: list[bytes] = []
    remaining = count
    while remaining > 0:
        chunk = conn.recv(remaining)
        if not chunk:
            raise ConnectionError("connection closed")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


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 read_frame(conn: socket.socket) -> tuple[int, dict[str, Any]]:
    header = read_exactly(conn, 8)
    msg_type, payload_len = struct.unpack("<II", header)
    payload = read_exactly(conn, payload_len).decode("utf-8") if payload_len else "{}"
    return msg_type, json.loads(payload)


@dataclass(frozen=True)
class SensorSpec:
    sensor_id: str
    sensor_type: int
    payload_type: int
    topic: str
    frame_id: str


@dataclass
class LatestFrame:
    sensor: SensorSpec
    hardware_timestamp_us: int
    sequence_id: int
    received_monotonic: float
    payload: dict[str, Any]
    raw_bytes: bytes = b""
    raw_field_path: tuple[str, str] | None = None
    uncompressed_size_bytes: int = 0
    digest: str = ""
    active_job_id: str = ""
    capture_intent: int = 0


class IsaacVehicleSensorAgentSim:
    def __init__(self, args: argparse.Namespace):
        self.args = args
        self.node = rclpy.create_node("isaac_vehicle_sensor_agent_sim")
        self.lock = threading.Lock()
        self.sequence_by_sensor: dict[str, int] = {}
        self.latest_frames: dict[str, LatestFrame] = {}
        self.sensors = self.build_sensor_specs(args)
        self.subscriptions = []
        self.stop_requested = False
        self.sensor_link_status_publisher = None

        self.add_subscription("front_camera", Image, self.on_image)
        self.add_subscription("down_camera", Image, self.on_image)
        self.add_subscription("lidar_3d", PointCloud2, self.on_point_cloud)
        self.add_subscription("lidar_2d", LaserScan, self.on_laser_scan)
        self.add_subscription("imu", Imu, self.on_imu)
        if SensorLinkStatus is not None:
            self.sensor_link_status_publisher = self.node.create_publisher(SensorLinkStatus, args.sensor_link_status_topic, 10)
            self.node.create_timer(0.5, self.publish_sensor_link_status)
        else:
            print("[WARN] vehicle_internal_interfaces is unavailable; sensor link status topic is disabled.")

    @staticmethod
    def build_sensor_specs(args: argparse.Namespace) -> dict[str, SensorSpec]:
        return {
            "front_camera": SensorSpec(args.front_camera_sensor_id, SENSOR_TYPE_FRONT_CAMERA, PAYLOAD_IMAGE, args.front_camera_topic, "front_camera_link"),
            "down_camera": SensorSpec(args.down_camera_sensor_id, SENSOR_TYPE_DOWNWARD_CAMERA, PAYLOAD_IMAGE, args.down_camera_topic, "down_camera_link"),
            "lidar_3d": SensorSpec(args.lidar_3d_sensor_id, SENSOR_TYPE_LIDAR_3D, PAYLOAD_POINT_CLOUD, args.lidar_3d_topic, "lidar_3d_link"),
            "lidar_2d": SensorSpec(args.lidar_2d_sensor_id, SENSOR_TYPE_LIDAR_2D, PAYLOAD_LASER_SCAN, args.lidar_2d_topic, "lidar_2d_link"),
            "imu": SensorSpec(args.imu_sensor_id, SENSOR_TYPE_IMU, PAYLOAD_IMU, args.imu_topic, "imu_link"),
        }

    def add_subscription(self, key: str, msg_type: Any, callback: Any) -> None:
        sensor = self.sensors[key]
        if not self.sensor_enabled(sensor):
            return
        self.subscriptions.append(self.node.create_subscription(msg_type, sensor.topic, lambda msg, key=key: callback(key, msg), 20))

    def sensor_enabled(self, sensor: SensorSpec) -> bool:
        return not self.args.enabled_sensor_ids or sensor.sensor_id in self.args.enabled_sensor_ids

    def next_sequence(self, sensor_id: str) -> int:
        with self.lock:
            next_value = self.sequence_by_sensor.get(sensor_id, 0) + 1
            self.sequence_by_sensor[sensor_id] = next_value
            return next_value

    def store_frame(self, frame: LatestFrame) -> None:
        with self.lock:
            self.latest_frames[frame.sensor.sensor_id] = frame

    def on_image(self, key: str, msg: Image) -> None:
        sensor = self.sensors[key]
        raw = bytes(msg.data)
        payload = {
            "width": int(msg.width),
            "height": int(msg.height),
            "encoding": str(msg.encoding),
            "is_bigendian": bool(msg.is_bigendian),
            "step": int(msg.step),
            "compression": "none",
        }
        frame = LatestFrame(
            sensor=sensor,
            hardware_timestamp_us=ros_stamp_to_us(msg.header.stamp),
            sequence_id=self.next_sequence(sensor.sensor_id),
            received_monotonic=time.monotonic(),
            payload={"image": payload},
            raw_bytes=raw,
            raw_field_path=("image", "data"),
            uncompressed_size_bytes=len(raw),
            digest=f"sha256:{hashlib.sha256(raw).hexdigest()}",
        )
        self.store_frame(frame)

    def on_point_cloud(self, key: str, msg: PointCloud2) -> None:
        sensor = self.sensors[key]
        raw = bytes(msg.data)
        fields = [
            {
                "name": str(field.name),
                "offset": int(field.offset),
                "datatype": int(field.datatype),
                "count": int(field.count),
            }
            for field in msg.fields
        ]
        payload = {
            "height": int(msg.height),
            "width": int(msg.width),
            "fields": fields,
            "is_bigendian": bool(msg.is_bigendian),
            "point_step": int(msg.point_step),
            "row_step": int(msg.row_step),
            "is_dense": bool(msg.is_dense),
            "compression": "none",
        }
        frame = LatestFrame(
            sensor=sensor,
            hardware_timestamp_us=ros_stamp_to_us(msg.header.stamp),
            sequence_id=self.next_sequence(sensor.sensor_id),
            received_monotonic=time.monotonic(),
            payload={"point_cloud": payload},
            raw_bytes=raw,
            raw_field_path=("point_cloud", "data"),
            uncompressed_size_bytes=len(raw),
            digest=f"sha256:{hashlib.sha256(raw).hexdigest()}",
        )
        self.store_frame(frame)

    def on_laser_scan(self, key: str, msg: LaserScan) -> None:
        sensor = self.sensors[key]
        payload = {
            "angle_min_rad": float(msg.angle_min),
            "angle_max_rad": float(msg.angle_max),
            "angle_increment_rad": float(msg.angle_increment),
            "time_increment_sec": float(msg.time_increment),
            "scan_time_sec": float(msg.scan_time),
            "range_min_m": float(msg.range_min),
            "range_max_m": float(msg.range_max),
            "ranges_m": [float(value) for value in msg.ranges],
            "intensities": [float(value) for value in msg.intensities],
        }
        frame = LatestFrame(
            sensor=sensor,
            hardware_timestamp_us=ros_stamp_to_us(msg.header.stamp),
            sequence_id=self.next_sequence(sensor.sensor_id),
            received_monotonic=time.monotonic(),
            payload={"laser_scan": payload},
        )
        self.store_frame(frame)

    def on_imu(self, key: str, msg: Imu) -> None:
        sensor = self.sensors[key]
        payload = {
            "orientation_x": float(msg.orientation.x),
            "orientation_y": float(msg.orientation.y),
            "orientation_z": float(msg.orientation.z),
            "orientation_w": float(msg.orientation.w),
            "orientation_covariance": [float(value) for value in msg.orientation_covariance],
            "angular_velocity": {
                "x": float(msg.angular_velocity.x),
                "y": float(msg.angular_velocity.y),
                "z": float(msg.angular_velocity.z),
            },
            "angular_velocity_covariance": [float(value) for value in msg.angular_velocity_covariance],
            "linear_acceleration": {
                "x": float(msg.linear_acceleration.x),
                "y": float(msg.linear_acceleration.y),
                "z": float(msg.linear_acceleration.z),
            },
            "linear_acceleration_covariance": [float(value) for value in msg.linear_acceleration_covariance],
        }
        frame = LatestFrame(
            sensor=sensor,
            hardware_timestamp_us=ros_stamp_to_us(msg.header.stamp),
            sequence_id=self.next_sequence(sensor.sensor_id),
            received_monotonic=time.monotonic(),
            payload={"imu": payload},
        )
        self.store_frame(frame)

    def list_sensors_payload(self) -> dict[str, Any]:
        sensors = []
        for sensor in self.sensors.values():
            sensors.append({
                "sensor_id": sensor.sensor_id,
                "sensor_type": sensor.sensor_type,
                "payload_type": sensor.payload_type,
                "topic": sensor.topic,
                "frame_id": sensor.frame_id,
                "enabled": self.sensor_enabled(sensor),
            })
        return {
            "success": True,
            "error_code": {"code": ERROR_OK},
            "message": "Isaac vehicle sensor agent sim: sensors listed.",
            "sensors": sensors,
        }

    def readiness_payload(self) -> dict[str, Any]:
        now = time.monotonic()
        ready_sensor_ids = []
        per_sensor = []
        with self.lock:
            latest_frames = dict(self.latest_frames)

        active_sensors = [sensor for sensor in self.sensors.values() if self.sensor_enabled(sensor)]
        for sensor in active_sensors:
            frame = latest_frames.get(sensor.sensor_id)
            age_sec = (now - frame.received_monotonic) if frame is not None else math.inf
            ready = frame is not None and age_sec <= self.args.frame_stale_timeout_sec
            if ready:
                ready_sensor_ids.append(sensor.sensor_id)
            per_sensor.append({
                "sensor_id": sensor.sensor_id,
                "sensor_type": sensor.sensor_type,
                "topic": sensor.topic,
                "ready": ready,
                "age_ms": None if not math.isfinite(age_sec) else age_sec * 1000.0,
                "sequence_id": 0 if frame is None else frame.sequence_id,
            })

        required_sensor_ids = [sensor.sensor_id for sensor in active_sensors]
        all_ready = all(sensor_id in ready_sensor_ids for sensor_id in required_sensor_ids)
        return {
            "success": True,
            "error_code": {"code": ERROR_OK},
            "message": "Isaac vehicle sensor agent sim: sensor readiness checked.",
            "agent_ready": True,
            "capture_pipeline_ready": all_ready,
            "storage_ready": True,
            "telemetry_ready": True,
            "vehicle_safe_to_move": True,
            "arm_ready": True,
            "ready_sensor_ids": ready_sensor_ids,
            "issues": [] if all_ready else [
                {
                    "resource_id": sensor["sensor_id"],
                    "message": f"sensor frame unavailable or stale on {sensor['topic']}",
                    "severity": 2,
                }
                for sensor in per_sensor
                if not sensor["ready"]
            ],
            "checked_timestamp_us": now_us(),
            "sensor_data_stream_ready": bool(ready_sensor_ids),
            "streaming_sensor_ids": ready_sensor_ids,
            "sensor_data_transport": "wifi6_sim_tcp",
            "per_sensor": per_sensor,
        }

    def publish_sensor_link_status(self) -> None:
        if self.sensor_link_status_publisher is None:
            return

        now = time.monotonic()
        with self.lock:
            latest_frames = dict(self.latest_frames)
        for sensor in self.sensors.values():
            if not self.sensor_enabled(sensor):
                continue
            frame = latest_frames.get(sensor.sensor_id)
            age_sec = math.inf if frame is None else now - frame.received_monotonic
            online = frame is not None and age_sec <= self.args.frame_stale_timeout_sec

            msg = SensorLinkStatus()
            msg.hardware_timestamp_us = now_us()
            msg.sensor_id = sensor.sensor_id
            msg.sensor_type.value = sensor.sensor_type
            msg.frame_id = sensor.frame_id
            msg.source_channel = sensor.topic
            msg.online = online
            msg.frame_rate_hz = 0.0 if frame is None else min(1.0 / max(age_sec, 1e-3), 1000.0)
            msg.latest_frame_age_ms = 0.0 if not math.isfinite(age_sec) else age_sec * 1000.0
            msg.frame_count = 0 if frame is None else frame.sequence_id
            msg.dropped_frame_ratio = 0.0
            msg.latest_transport_latency_ms = msg.latest_frame_age_ms
            msg.status_message = "online" if online else "sensor frame unavailable or stale"
            self.sensor_link_status_publisher.publish(msg)

    def frame_payload(self, request: dict[str, Any]) -> dict[str, Any]:
        sensor_id = str(request.get("sensor_id", ""))
        if not sensor_id:
            raise ValueError("sensor_id is required")
        include_payload = bool(request.get("include_payload", True))
        fragment_index = int(request.get("fragment_index", 0) or 0)
        if fragment_index < 0:
            raise ValueError("fragment_index must be non-negative")
        max_payload_bytes = int(request.get("max_payload_bytes", 0) or 0)
        if max_payload_bytes <= 0:
            max_payload_bytes = self.args.max_frame_payload_bytes

        with self.lock:
            frame = self.latest_frames.get(sensor_id)
        if frame is None:
            return {
                "success": False,
                "error_code": {"code": ERROR_INVALID_STATE},
                "message": f"no frame has been received for sensor_id={sensor_id}",
            }

        sensor_frame = self.encode_vehicle_sensor_frame(frame, include_payload, max_payload_bytes, fragment_index)
        return {
            "success": True,
            "error_code": {"code": ERROR_OK},
            "message": "Isaac vehicle sensor agent sim: latest frame returned.",
            "sensor_frame": sensor_frame,
        }

    def encode_vehicle_sensor_frame(
        self,
        frame: LatestFrame,
        include_payload: bool,
        max_payload_bytes: int,
        fragment_index: int,
    ) -> dict[str, Any]:
        payload = json.loads(json.dumps(frame.payload, separators=(",", ":")))
        total_fragments = 1
        fragment = b""
        if include_payload and frame.raw_bytes:
            max_payload_bytes = max(1, max_payload_bytes)
            total_fragments = max(1, math.ceil(len(frame.raw_bytes) / max_payload_bytes))
            if fragment_index >= total_fragments:
                raise ValueError(f"fragment_index {fragment_index} exceeds total_fragments {total_fragments}")
            start = fragment_index * max_payload_bytes
            end = min(len(frame.raw_bytes), start + max_payload_bytes)
            fragment = frame.raw_bytes[start:end]
            if frame.raw_field_path is not None:
                object_name, field_name = frame.raw_field_path
                payload[object_name][field_name] = base64.b64encode(fragment).decode("ascii")
        elif frame.raw_field_path is not None:
            object_name, field_name = frame.raw_field_path
            payload[object_name][field_name] = ""

        sensor_frame = {
            "header": {
                "sensor_id": frame.sensor.sensor_id,
                "sensor_type": frame.sensor.sensor_type,
                "frame_id": frame.sensor.frame_id,
                "hardware_timestamp_us": frame.hardware_timestamp_us,
                "sequence_id": frame.sequence_id,
                "active_job_id": frame.active_job_id,
                "capture_intent": frame.capture_intent,
            },
            "payload_type": frame.sensor.payload_type,
            "stream_id": f"{frame.sensor.sensor_id}:{frame.sequence_id}",
            "fragment_index": fragment_index,
            "total_fragments": total_fragments,
            "is_final_fragment": fragment_index + 1 >= total_fragments,
            "uncompressed_size_bytes": frame.uncompressed_size_bytes,
            "payload_digest": frame.digest,
        }
        sensor_frame.update(payload)
        sensor_frame["fragment_payload_bytes"] = len(fragment)
        sensor_frame["received_age_ms"] = (time.monotonic() - frame.received_monotonic) * 1000.0
        return sensor_frame

    def handle_request(self, msg_type: int, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]:
        if msg_type == SENSOR_GET_READINESS_REQ:
            return SENSOR_GET_READINESS_RSP, self.readiness_payload()
        if msg_type == SENSOR_GET_LATEST_FRAME_REQ:
            return SENSOR_GET_LATEST_FRAME_RSP, self.frame_payload(payload)
        if msg_type == SENSOR_LIST_SENSORS_REQ:
            return SENSOR_LIST_SENSORS_RSP, self.list_sensors_payload()
        return SENSOR_GET_READINESS_RSP, {
            "success": False,
            "error_code": {"code": ERROR_INVALID_ARGUMENT},
            "message": f"unsupported sensor msg_type={msg_type}",
        }


class AgentTcpHandler(socketserver.BaseRequestHandler):
    def handle(self) -> None:
        agent: IsaacVehicleSensorAgentSim = 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 = SENSOR_GET_READINESS_RSP
            response_payload = {
                "success": False,
                "error_code": {"code": ERROR_INVALID_ARGUMENT},
                "message": str(exc),
            }
        send_frame(self.request, response_type, response_payload)


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


def create_server(agent: IsaacVehicleSensorAgentSim, host: str, port: int) -> ThreadingTcpServer:
    server = ThreadingTcpServer((host, port), AgentTcpHandler)
    server.agent = agent
    return server


def serve(server: ThreadingTcpServer) -> None:
    host, port = server.server_address
    print(f"[*] sensor vehicle agent sim listening on {host}:{port}")
    server.serve_forever()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Isaac vehicle-side sensor agent simulator")
    parser.add_argument("--bind-host", default="0.0.0.0")
    parser.add_argument("--sensor-port", type=int, default=9003)
    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.parse_args()


def main() -> int:
    args = parse_args()
    rclpy.init(args=None)
    agent = IsaacVehicleSensorAgentSim(args)
    server = None
    try:
        server = create_server(agent, args.bind_host, args.sensor_port)
        threading.Thread(target=serve, args=(server,), daemon=True).start()
        print("[*] subscribed vehicle sensor topics:")
        for sensor in agent.sensors.values():
            print(f"  - {sensor.sensor_id}: {sensor.topic}")
        while rclpy.ok():
            rclpy.spin_once(agent.node, timeout_sec=0.1)
    except OSError as exc:
        print(f"[ERROR] vehicle_sensor_agent_sim network startup failed: {exc}")
        return 1
    except KeyboardInterrupt:
        print("[*] vehicle_sensor_agent_sim shutdown requested.")
    finally:
        if server is not None:
            server.shutdown()
            server.server_close()
        agent.node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()
    return 0


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