"""Scene state and handlers for object management."""

from __future__ import annotations

import json
import os
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, TYPE_CHECKING

try:
    import tomllib
except ModuleNotFoundError:
    import tomli as tomllib  # type: ignore

if TYPE_CHECKING:
    from .node import OutputPort, VoiceControlLogger


@dataclass
class SceneObject:
    """An object in the scene with spatial properties."""

    id: str  # Unique identifier
    object_type: str  # cube, cylinder, box, etc.
    color: str  # red, blue, etc.
    size: str = "normal"  # big, small, normal

    # Spatial properties
    position_mm: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
    bbox_mm: Optional[List[float]] = None  # [x1, y1, x2, y2] bounding box
    height_mm: float = 30.0  # Object height
    area_mm2: Optional[float] = None  # Footprint area

    # Relationships
    on_top_of: Optional[str] = None  # ID of object this is stacked on

    # Metadata
    source: str = "detected"  # "detected", "manual", or "config"
    confidence: float = 1.0
    last_seen_ns: int = 0

    @property
    def center(self) -> List[float]:
        """Get center point (alias for position_mm)."""
        return self.position_mm

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "id": self.id,
            "object_type": self.object_type,
            "color": self.color,
            "size": self.size,
            "position_mm": self.position_mm,
            "bbox_mm": self.bbox_mm,
            "height_mm": self.height_mm,
            "area_mm2": self.area_mm2,
            "on_top_of": self.on_top_of,
            "source": self.source,
            "confidence": self.confidence,
            "last_seen_ns": self.last_seen_ns,
        }


