#!/usr/bin/env python3
"""Publish or receive the complete robot model used by MVT visualization."""

import argparse
import hashlib
import json
import os
import sys
import xml.etree.ElementTree as ET
from pathlib import Path

import rclpy
from rclpy.node import Node
from rclpy.qos import DurabilityPolicy, HistoryPolicy, QoSProfile, ReliabilityPolicy

from stardust_mvt_protocol.msg import RobotDescriptionSnapshot
from stardust_mvt_core.robot_description_resolver import (
    local_robot_description_paths,
    validate_arm_urdf,
)


TOPIC_NAME = "/mvt_core/robot_description_snapshot"
SCHEMA_VERSION = 1


def _sha256(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _snapshot_qos() -> QoSProfile:
    return QoSProfile(
        history=HistoryPolicy.KEEP_LAST,
        depth=1,
        reliability=ReliabilityPolicy.RELIABLE,
        durability=DurabilityPolicy.TRANSIENT_LOCAL,
    )


def _atomic_write(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temp_path = path.with_name(f"{path.name}.tmp")
    temp_path.write_text(content, encoding="utf-8")
    os.replace(temp_path, path)


def _validate_snapshot(message: RobotDescriptionSnapshot) -> dict:
    if message.schema_version != SCHEMA_VERSION:
        raise RuntimeError(
            f"unsupported robot description snapshot schema: {message.schema_version}"
        )
    if not message.urdf or not message.srdf:
        raise RuntimeError("robot description snapshot has empty URDF or SRDF")
    if _sha256(message.urdf) != message.urdf_sha256:
        raise RuntimeError("robot description snapshot URDF sha256 mismatch")
    if _sha256(message.srdf) != message.srdf_sha256:
        raise RuntimeError("robot description snapshot SRDF sha256 mismatch")

    errors = validate_arm_urdf(message.urdf)
    if errors:
        raise RuntimeError("robot description snapshot URDF invalid: " + "; ".join(errors))
    try:
        srdf_root = ET.fromstring(message.srdf)
    except ET.ParseError as exc:
        raise RuntimeError(f"robot description snapshot SRDF XML parse failed: {exc}") from exc
    if srdf_root.tag != "robot":
        raise RuntimeError("robot description snapshot SRDF root must be <robot>")
    try:
        metadata = json.loads(message.metadata_json)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"robot description snapshot metadata is not JSON: {exc}") from exc
    if not isinstance(metadata, dict):
        raise RuntimeError("robot description snapshot metadata must be an object")
    metadata_hash = metadata.get("urdf_hash")
    if metadata_hash is not None and metadata_hash != message.urdf_sha256:
        raise RuntimeError(
            "robot description snapshot metadata urdf_hash does not match payload"
        )
    return metadata


class SnapshotPublisher(Node):
    def __init__(self):
        super().__init__("robot_description_snapshot_publisher", namespace="/mvt_core")
        self.declare_parameter("topic_name", TOPIC_NAME)
        self.declare_parameter("urdf", "")
        self.declare_parameter("srdf", "")
        self.declare_parameter("metadata_json", "{}")
        topic_name = str(self.get_parameter("topic_name").value)
        urdf = str(self.get_parameter("urdf").value)
        srdf = str(self.get_parameter("srdf").value)
        metadata_json = str(self.get_parameter("metadata_json").value)

        metadata = json.loads(metadata_json)
        if not isinstance(metadata, dict):
            raise RuntimeError("metadata_json must encode an object")
        metadata["urdf_hash"] = _sha256(urdf)
        self.message = RobotDescriptionSnapshot(
            schema_version=SCHEMA_VERSION,
            generated_at=self.get_clock().now().to_msg(),
            urdf=urdf,
            srdf=srdf,
            metadata_json=json.dumps(metadata, separators=(",", ":"), sort_keys=True),
            urdf_sha256=_sha256(urdf),
            srdf_sha256=_sha256(srdf),
        )
        _validate_snapshot(self.message)
        self.publisher = self.create_publisher(
            RobotDescriptionSnapshot, topic_name, _snapshot_qos()
        )
        self.timer = self.create_timer(0.1, self._publish_once)

    def _publish_once(self) -> None:
        self.publisher.publish(self.message)
        self.timer.cancel()
        self.get_logger().info(
            "published robot description snapshot "
            f"urdf_sha256={self.message.urdf_sha256}"
        )


class SnapshotWaiter(Node):
    def __init__(self, runtime_workspace_dir: str, timeout_sec: float):
        super().__init__("robot_description_snapshot_waiter", namespace="/mvt_core")
        self.timeout_sec = timeout_sec
        self.result = None
        self.error = None
        local_paths = local_robot_description_paths(
            workspace_dir=runtime_workspace_dir
        )
        self.output_dir = Path(str(local_paths["robot_description_dir"]))
        self.subscription = self.create_subscription(
            RobotDescriptionSnapshot,
            TOPIC_NAME,
            self._callback,
            _snapshot_qos(),
        )

    def _callback(self, message: RobotDescriptionSnapshot) -> None:
        if self.result is not None or self.error is not None:
            return
        try:
            metadata = _validate_snapshot(message)
            _atomic_write(
                self.output_dir / "resolved_robot_description.urdf", message.urdf
            )
            _atomic_write(
                self.output_dir / "resolved_robot_description_semantic.srdf",
                message.srdf,
            )
            _atomic_write(
                self.output_dir / "resolved_robot_description.json",
                json.dumps(metadata, indent=2, sort_keys=True) + "\n",
            )
            self.result = message
            self.get_logger().info(
                f"robot description snapshot cached in {self.output_dir}"
            )
        except Exception as exc:  # noqa: BLE001 - report validation to CLI
            self.error = exc


def _wait_for_snapshot(runtime_workspace_dir: str, timeout_sec: float) -> int:
    rclpy.init(args=None)
    node = SnapshotWaiter(runtime_workspace_dir, timeout_sec)
    try:
        deadline = node.get_clock().now().nanoseconds + int(timeout_sec * 1e9)
        while rclpy.ok() and node.result is None and node.error is None:
            if node.get_clock().now().nanoseconds >= deadline:
                break
            rclpy.spin_once(node, timeout_sec=0.2)
        if node.error is not None:
            print(f"robot description snapshot rejected: {node.error}", file=sys.stderr)
            return 1
        if node.result is None:
            print(
                f"timed out after {timeout_sec:.1f}s waiting for {TOPIC_NAME}",
                file=sys.stderr,
            )
            return 1
        return 0
    finally:
        node.destroy_node()
        rclpy.shutdown()


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=("publisher", "wait"), default="publisher")
    parser.add_argument("--runtime-workspace-dir", default="")
    parser.add_argument("--timeout-sec", type=float, default=10.0)
    return parser.parse_known_args()


def main() -> int:
    args, ros_args = _parse_args()
    if args.mode == "wait":
        if not args.runtime_workspace_dir:
            print("--runtime-workspace-dir is required in wait mode", file=sys.stderr)
            return 2
        return _wait_for_snapshot(args.runtime_workspace_dir, args.timeout_sec)

    rclpy.init(args=ros_args)
    node = SnapshotPublisher()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()
    return 0


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