#!/usr/bin/env python3
"""ROS2 image topic to MJPEG web bridge.

Supports the image publishing modes used by car_usb_cam:
- sensor_msgs/msg/Image on /image
- sensor_msgs/msg/CompressedImage style JPEG topics
- hbm_img_msgs/msg/HbmMsg1080P/540P/480P zero-copy payload messages
"""

from __future__ import annotations

import argparse
import importlib
import json
import signal
import threading
import time
from collections import deque
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs, urlsplit

import cv2
import numpy as np


ROOT = Path(__file__).resolve().parent
JPEG_LOCK = threading.Condition()
LATEST_JPEG: bytes | None = None
LATEST_META: dict[str, Any] = {
    "ok": False,
    "message": "ROS thread not started",
    "topic": None,
    "topic_type": None,
    "frames": 0,
    "fps": 0.0,
    "last_frame_time": None,
    "compress": False,
    "quality": None,
}


def normalize_encoding(value: Any) -> str:
    if isinstance(value, str):
        return value.strip("\x00").lower()
    if isinstance(value, (bytes, bytearray)):
        return bytes(value).split(b"\x00", 1)[0].decode("ascii", errors="ignore").lower()
    if isinstance(value, (list, tuple)):
        return bytes(int(item) & 0xFF for item in value).split(b"\x00", 1)[0].decode(
            "ascii", errors="ignore"
        ).lower()
    return str(value).strip("\x00").lower()


def _valid_bytes(data: Any, size: int | None = None) -> bytes:
    if isinstance(data, bytes):
        raw = data
    elif isinstance(data, bytearray):
        raw = bytes(data)
    else:
        raw = bytes(int(item) & 0xFF for item in data)
    return raw[:size] if size is not None else raw


def _rows_from_buffer(raw: bytes, height: int, step: int, row_bytes: int) -> np.ndarray:
    expected = height * step
    if len(raw) < expected:
        raise ValueError(f"image data too short: got {len(raw)} bytes, need {expected}")
    rows = np.frombuffer(raw[:expected], dtype=np.uint8).reshape((height, step))
    return rows[:, :row_bytes]