class SceneState:
    """Thread-safe scene state with spatial relationships."""

    # Tolerance for stacking detection (mm)
    Z_TOLERANCE_MM: float = 15.0

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._objects: Dict[str, SceneObject] = {}
        # Scene update mode: "static" (default) or "dynamic"
        self._update_mode: str = os.getenv("SCENE_UPDATE_MODE", "static").lower()
        # Whether scene has been captured (STATIC mode only)
        self._scene_captured: bool = False

    # === Core Operations ===

    def add(self, obj: SceneObject) -> bool:
        """Add object to scene.

        Returns True if added, False if ID already exists.
        """
        with self._lock:
            if obj.id in self._objects:
                return False
            self._objects[obj.id] = obj
            return True

    def delete(self, object_id: str) -> bool:
        """Remove object and clear relationships pointing to it.

        Returns True if deleted, False if not found.
        """
        with self._lock:
            if object_id not in self._objects:
                return False
            del self._objects[object_id]
            # Clear relationships pointing to deleted object
            for obj in self._objects.values():
                if obj.on_top_of == object_id:
                    obj.on_top_of = None
            return True

    def update(self, object_id: str, **properties: Any) -> bool:
        """Update object properties.

        Returns True if updated, False if not found.
        """
        with self._lock:
            if object_id not in self._objects:
                return False
            obj = self._objects[object_id]
            for key, value in properties.items():
                if hasattr(obj, key):
                    setattr(obj, key, value)
            return True

    def get(self, object_id: str) -> Optional[SceneObject]:
        """Get object by ID."""
        with self._lock:
            obj = self._objects.get(object_id)
            return obj  # Return reference (caller should not mutate)

    def get_all(self) -> List[SceneObject]:
        """Get all objects."""
        with self._lock:
            return list(self._objects.values())

    def clear(self) -> None:
        """Remove all objects."""
        with self._lock:
            self._objects.clear()

    def clear_detected(self) -> None:
        """Remove only detected objects (keep manual and config)."""
        with self._lock:
            to_delete = [
                obj_id
                for obj_id, obj in self._objects.items()
                if obj.source == "detected"
            ]
            for obj_id in to_delete:
                del self._objects[obj_id]
            # Clear broken relationships
            for obj in self._objects.values():
                if obj.on_top_of and obj.on_top_of not in self._objects:
                    obj.on_top_of = None
            # Reset capture flag to allow next detection
            self._scene_captured = False

    # === Query ===

    def query(
        self,
        object_type: Optional[str] = None,
        color: Optional[str] = None,
        size: Optional[str] = None,
        source: Optional[str] = None,
    ) -> List[SceneObject]:
        """Find objects matching criteria."""
        with self._lock:
            results: List[SceneObject] = []
            for obj in self._objects.values():
                if object_type is not None and obj.object_type != object_type:
                    continue
                if color is not None and obj.color != color:
                    continue
                if size is not None and obj.size != size:
                    continue
                if source is not None and obj.source != source:
                    continue
                results.append(obj)
            return results

    def count(self) -> int:
        """Get total object count."""
        with self._lock:
            return len(self._objects)

    # === Scene Update Mode ===

    def is_static_mode(self) -> bool:
        """Check if in STATIC update mode."""
        return self._update_mode == "static"

    def is_captured(self) -> bool:
        """Check if scene has been captured (STATIC mode)."""
        with self._lock:
            return self._scene_captured

    def reset_capture(self) -> None:
        """Reset capture flag to allow next detection to update scene."""
        with self._lock:
            self._scene_captured = False

    # === Spatial Relationships ===

    def set_on_top_of(self, object_id: str, below_id: Optional[str]) -> bool:
        """Set stacking relationship.

        Returns True if set, False if object not found.
        """
        with self._lock:
            if object_id not in self._objects:
                return False
            if below_id is not None and below_id not in self._objects:
                return False
            self._objects[object_id].on_top_of = below_id
            return True

    def infer_relationships(self) -> None:
        """Infer stacking from positions.

        Call after replace_detected() to automatically detect stacking.
        A is on B if: A.z ≈ B.z + B.height AND A overlaps B horizontally.
        """
        with self._lock:
            self._infer_relationships_unlocked()

    def _infer_relationships_unlocked(self) -> None:
        """Infer relationships without holding lock (internal use)."""
        # Reset all relationships first
        for obj in self._objects.values():
            obj.on_top_of = None

        for obj in self._objects.values():
            best_match: Optional[SceneObject] = None
            best_distance = float("inf")

            for candidate in self._objects.values():
                if candidate.id == obj.id:
                    continue

                # Check if obj is on top of candidate
                expected_z = candidate.position_mm[2] + candidate.height_mm
                actual_z = obj.position_mm[2]
                z_diff = abs(actual_z - expected_z)

                if z_diff < self.Z_TOLERANCE_MM:
                    if self._overlaps_horizontally_unlocked(obj, candidate):
                        # Use closest match in z
                        if z_diff < best_distance:
                            best_match = candidate
                            best_distance = z_diff

            if best_match is not None:
                obj.on_top_of = best_match.id

    def get_object_below(self, object_id: str) -> Optional[SceneObject]:
        """Get object this one is stacked on."""
        with self._lock:
            obj = self._objects.get(object_id)
            if obj is None or obj.on_top_of is None:
                return None
            return self._objects.get(obj.on_top_of)

    def get_objects_above(self, object_id: str) -> List[SceneObject]:
        """Get objects stacked directly on this one."""
        with self._lock:
            results: List[SceneObject] = []
            for obj in self._objects.values():
                if obj.on_top_of == object_id:
                    results.append(obj)
            return results

    def get_stack(self, object_id: str) -> List[SceneObject]:
        """Get all objects in stack (bottom to top).

        Returns empty list if object not found.
        """
        with self._lock:
            obj = self._objects.get(object_id)
            if obj is None:
                return []

            # Find bottom of stack
            bottom = obj
            while bottom.on_top_of is not None:
                below = self._objects.get(bottom.on_top_of)
                if below is None:
                    break
                bottom = below

            # Build stack from bottom up
            stack = [bottom]
            visited = {bottom.id}

            while True:
                # Find object on top of current top
                current_top = stack[-1]
                found_next = False
                for candidate in self._objects.values():
                    if candidate.on_top_of == current_top.id and candidate.id not in visited:
                        stack.append(candidate)
                        visited.add(candidate.id)
                        found_next = True
                        break
                if not found_next:
                    break

            return stack

    def effective_height(self, object_id: str) -> float:
        """Get height including objects below (if stacked).

        Returns 0 if object not found.
        """
        with self._lock:
            obj = self._objects.get(object_id)
            if obj is None:
                return 0.0
            return self._effective_height_unlocked(obj)

    def _effective_height_unlocked(self, obj: SceneObject) -> float:
        """Calculate effective height without holding lock."""
        if obj.on_top_of is None:
            return obj.height_mm
        below = self._objects.get(obj.on_top_of)
        if below is None:
            return obj.height_mm
        return obj.height_mm + self._effective_height_unlocked(below)

    def _overlaps_horizontally(self, a: SceneObject, b: SceneObject) -> bool:
        """Check if two objects overlap in x,y plane (thread-safe)."""
        with self._lock:
            return self._overlaps_horizontally_unlocked(a, b)

    def _overlaps_horizontally_unlocked(self, a: SceneObject, b: SceneObject) -> bool:
        """Check horizontal overlap without lock.

        If bounding boxes are available, use them.
        Otherwise, check if centers are within a distance threshold.
        """
        # If both have bounding boxes, use proper overlap check
        if a.bbox_mm is not None and b.bbox_mm is not None:
            # bbox format: [x1, y1, x2, y2]
            ax1, ay1, ax2, ay2 = a.bbox_mm
            bx1, by1, bx2, by2 = b.bbox_mm
            # Check for no overlap (if any of these is true, no overlap)
            if ax2 < bx1 or bx2 < ax1:
                return False
            if ay2 < by1 or by2 < ay1:
                return False
            return True

        # Fall back to center distance check
        # Default threshold: ~30mm (typical small object radius)
        threshold = 40.0
        dx = abs(a.position_mm[0] - b.position_mm[0])
        dy = abs(a.position_mm[1] - b.position_mm[1])
        return dx < threshold and dy < threshold

    # === Detector Integration ===

    def replace_detected(self, objects: List[Dict[str, Any]]) -> None:
        """Replace detected objects (keeps manual and config objects).

        Automatically generates IDs and infers stacking relationships.
        """
        with self._lock:
            # Remove existing detected objects
            to_delete = [
                obj_id
                for obj_id, obj in self._objects.items()
                if obj.source == "detected"
            ]
            for obj_id in to_delete:
                del self._objects[obj_id]

            # Add new detected objects
            now_ns = time.time_ns()
            counters: Dict[str, int] = {}

            for obj_data in objects:
                obj_type = obj_data.get("object_type", "unknown")
                color = obj_data.get("color", "unknown")

                # Generate unique ID
                key = f"{obj_type}_{color}"
                counters[key] = counters.get(key, 0) + 1
                obj_id = f"{key}_{counters[key]}"

                position = obj_data.get("position_mm", [0.0, 0.0, 0.0])
                if len(position) < 3:
                    position = position + [0.0] * (3 - len(position))

                scene_obj = SceneObject(
                    id=obj_id,
                    object_type=obj_type,
                    color=color,
                    size=obj_data.get("size", "normal"),
                    position_mm=position,
                    bbox_mm=obj_data.get("bbox_mm"),
                    height_mm=obj_data.get("height_mm", 30.0),
                    area_mm2=obj_data.get("area_mm2"),
                    source="detected",
                    confidence=obj_data.get("confidence", 1.0),
                    last_seen_ns=now_ns,
                )
                self._objects[obj_id] = scene_obj

            # Clear broken relationships
            for obj in self._objects.values():
                if obj.on_top_of and obj.on_top_of not in self._objects:
                    obj.on_top_of = None

            # Infer stacking relationships
            self._infer_relationships_unlocked()

    # === Serialization ===

    def to_list(self) -> List[Dict[str, Any]]:
        """Convert all objects to list of dictionaries."""
        with self._lock:
            return [obj.to_dict() for obj in self._objects.values()]

    def get_objects_as_dicts(self) -> List[Dict[str, Any]]:
        """Get all objects as dictionaries (for compatibility)."""
        return self.to_list()

    # === Object Matching (voice criteria) ===

    def find_by_criteria(
        self,
        obj_type: str,
        color: str,
        size: str,
        class_map: Optional[Dict[str, str]] = None,
    ) -> Optional[SceneObject]:
        """Find object matching voice command criteria.

        Args:
            obj_type: Object type from voice command (e.g., "cubo", "cilindro")
            color: Color from voice command (e.g., "rojo", "azul")
            size: Size from voice command (e.g., "grande", "pequeno")
            class_map: Optional mapping to translate voice terms to detector terms

        Returns:
            First matching SceneObject, or None if not found.
        """
        if obj_type == "no especificado":
            return None

        class_map = class_map or {}

        def translate(token: str) -> str:
            return class_map.get(token, token)

        target_type = translate(obj_type)
        target_color = translate(color) if color != "no especificado" else None
        target_size = translate(size) if size != "no especificado" else None

        results = self.query(
            object_type=target_type,
            color=target_color,
            size=target_size,
        )

        return results[0] if results else None

    # === Config Loading ===

    def load_from_config(self, config_path: Optional[str] = None) -> int:
        """Load static objects from TOML config file.

        Config format:
            [object_parameters]
            normal_height = 220.0

            [[static_objects]]
            type = "box"
            color = "blue"
            position = [100.0, -200.0]
            size = "big"
            height = 50.0

        Args:
            config_path: Path to config file. Defaults to CONFIG_FILE env var or "config.toml".

        Returns:
            Number of objects loaded.
        """
        if config_path is None:
            config_path = os.getenv("CONFIG_FILE", "config.toml")

        cfg = self._load_toml(config_path)
        if not cfg:
            return 0

        obj_cfg = cfg.get("object_parameters", {})
        base_z = float(obj_cfg.get("normal_height", 220.0))

        count = 0
        counters: Dict[str, int] = {}

        for obj in cfg.get("static_objects", []):
            obj_type = obj.get("type", "box")
            color = obj.get("color", "unknown")
            pos = obj.get("position", [0, 0])

            # Generate unique ID
            key = f"{obj_type}_{color}"
            counters[key] = counters.get(key, 0) + 1
            obj_id = f"config_{key}_{counters[key]}"

            scene_obj = SceneObject(
                id=obj_id,
                object_type=obj_type,
                color=color,
                size=obj.get("size", "big"),
                position_mm=[float(pos[0]), float(pos[1]), base_z],
                height_mm=obj.get("height", 30.0),
                source="config",
            )
            if self.add(scene_obj):
                count += 1

        return count

    @staticmethod
    def _load_toml(path: str) -> Dict[str, Any]:
        """Load TOML file, returning empty dict on error."""
        if not path or not os.path.exists(path):
            return {}
        try:
            with open(path, "rb") as f:
                return tomllib.load(f)
        except Exception:
            return {}


