#!/usr/bin/env python3
"""ROS2 TF to browser bridge with optional static map image overlay."""

from __future__ import annotations

import argparse
import struct
import json
import mimetypes
import queue
import signal
import threading
import time
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit


ROOT = Path(__file__).resolve().parent
CLIENTS: list[queue.Queue[str]] = []
CLIENTS_LOCK = threading.Lock()
MAP_ASSET: "MapAsset | None" = None
LATEST: dict[str, Any] = {
    "tf": None,
    "tf_static": None,
    "map": None,
    "bridge_status": {"ok": False, "message": "ROS thread not started"},
    "connected_at": time.time(),
}


@dataclass(frozen=True)
class MapAsset:
    path: Path
    content_type: str
    metadata: dict[str, Any]


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


def stamp_to_float(stamp: Any) -> float | None:
    if stamp is None:
        return None
    return float(getattr(stamp, "sec", 0)) + float(getattr(stamp, "nanosec", 0)) / 1_000_000_000.0


def vector_to_dict(value: Any) -> dict[str, float]:
    return {
        "x": float(getattr(value, "x", 0.0)),
        "y": float(getattr(value, "y", 0.0)),
        "z": float(getattr(value, "z", 0.0)),
    }


def quaternion_to_dict(value: Any) -> dict[str, float]:
    return {
        "x": float(getattr(value, "x", 0.0)),
        "y": float(getattr(value, "y", 0.0)),
        "z": float(getattr(value, "z", 0.0)),
        "w": float(getattr(value, "w", 1.0)),
    }


def tf_message_to_dict(msg: Any) -> dict[str, Any]:
    transforms = []
    for transform in msg.transforms:
        header = transform.header
        transforms.append(
            {
                "parent": str(header.frame_id),
                "child": str(transform.child_frame_id),
                "stamp": stamp_to_float(getattr(header, "stamp", None)),
                "translation": vector_to_dict(transform.transform.translation),
                "rotation": quaternion_to_dict(transform.transform.rotation),
            }
        )
    return {"count": len(transforms), "transforms": transforms}


def map_message_to_dict(msg: Any) -> dict[str, Any]:
    return {
        "header": {
            "frame_id": msg.header.frame_id,
            "stamp": stamp_to_float(msg.header.stamp),
        },
        "info": {
            "width": int(msg.info.width),
            "height": int(msg.info.height),
            "resolution": float(msg.info.resolution),
            "origin": {
                "position": vector_to_dict(msg.info.origin.position),
                "orientation": quaternion_to_dict(msg.info.origin.orientation),
            },
        },
    }


def parse_yaml_scalar(value: str) -> Any:
    value = value.strip().strip('"').strip("'")
    if value.startswith("[") and value.endswith("]"):
        return [parse_yaml_scalar(part) for part in value[1:-1].split(",")]
    try:
        return int(value)
    except ValueError:
        pass
    try:
        return float(value)
    except ValueError:
        return value


def load_simple_yaml(path: Path) -> dict[str, Any]:
    data: dict[str, Any] = {}
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.split("#", 1)[0].strip()
        if not line or ":" not in line:
            continue
        key, value = line.split(":", 1)
        data[key.strip()] = parse_yaml_scalar(value)
    return data


def read_png_size(path: Path) -> tuple[int, int] | None:
    data = path.read_bytes()[:24]
    if len(data) >= 24 and data.startswith(b"\x89PNG\r\n\x1a\n"):
        return int.from_bytes(data[16:20], "big"), int.from_bytes(data[20:24], "big")
    return None


def read_jpeg_size(path: Path) -> tuple[int, int] | None:
    data = path.read_bytes()
    if not data.startswith(b"\xff\xd8"):
        return None
    index = 2
    while index + 9 < len(data):
        if data[index] != 0xFF:
            index += 1
            continue
        marker = data[index + 1]
        index += 2
        if marker in (0xD8, 0xD9):
            continue
        if index + 2 > len(data):
            return None
        length = struct.unpack(">H", data[index : index + 2])[0]
        if marker in range(0xC0, 0xC4) and index + 7 < len(data):
            height = struct.unpack(">H", data[index + 3 : index + 5])[0]
            width = struct.unpack(">H", data[index + 5 : index + 7])[0]
            return width, height
        index += length
    return None


def read_image_size(path: Path) -> tuple[int, int] | None:
    return read_png_size(path) or read_jpeg_size(path)


def detect_content_type(path: Path) -> str:
    header = path.read_bytes()[:12]
    if header.startswith(b"\x89PNG\r\n\x1a\n"):
        return "image/png"
    if header.startswith(b"\xff\xd8"):
        return "image/jpeg"
    return mimetypes.guess_type(path.name)[0] or "application/octet-stream"


def load_map_asset(
    image_path: str | None,
    yaml_path: str | None = None,
    map_width_meters: float = 5.0,
) -> MapAsset | None:
    if not image_path:
        return None
    path = Path(image_path).expanduser().resolve()
    if not path.is_file():
        raise FileNotFoundError(f"map image not found: {path}")
    content_type = detect_content_type(path)
    metadata: dict[str, Any] = {"image": path.name, "placement": "publish_map_windows"}
    image_size = read_image_size(path)
    if image_size is not None:
        width, height = image_size
        metadata.update(
            {
                "width": width,
                "height": height,
                "resolution": float(map_width_meters) / float(width),
                "origin": [0.0, 0.0, 0.0],
            }
        )
    metadata_path = Path(yaml_path).expanduser().resolve() if yaml_path else path.with_suffix(".yaml")
    if metadata_path.is_file():
        metadata.update(load_simple_yaml(metadata_path))
    return MapAsset(path=path, content_type=content_type, metadata=metadata)