def image_message_to_bgr(msg: Any) -> np.ndarray:
    height = int(msg.height)
    width = int(msg.width)
    step = int(msg.step)
    encoding = normalize_encoding(msg.encoding)
    raw = _valid_bytes(msg.data)

    if encoding in ("rgb8", "bgr8"):
        rows = _rows_from_buffer(raw, height, step, width * 3)
        image = rows.reshape((height, width, 3))
        return cv2.cvtColor(image, cv2.COLOR_RGB2BGR) if encoding == "rgb8" else image

    if encoding in ("mono8", "8uc1"):
        rows = _rows_from_buffer(raw, height, step, width)
        return cv2.cvtColor(rows.reshape((height, width)), cv2.COLOR_GRAY2BGR)

    if encoding in ("rgba8", "bgra8"):
        rows = _rows_from_buffer(raw, height, step, width * 4)
        image = rows.reshape((height, width, 4))
        code = cv2.COLOR_RGBA2BGR if encoding == "rgba8" else cv2.COLOR_BGRA2BGR
        return cv2.cvtColor(image, code)

    if encoding in ("yuyv", "yuyv2rgb", "yuv422_yuy2"):
        rows = _rows_from_buffer(raw, height, step, width * 2)
        image = rows.reshape((height, width, 2))
        return cv2.cvtColor(image, cv2.COLOR_YUV2BGR_YUY2)

    if encoding == "nv12":
        expected = height * width * 3 // 2
        if len(raw) < expected:
            raise ValueError(f"nv12 data too short: got {len(raw)} bytes, need {expected}")
        image = np.frombuffer(raw[:expected], dtype=np.uint8).reshape((height * 3 // 2, width))
        return cv2.cvtColor(image, cv2.COLOR_YUV2BGR_NV12)

    raise ValueError(f"unsupported image encoding: {encoding or '<empty>'}")


def hbm_message_to_bgr(msg: Any) -> np.ndarray:
    data_size = int(getattr(msg, "data_size", 0)) or None
    image_like = type(
        "ImageLike",
        (),
        {
            "height": int(msg.height),
            "width": int(msg.width),
            "step": int(msg.step),
            "encoding": normalize_encoding(msg.encoding),
            "data": _valid_bytes(msg.data, data_size),
        },
    )()
    return image_message_to_bgr(image_like)


def encode_bgr_to_jpeg(image: np.ndarray, quality: int = 75) -> bytes:
    quality = max(1, min(100, int(quality)))
    ok, encoded = cv2.imencode(".jpg", image, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
    if not ok:
        raise ValueError("cv2 failed to encode jpeg")
    return encoded.tobytes()


def compressed_message_to_jpeg(msg: Any, quality: int = 75, force_compress: bool = False) -> bytes:
    data = _valid_bytes(msg.data)
    fmt = str(getattr(msg, "format", "")).lower()
    is_jpeg = "jpeg" in fmt or "jpg" in fmt or data.startswith(b"\xff\xd8")
    if is_jpeg and not force_compress:
        return data
    array = np.frombuffer(data, dtype=np.uint8)
    decoded = cv2.imdecode(array, cv2.IMREAD_COLOR)
    if decoded is None:
        raise ValueError(f"unsupported compressed image format: {fmt or '<empty>'}")
    return encode_bgr_to_jpeg(decoded, quality)


def message_to_jpeg(
    msg: Any,
    topic_type: str,
    quality: int = 75,
    force_compress: bool = False,
) -> bytes:
    if topic_type == "sensor_msgs/msg/CompressedImage":
        return compressed_message_to_jpeg(msg, quality, force_compress)
    if topic_type == "sensor_msgs/msg/Image":
        return encode_bgr_to_jpeg(image_message_to_bgr(msg), quality)
    if topic_type.startswith("hbm_img_msgs/msg/HbmMsg"):
        return encode_bgr_to_jpeg(hbm_message_to_bgr(msg), quality)
    raise ValueError(f"unsupported topic type: {topic_type}")


class FrameRateTracker:
    def __init__(self, window_seconds: float = 2.0) -> None:
        self.window_seconds = float(window_seconds)
        self.samples: deque[float] = deque()

    def record(self, timestamp: float | None = None) -> float:
        now = time.time() if timestamp is None else float(timestamp)
        self.samples.append(now)
        oldest = now - self.window_seconds
        while self.samples and self.samples[0] < oldest:
            self.samples.popleft()
        if len(self.samples) < 2:
            return 0.0
        elapsed = self.samples[-1] - self.samples[0]
        if elapsed <= 0.0:
            return 0.0
        return round((len(self.samples) - 1) / elapsed, 1)


def publish_frame(jpeg: bytes, topic: str, topic_type: str, fps: float) -> None:
    with JPEG_LOCK:
        global LATEST_JPEG
        LATEST_JPEG = jpeg
        LATEST_META.update(
            {
                "ok": True,
                "message": "streaming",
                "topic": topic,
                "topic_type": topic_type,
                "frames": int(LATEST_META.get("frames") or 0) + 1,
                "fps": float(fps),
                "last_frame_time": time.time(),
            }
        )
        JPEG_LOCK.notify_all()


def update_status(ok: bool, message: str, **extra: Any) -> None:
    with JPEG_LOCK:
        LATEST_META.update({"ok": ok, "message": message, **extra})
        JPEG_LOCK.notify_all()


def request_path(raw_path: str) -> str:
    return urlsplit(raw_path).path


def request_query(raw_path: str) -> dict[str, list[str]]:
    return parse_qs(urlsplit(raw_path).query)


class ImageRequestHandler(BaseHTTPRequestHandler):
    server_version = "ImageWebBridge/1.0"

    def log_message(self, fmt: str, *args: Any) -> None:
        return

    def do_GET(self) -> None:
        path = request_path(self.path)
        if path in ("/", "/index.html"):
            self.serve_index()
        elif path == "/stream.mjpg":
            self.serve_stream()
        elif path == "/snapshot.jpg":
            self.serve_snapshot()
        elif path == "/status":
            self.serve_json(LATEST_META)
        else:
            self.send_error(404)

    def serve_index(self) -> None:
        topic = LATEST_META.get("topic") or ""
        body = f"""<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>ROS Image Bridge</title>
  <style>
    html, body {{ margin: 0; height: 100%; background: #111; color: #eee; font-family: Arial, sans-serif; }}
    main {{ height: 100%; display: grid; grid-template-rows: auto 1fr; }}
    header {{ display: flex; gap: 18px; align-items: center; padding: 10px 14px; background: #1d1f23; font-size: 14px; }}
    img {{ width: 100%; height: 100%; object-fit: contain; background: #050505; }}
    code {{ color: #9ad; }}
  </style>
</head>
<body>
  <main>
    <header>
      <strong>ROS Image Bridge</strong>
      <span>topic <code id="topic">{topic}</code></span>
      <span id="status">connecting</span>
    </header>
    <img src="/stream.mjpg" alt="ROS image stream">
  </main>
  <script>
    async function tick() {{
      const res = await fetch('/status', {{ cache: 'no-store' }});
      const data = await res.json();
      document.getElementById('topic').textContent = data.topic || '';
      const fps = Number(data.fps || 0).toFixed(1);
      document.getElementById('status').textContent =
        (data.ok ? 'online' : 'waiting') + ' | fps ' + fps +
        ' | frames ' + (data.frames || 0) + ' | ' + data.message;
    }}
    setInterval(tick, 1000);
    tick();
  </script>
</body>
</html>""".encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def serve_json(self, data: Any) -> None:
        body = json.dumps(data, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def serve_snapshot(self) -> None:
        with JPEG_LOCK:
            if LATEST_JPEG is None:
                self.send_error(503, "no frame yet")
                return
            jpeg = LATEST_JPEG
        self.send_response(200)
        self.send_header("Content-Type", "image/jpeg")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(jpeg)))
        self.end_headers()
        self.wfile.write(jpeg)

    def serve_stream(self) -> None:
        boundary = "frame"
        self.send_response(200)
        self.send_header("Content-Type", f"multipart/x-mixed-replace; boundary={boundary}")
        self.send_header("Cache-Control", "no-store")
        self.end_headers()

        last_frame: bytes | None = None
        try:
            while True:
                with JPEG_LOCK:
                    JPEG_LOCK.wait_for(
                        lambda: LATEST_JPEG is not None and LATEST_JPEG is not last_frame,
                        timeout=2.0,
                    )
                    jpeg = LATEST_JPEG
                if jpeg is None or jpeg is last_frame:
                    self.wfile.write(b"\r\n")
                    self.wfile.flush()
                    continue
                last_frame = jpeg
                header = (
                    f"--{boundary}\r\n"
                    "Content-Type: image/jpeg\r\n"
                    f"Content-Length: {len(jpeg)}\r\n\r\n"
                ).encode("ascii")
                self.wfile.write(header)
                self.wfile.write(jpeg)
                self.wfile.write(b"\r\n")
                self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError):
            return


class ReusableThreadingHTTPServer(ThreadingHTTPServer):
    allow_reuse_address = True


def start_http(host: str, port: int) -> ThreadingHTTPServer:
    server = ReusableThreadingHTTPServer((host, port), ImageRequestHandler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return server


def import_message_class(topic_type: str) -> Any:
    parts = topic_type.split("/")
    if len(parts) != 3 or parts[1] != "msg":
        raise ValueError(f"expected ROS2 message type like pkg/msg/Name, got {topic_type}")
    module = importlib.import_module(f"{parts[0]}.msg")
    return getattr(module, parts[2])


def shutdown_rclpy_once(rclpy_module: Any) -> None:
    if rclpy_module.ok():
        rclpy_module.shutdown()


def discover_topic_type(node: Any, topic: str, timeout: float) -> str:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        for name, types in node.get_topic_names_and_types():
            if name == topic and types:
                return types[0]
        time.sleep(0.2)
    raise TimeoutError(f"topic {topic} not found within {timeout:.1f}s")


def spin_ros(
    topic: str,
    topic_type: str,
    quality: int,
    max_fps: float,
    discovery_timeout: float,
    force_compress: bool,
) -> None:
    try:
        import rclpy
        from rclpy.executors import ExternalShutdownException
        from rclpy.node import Node
    except Exception as exc:
        update_status(False, f"ROS imports failed: {exc}", topic=topic, topic_type=topic_type)
        return

    class BridgeNode(Node):
        def __init__(self) -> None:
            super().__init__("image_web_bridge")
            actual_type = topic_type
            if actual_type == "auto":
                actual_type = discover_topic_type(self, topic, discovery_timeout)
            msg_class = import_message_class(actual_type)
            self.actual_type = actual_type
            self.last_emit = 0.0
            self.fps = FrameRateTracker()
            self.create_subscription(msg_class, topic, self.on_image, 10)
            update_status(
                False,
                "subscribed, waiting for first frame",
                topic=topic,
                topic_type=actual_type,
                compress=force_compress,
                quality=quality,
            )

        def on_image(self, msg: Any) -> None:
            now = time.monotonic()
            if max_fps > 0 and now - self.last_emit < 1.0 / max_fps:
                return
            try:
                jpeg = message_to_jpeg(msg, self.actual_type, quality, force_compress)
            except Exception as exc:
                update_status(False, f"frame conversion failed: {exc}", topic=topic, topic_type=self.actual_type)
                return
            self.last_emit = now
            publish_frame(jpeg, topic, self.actual_type, self.fps.record())

    try:
        rclpy.init()
        node = BridgeNode()
    except Exception as exc:
        update_status(False, f"ROS setup failed: {exc}", topic=topic, topic_type=topic_type)
        return

    try:
        rclpy.spin(node)
    except ExternalShutdownException:
        pass
    finally:
        node.destroy_node()
        shutdown_rclpy_once(rclpy)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8080)
    parser.add_argument("--topic", default="/image_mjpeg")
    parser.add_argument("--topic-type", default="auto")
    parser.add_argument("--quality", type=int, default=75)
    parser.add_argument("--compress", action="store_true")
    parser.add_argument("--max-fps", type=float, default=15.0)
    parser.add_argument("--discovery-timeout", type=float, default=10.0)
    args = parser.parse_args()

    update_status(
        False,
        "HTTP started, waiting for ROS",
        topic=args.topic,
        topic_type=args.topic_type,
        compress=args.compress,
        quality=args.quality,
    )
    server = start_http(args.host, args.port)
    print(f"image web bridge: http://{args.host}:{args.port}", flush=True)
    print(f"topic: {args.topic} type: {args.topic_type}", flush=True)

    stop = threading.Event()
    signal.signal(signal.SIGINT, lambda *_: stop.set())
    signal.signal(signal.SIGTERM, lambda *_: stop.set())

    ros_thread = threading.Thread(
        target=spin_ros,
        args=(
            args.topic,
            args.topic_type,
            args.quality,
            args.max_fps,
            args.discovery_timeout,
            args.compress,
        ),
        daemon=True,
    )
    ros_thread.start()

    while not stop.is_set():
        time.sleep(0.2)
    server.shutdown()


if __name__ == "__main__":
    main()
