#!/usr/bin/env python3
"""Workshop-side sensor ingest for the simulated WiFi6 TCP vehicle link."""

from __future__ import annotations

import argparse
import base64
import json
import socket
import struct
import time
from typing import Any

import rclpy
from rclpy.executors import ExternalShutdownException
from sensor_msgs.msg import Image, Imu, LaserScan, PointCloud2, PointField


SENSOR_GET_LATEST_FRAME_REQ = 23
SENSOR_GET_LATEST_FRAME_RSP = 24
SENSOR_LIST_SENSORS_REQ = 25
SENSOR_LIST_SENSORS_RSP = 26

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


def is_shutdown_context_error(exc: Exception) -> bool:
    text = str(exc)
    return "publisher's context is invalid" in text or "context is invalid" in text


def read_exactly(conn: socket.socket, size: int) -> bytes:
    chunks: list[bytes] = []
    remaining = size
    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_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 = read_exactly(conn, 8)
        rsp_type, payload_len = struct.unpack("<II", header)
        raw_payload = read_exactly(conn, payload_len) if payload_len else b"{}"
    return rsp_type, json.loads(raw_payload.decode("utf-8"))


def sanitize_topic_part(value: str) -> str:
    cleaned = "".join(ch if ch.isalnum() or ch == "_" else "_" for ch in value)
    return cleaned.strip("_") or "sensor"


def topic_for(prefix: str, sensor_id: str, payload_type: int) -> str:
    suffix_by_payload = {
        PAYLOAD_IMAGE: "image_raw",
        PAYLOAD_POINT_CLOUD: "pointcloud",
        PAYLOAD_LASER_SCAN: "scan",
        PAYLOAD_IMU: "data",
    }
    suffix = suffix_by_payload.get(payload_type, "frame")
    return f"{prefix.rstrip('/')}/{sanitize_topic_part(sensor_id)}/{suffix}"