class SceneNotifier:
    """Emits scene updates to clients."""

    def __init__(
        self,
        scene: SceneState,
        port: "OutputPort",
        logger: "VoiceControlLogger",
    ) -> None:
        self._scene = scene
        self._port = port
        self._logger = logger

    def send_scene_update(self) -> None:
        """Send current scene state to clients."""
        self._port.send_json({"objects": self._scene.to_list()})

    def send_scene_reset(self) -> None:
        """Send scene reset notification to clients."""
        self._scene.clear_detected()
        self._port.send_json({"objects": self._scene.to_list(), "reset": True})
        self._logger.log("Sent scene reset notification")

    def send_initial_scene(self) -> None:
        """Send initial scene state on startup."""
        self._port.send_json({"objects": self._scene.to_list(), "reset": True})
        self._logger.log("Sent initial scene reset notification")


class ObjectsHandler:
    """Handles objects detection events from the Dora node."""

    def __init__(
        self,
        scene: SceneState,
        notifier: SceneNotifier,
        logger: "VoiceControlLogger",
    ) -> None:
        self._scene = scene
        self._notifier = notifier
        self._logger = logger

    def handle(self, event: Dict[str, Any]) -> None:
        """Handle an objects detection event."""
        raw = event["value"][0].as_py() if len(event["value"]) else ""
        if not raw:
            return

        # In STATIC mode, ignore updates after scene is captured
        if self._scene.is_static_mode() and self._scene.is_captured():
            return

        try:
            payload = json.loads(raw)
            objects = payload.get("objects", [])
        except Exception as e:
            self._logger.warn(f"failed to parse objects payload: {e}")
            objects = []

        self._scene.replace_detected(objects)

        # Mark scene as captured after first successful update (STATIC mode)
        if self._scene.is_static_mode():
            with self._scene._lock:
                self._scene._scene_captured = True
            self._logger.log("Scene captured (STATIC mode)")

        # Emit scene update
        self._notifier.send_scene_update()
