"""JSON serialization for SimpleSceneGraph.

Provides save_scene_graph() and load_scene_graph() for persisting
scene graphs to/from JSON files. Handles numpy arrays, bytes (hashes),
and optional fields.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np

from .objectnode import ObjectNode
from .roomnode import RoomNode
from .scenegraph import SimpleSceneGraph


def _ndarray_to_list(arr: np.ndarray) -> list[float]:
    return arr.tolist()


def _bytes_to_hex(b: bytes) -> str:
    return b.hex()


def _room_node_to_dict(node: RoomNode) -> dict[str, Any]:
    return {
        "room_id": node.room_id,
        "center": _ndarray_to_list(node.center),
        "bbox_extent": _ndarray_to_list(node.bbox_extent),
    }


def _object_node_to_dict(node: ObjectNode) -> dict[str, Any]:
    return {
        "obj_id": node.obj_id,
        "room_id": node.room_id,
        "position": _ndarray_to_list(node.position),
        "visual_hash": _bytes_to_hex(node.visual_hash),
        "semantic_hash": _bytes_to_hex(node.semantic_hash),
        "hit_count": node.hit_count,
        "last_seen_frame": node.last_seen_frame,
        "label": node.label,
        "confidence": node.confidence,
        "bbox_xyxy": list(node.bbox_xyxy) if node.bbox_xyxy is not None else None,
        "source_view_id": node.source_view_id,
        "position_confidence": node.position_confidence,
    }


def _room_node_from_dict(d: dict[str, Any]) -> RoomNode:
    return RoomNode(
        room_id=d["room_id"],
        center=np.asarray(d["center"], dtype=np.float32),
        bbox_extent=np.asarray(d["bbox_extent"], dtype=np.float32),
    )


def _object_node_from_dict(d: dict[str, Any]) -> ObjectNode:
    bbox = d.get("bbox_xyxy")
    return ObjectNode(
        obj_id=d["obj_id"],
        room_id=d["room_id"],
        position=np.asarray(d["position"], dtype=np.float32),
        visual_hash=bytes.fromhex(d["visual_hash"]),
        semantic_hash=bytes.fromhex(d["semantic_hash"]),
        hit_count=d["hit_count"],
        last_seen_frame=d["last_seen_frame"],
        label=d.get("label"),
        confidence=d.get("confidence"),
        bbox_xyxy=tuple(bbox) if bbox is not None else None,
        source_view_id=d.get("source_view_id"),
        position_confidence=d.get("position_confidence"),
    )


def save_scene_graph(path: str | Path, graph: SimpleSceneGraph) -> None:
    """Save a SimpleSceneGraph to a JSON file.

    Args:
        path: Output file path.
        graph: The scene graph to serialize.
    """
    data = {
        "version": 1,
        "rooms": {
            room_id: _room_node_to_dict(node)
            for room_id, node in graph.rooms.items()
        },
        "objects": {
            obj_id: _object_node_to_dict(node)
            for obj_id, node in graph.objects.items()
        },
    }
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)


def load_scene_graph(path: str | Path) -> SimpleSceneGraph:
    """Load a SimpleSceneGraph from a JSON file.

    Args:
        path: Input file path.

    Returns:
        Deserialized SimpleSceneGraph.
    """
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)

    rooms = {
        room_id: _room_node_from_dict(node_dict)
        for room_id, node_dict in data["rooms"].items()
    }
    objects = {
        obj_id: _object_node_from_dict(node_dict)
        for obj_id, node_dict in data["objects"].items()
    }
    return SimpleSceneGraph(rooms=rooms, objects=objects)