def stamp_from_us(msg: Any, timestamp_us: int, frame_id: str) -> None:
    msg.header.frame_id = frame_id
    msg.header.stamp.sec = int(timestamp_us // 1_000_000)
    msg.header.stamp.nanosec = int((timestamp_us % 1_000_000) * 1000)


def decode_data(payload: dict[str, Any], field_name: str = "data") -> bytes:
    encoded = payload.get(field_name, "")
    if not encoded:
        return b""
    return base64.b64decode(encoded.encode("ascii"))


def fixed_covariance(values: Any) -> list[float]:
    result = [float(value) for value in list(values or [])[:9]]
    while len(result) < 9:
        result.append(0.0)
    return result


class WorkshopSensorIngestSim:
    def __init__(self, args: argparse.Namespace):
        self.args = args
        self.node = rclpy.create_node("workshop_sensor_ingest_sim")
        self.publishers: dict[str, Any] = {}
        self.sensors: dict[str, dict[str, Any]] = {}
        self.last_sensor_refresh_monotonic = 0.0
        self.last_error_log_monotonic = 0.0
        self.published_count = 0
        self.fail_count = 0
        self.timer = self.node.create_timer(1.0 / max(args.poll_hz, 0.1), self.poll_once)
        self.node.get_logger().info(
            f"workshop sensor ingest: {args.sensor_host}:{args.sensor_port} -> {args.publish_prefix}"
        )

    def log_error(self, message: str) -> None:
        if not rclpy.ok():
            return
        self.fail_count += 1
        now = time.monotonic()
        if now - self.last_error_log_monotonic >= self.args.error_log_interval_sec:
            self.last_error_log_monotonic = now
            self.node.get_logger().warning(f"{message}; fail_count={self.fail_count}")

    def refresh_sensors(self) -> None:
        now = time.monotonic()
        if self.sensors and now - self.last_sensor_refresh_monotonic < self.args.sensor_refresh_interval_sec:
            return
        self.last_sensor_refresh_monotonic = now

        rsp_type, payload = send_request(
            self.args.sensor_host,
            self.args.sensor_port,
            SENSOR_LIST_SENSORS_REQ,
            {},
            self.args.timeout_sec,
        )
        if rsp_type != SENSOR_LIST_SENSORS_RSP or not payload.get("success", False):
            raise RuntimeError(f"sensor list failed: type={rsp_type}, payload={payload}")

        allowed = set(self.args.sensor_id)
        sensors = {}
        for sensor in payload.get("sensors", []):
            sensor_id = str(sensor.get("sensor_id", ""))
            if not sensor_id or (allowed and sensor_id not in allowed):
                continue
            if not bool(sensor.get("enabled", True)):
                continue
            sensors[sensor_id] = sensor
            self.ensure_publisher(sensor_id, int(sensor.get("payload_type", 0) or 0))

        self.sensors = sensors

    def ensure_publisher(self, sensor_id: str, payload_type: int) -> None:
        if sensor_id in self.publishers:
            return
        topic = topic_for(self.args.publish_prefix, sensor_id, payload_type)
        if payload_type == PAYLOAD_IMAGE:
            publisher = self.node.create_publisher(Image, topic, 10)
        elif payload_type == PAYLOAD_POINT_CLOUD:
            publisher = self.node.create_publisher(PointCloud2, topic, 10)
        elif payload_type == PAYLOAD_LASER_SCAN:
            publisher = self.node.create_publisher(LaserScan, topic, 10)
        elif payload_type == PAYLOAD_IMU:
            publisher = self.node.create_publisher(Imu, topic, 10)
        else:
            self.node.get_logger().warning(f"unsupported sensor payload_type={payload_type} for sensor_id={sensor_id}")
            return
        self.publishers[sensor_id] = publisher
        self.node.get_logger().info(f"publish {sensor_id} frames to {topic}")

    def poll_once(self) -> None:
        if not rclpy.ok():
            return
        try:
            self.refresh_sensors()
            for sensor_id in list(self.sensors):
                if not rclpy.ok():
                    return
                frame = self.fetch_latest_frame(sensor_id)
                self.publish_frame(frame)
        except Exception as exc:
            if not rclpy.ok() or is_shutdown_context_error(exc):
                return
            self.log_error(str(exc))

    def fetch_latest_frame(self, sensor_id: str) -> dict[str, Any]:
        first = self.request_frame_fragment(sensor_id, 0)
        total_fragments = int(first.get("total_fragments", 1) or 1)
        if total_fragments <= 1:
            return first

        payload_type = int(first.get("payload_type", 0) or 0)
        if payload_type not in (PAYLOAD_IMAGE, PAYLOAD_POINT_CLOUD):
            return first

        stream_id = str(first.get("stream_id", ""))
        raw_parts = [self.extract_raw_payload(first)]
        for fragment_index in range(1, total_fragments):
            fragment = self.request_frame_fragment(sensor_id, fragment_index)
            if str(fragment.get("stream_id", "")) != stream_id:
                raise RuntimeError(f"sensor_id={sensor_id} fragment stream changed while reassembling")
            raw_parts.append(self.extract_raw_payload(fragment))
        self.inject_raw_payload(first, b"".join(raw_parts))
        first["fragment_index"] = 0
        first["total_fragments"] = 1
        first["is_final_fragment"] = True
        return first

    def request_frame_fragment(self, sensor_id: str, fragment_index: int) -> dict[str, Any]:
        rsp_type, payload = send_request(
            self.args.sensor_host,
            self.args.sensor_port,
            SENSOR_GET_LATEST_FRAME_REQ,
            {
                "sensor_id": sensor_id,
                "include_payload": True,
                "max_payload_bytes": self.args.max_payload_bytes,
                "fragment_index": fragment_index,
            },
            self.args.timeout_sec,
        )
        if rsp_type != SENSOR_GET_LATEST_FRAME_RSP or not payload.get("success", False):
            raise RuntimeError(f"latest frame failed for {sensor_id}: type={rsp_type}, payload={payload}")
        return payload.get("sensor_frame", {})

    @staticmethod
    def extract_raw_payload(frame: dict[str, Any]) -> bytes:
        payload_type = int(frame.get("payload_type", 0) or 0)
        if payload_type == PAYLOAD_IMAGE:
            return decode_data(frame.get("image", {}))
        if payload_type == PAYLOAD_POINT_CLOUD:
            return decode_data(frame.get("point_cloud", {}))
        return b""

    @staticmethod
    def inject_raw_payload(frame: dict[str, Any], raw: bytes) -> None:
        payload_type = int(frame.get("payload_type", 0) or 0)
        encoded = base64.b64encode(raw).decode("ascii")
        if payload_type == PAYLOAD_IMAGE:
            frame.setdefault("image", {})["data"] = encoded
        elif payload_type == PAYLOAD_POINT_CLOUD:
            frame.setdefault("point_cloud", {})["data"] = encoded

    def publish_frame(self, frame: dict[str, Any]) -> None:
        header = frame.get("header", {})
        sensor_id = str(header.get("sensor_id", ""))
        payload_type = int(frame.get("payload_type", 0) or 0)
        publisher = self.publishers.get(sensor_id)
        if publisher is None:
            self.ensure_publisher(sensor_id, payload_type)
            publisher = self.publishers.get(sensor_id)
        if publisher is None:
            return

        if payload_type == PAYLOAD_IMAGE:
            publisher.publish(self.to_image(frame))
        elif payload_type == PAYLOAD_POINT_CLOUD:
            publisher.publish(self.to_point_cloud(frame))
        elif payload_type == PAYLOAD_LASER_SCAN:
            publisher.publish(self.to_laser_scan(frame))
        elif payload_type == PAYLOAD_IMU:
            publisher.publish(self.to_imu(frame))
        else:
            raise RuntimeError(f"unsupported payload_type={payload_type} for sensor_id={sensor_id}")
        self.published_count += 1

    def to_image(self, frame: dict[str, Any]) -> Image:
        header = frame.get("header", {})
        payload = frame.get("image", {})
        msg = Image()
        stamp_from_us(msg, int(header.get("hardware_timestamp_us", 0) or 0), str(header.get("frame_id", "")))
        msg.width = int(payload.get("width", 0) or 0)
        msg.height = int(payload.get("height", 0) or 0)
        msg.encoding = str(payload.get("encoding", ""))
        msg.is_bigendian = int(1 if payload.get("is_bigendian", False) else 0)
        msg.step = int(payload.get("step", 0) or 0)
        msg.data = decode_data(payload)
        return msg

    def to_point_cloud(self, frame: dict[str, Any]) -> PointCloud2:
        header = frame.get("header", {})
        payload = frame.get("point_cloud", {})
        msg = PointCloud2()
        stamp_from_us(msg, int(header.get("hardware_timestamp_us", 0) or 0), str(header.get("frame_id", "")))
        msg.height = int(payload.get("height", 0) or 0)
        msg.width = int(payload.get("width", 0) or 0)
        msg.fields = [
            PointField(
                name=str(field.get("name", "")),
                offset=int(field.get("offset", 0) or 0),
                datatype=int(field.get("datatype", 0) or 0),
                count=int(field.get("count", 0) or 0),
            )
            for field in payload.get("fields", [])
        ]
        msg.is_bigendian = bool(payload.get("is_bigendian", False))
        msg.point_step = int(payload.get("point_step", 0) or 0)
        msg.row_step = int(payload.get("row_step", 0) or 0)
        msg.is_dense = bool(payload.get("is_dense", True))
        msg.data = decode_data(payload)
        return msg

    def to_laser_scan(self, frame: dict[str, Any]) -> LaserScan:
        header = frame.get("header", {})
        payload = frame.get("laser_scan", {})
        msg = LaserScan()
        stamp_from_us(msg, int(header.get("hardware_timestamp_us", 0) or 0), str(header.get("frame_id", "")))
        msg.angle_min = float(payload.get("angle_min_rad", 0.0) or 0.0)
        msg.angle_max = float(payload.get("angle_max_rad", 0.0) or 0.0)
        msg.angle_increment = float(payload.get("angle_increment_rad", 0.0) or 0.0)
        msg.time_increment = float(payload.get("time_increment_sec", 0.0) or 0.0)
        msg.scan_time = float(payload.get("scan_time_sec", 0.0) or 0.0)
        msg.range_min = float(payload.get("range_min_m", 0.0) or 0.0)
        msg.range_max = float(payload.get("range_max_m", 0.0) or 0.0)
        msg.ranges = [float(value) for value in payload.get("ranges_m", [])]
        msg.intensities = [float(value) for value in payload.get("intensities", [])]
        return msg

    def to_imu(self, frame: dict[str, Any]) -> Imu:
        header = frame.get("header", {})
        payload = frame.get("imu", {})
        msg = Imu()
        stamp_from_us(msg, int(header.get("hardware_timestamp_us", 0) or 0), str(header.get("frame_id", "")))
        msg.orientation.x = float(payload.get("orientation_x", 0.0) or 0.0)
        msg.orientation.y = float(payload.get("orientation_y", 0.0) or 0.0)
        msg.orientation.z = float(payload.get("orientation_z", 0.0) or 0.0)
        msg.orientation.w = float(payload.get("orientation_w", 1.0) or 1.0)
        msg.orientation_covariance = fixed_covariance(payload.get("orientation_covariance", []))
        angular_velocity = payload.get("angular_velocity", {})
        msg.angular_velocity.x = float(angular_velocity.get("x", 0.0) or 0.0)
        msg.angular_velocity.y = float(angular_velocity.get("y", 0.0) or 0.0)
        msg.angular_velocity.z = float(angular_velocity.get("z", 0.0) or 0.0)
        msg.angular_velocity_covariance = fixed_covariance(payload.get("angular_velocity_covariance", []))
        linear_acceleration = payload.get("linear_acceleration", {})
        msg.linear_acceleration.x = float(linear_acceleration.get("x", 0.0) or 0.0)
        msg.linear_acceleration.y = float(linear_acceleration.get("y", 0.0) or 0.0)
        msg.linear_acceleration.z = float(linear_acceleration.get("z", 0.0) or 0.0)
        msg.linear_acceleration_covariance = fixed_covariance(payload.get("linear_acceleration_covariance", []))
        return msg


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="车间侧传感器 WiFi6/TCP 接入仿真")
    parser.add_argument("--sensor-host", default="127.0.0.1")
    parser.add_argument("--sensor-port", type=int, default=9000)
    parser.add_argument("--publish-prefix", default="/workshop/vehicle_sensor")
    parser.add_argument("--poll-hz", type=float, default=15.0)
    parser.add_argument("--max-payload-bytes", type=int, default=4 * 1024 * 1024)
    parser.add_argument("--timeout-sec", type=float, default=2.0)
    parser.add_argument("--sensor-id", action="append", default=[], help="只拉取指定 sensor_id，可重复")
    parser.add_argument("--sensor-refresh-interval-sec", type=float, default=5.0)
    parser.add_argument("--error-log-interval-sec", type=float, default=2.0)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    rclpy.init(args=None)
    ingest = None
    try:
        ingest = WorkshopSensorIngestSim(args)
        rclpy.spin(ingest.node)
    except (KeyboardInterrupt, ExternalShutdownException):
        pass
    finally:
        if ingest is not None:
            ingest.node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()
    return 0


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