def publish_event(event_type: str, data: Any) -> None:
    LATEST[event_type] = data
    payload = json.dumps(
        {"type": event_type, "data": data, "time": time.time()},
        ensure_ascii=False,
        separators=(",", ":"),
    )
    with CLIENTS_LOCK:
        clients = list(CLIENTS)
    for client in clients:
        client.put(payload)


class TfRequestHandler(BaseHTTPRequestHandler):
    server_version = "TfWebBridge/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_file(ROOT / "tf_viewer.html", "text/html; charset=utf-8")
            return
        if path == "/events":
            self.serve_events()
            return
        if path == "/latest":
            self.serve_json(LATEST)
            return
        if path == "/map-image" and MAP_ASSET is not None:
            self.serve_file(MAP_ASSET.path, MAP_ASSET.content_type)
            return
        self.send_error(404)

    def serve_file(self, path: Path, content_type: str) -> None:
        data = path.read_bytes()
        self.send_response(200)
        self.send_header("Content-Type", content_type)
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    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("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def serve_events(self) -> None:
        client: queue.Queue[str] = queue.Queue()
        with CLIENTS_LOCK:
            CLIENTS.append(client)

        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("Connection", "keep-alive")
        self.end_headers()

        try:
            client.put(json.dumps({"type": "snapshot", "data": LATEST, "time": time.time()}))
            while True:
                try:
                    payload = client.get(timeout=10.0)
                    self.wfile.write(f"data: {payload}\n\n".encode("utf-8"))
                except queue.Empty:
                    self.wfile.write(b": keepalive\n\n")
                self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError):
            pass
        finally:
            with CLIENTS_LOCK:
                if client in CLIENTS:
                    CLIENTS.remove(client)


class ReusableThreadingHTTPServer(ThreadingHTTPServer):
    allow_reuse_address = True


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


def spin_ros(tf_topic: str, tf_static_topic: str, map_topic: str | None) -> None:
    try:
        import rclpy
        from nav_msgs.msg import OccupancyGrid
        from rclpy.node import Node
        from rclpy.qos import DurabilityPolicy, QoSProfile, ReliabilityPolicy
        from tf2_msgs.msg import TFMessage
    except Exception as exc:
        message = f"ROS imports failed: {exc}"
        print(message, flush=True)
        publish_event("bridge_status", {"ok": False, "message": message})
        return

    class BridgeNode(Node):
        def __init__(self) -> None:
            super().__init__("tf_web_bridge")
            tf_qos = QoSProfile(depth=100)
            tf_qos.reliability = ReliabilityPolicy.RELIABLE
            static_qos = QoSProfile(depth=1)
            static_qos.reliability = ReliabilityPolicy.RELIABLE
            static_qos.durability = DurabilityPolicy.TRANSIENT_LOCAL
            self.create_subscription(TFMessage, tf_topic, self.on_tf, tf_qos)
            self.create_subscription(TFMessage, tf_static_topic, self.on_tf_static, static_qos)
            if map_topic:
                self.create_subscription(OccupancyGrid, map_topic, self.on_map, static_qos)

        def on_tf(self, msg: Any) -> None:
            publish_event("tf", tf_message_to_dict(msg))

        def on_tf_static(self, msg: Any) -> None:
            publish_event("tf_static", tf_message_to_dict(msg))

        def on_map(self, msg: Any) -> None:
            publish_event("map", map_message_to_dict(msg))

    rclpy.init()
    node = BridgeNode()
    publish_event("bridge_status", {"ok": True, "message": "subscribed"})
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


def bind_help() -> str:
    return "bind address, default 127.0.0.1; use 192.168.10.164 to expose on that interface"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1", help=bind_help())
    parser.add_argument("--port", type=int, default=8770)
    parser.add_argument("--tf-topic", default="/tf")
    parser.add_argument("--tf-static-topic", default="/tf_static")
    parser.add_argument("--map-topic", default="/map")
    parser.add_argument("--no-map-topic", action="store_true")
    parser.add_argument("--map-image", default=None)
    parser.add_argument("--map-yaml", default=None)
    parser.add_argument("--map-width-meters", type=float, default=5.0)
    args = parser.parse_args()

    global MAP_ASSET
    MAP_ASSET = load_map_asset(args.map_image, args.map_yaml, args.map_width_meters)
    if MAP_ASSET:
        LATEST["map_asset"] = {**MAP_ASSET.metadata, "url": "/map-image"}

    server = start_http(args.host, args.port)
    print(f"tf web bridge: http://{args.host}:{args.port}", flush=True)
    if MAP_ASSET:
        print(f"map image: {MAP_ASSET.path}", flush=True)
    if args.host == "127.0.0.1":
        print("bind: local loopback only; add --host 192.168.10.164 to expose that interface", flush=True)

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

    map_topic = None if args.no_map_topic else args.map_topic
    ros_thread = threading.Thread(
        target=spin_ros,
        args=(args.tf_topic, args.tf_static_topic, map_topic),
        daemon=True,
    )
    ros_thread.start()

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


if __name__ == "__main__":
    main()

