"""Shared live-scene residency management for desktop and notebook viewers."""

from __future__ import annotations

import logging
from collections import OrderedDict
from dataclasses import dataclass, field
from typing import Any, Callable

import numpy as np

from .session import (
    ViewerSession,
    ViewerSessionResidencySnapshot,
    adapt_viewer_host,
)

logger = logging.getLogger(__name__)


def _scene_supports_system(scene_cls: Any, system: Any) -> bool:
    supports = getattr(scene_cls, "supports_system", None)
    if not callable(supports):
        return True
    return bool(supports(system))


def _system_attribute_key(value: Any) -> Any:
    if value is None:
        return None
    array = np.asarray(value)
    if array.ndim == 0:
        return array.item()
    return tuple(int(v) for v in array.shape), tuple(array.reshape(-1).tolist())


def system_cache_key(system: Any) -> tuple[Any, ...] | None:
    """Compute the structural cache key shared by Qt and marimo viewers."""
    if system is None:
        return None
    return (
        type(system),
        *(
            _system_attribute_key(getattr(system, name, None))
            for name in ("size", "primitives", "representatives", "dim", "rank")
        ),
    )


@dataclass(slots=True)
class _SystemResidencyRecord:
    system: Any
    snapshot: ViewerSessionResidencySnapshot | None = None
    live_scenes: dict[str, Any] = field(default_factory=dict)


class SceneResidencyManager:
    """Manage per-system snapshots and a global LRU of live scenes."""

    def __init__(
        self,
        host: Any,
        *,
        scene_registry_loader: Callable[[], dict[str, Any]],
        session_factory: Callable[..., ViewerSession] = ViewerSession,
        session_kwargs: dict[str, Any] | None = None,
        max_live_scenes: int = 2,
    ):
        if int(max_live_scenes) < 1:
            raise ValueError("max_live_scenes must be at least 1.")
        self._host = adapt_viewer_host(host)
        self._scene_registry_loader = scene_registry_loader
        self._session_factory = session_factory
        self._session_kwargs = {} if session_kwargs is None else dict(session_kwargs)
        self._max_live_scenes = int(max_live_scenes)
        self._records: dict[tuple[Any, ...], _SystemResidencyRecord] = {}
        self._active_system_key: tuple[Any, ...] | None = None
        self._active_session: ViewerSession | None = None
        self._live_scene_lru: OrderedDict[tuple[tuple[Any, ...], str], None] = (
            OrderedDict()
        )
        self._last_activation_mode = "new_scene"

    @property
    def max_live_scenes(self) -> int:
        return self._max_live_scenes

    @property
    def session(self) -> ViewerSession | None:
        return self._active_session

    @property
    def last_activation_mode(self) -> str:
        return self._last_activation_mode

    def _scene_registry(self) -> dict[str, Any]:
        return dict(self._scene_registry_loader())

    def _resolve_scene_name(
        self,
        *,
        system: Any,
        record: _SystemResidencyRecord,
        preferred_scene_name: str | None,
    ) -> str:
        registry = self._scene_registry()
        supported = [
            str(name)
            for name, scene_cls in registry.items()
            if _scene_supports_system(scene_cls, system)
        ]
        if not supported:
            raise RuntimeError("No built-in scenes are compatible with this system.")

        snapshot_scene = (
            None if record.snapshot is None else record.snapshot.selected_scene
        )
        for candidate in (snapshot_scene, preferred_scene_name):
            if candidate and candidate in supported:
                return str(candidate)
        return supported[0]

    def _actual_live_scene_keys(self) -> set[tuple[tuple[Any, ...], str]]:
        keys: set[tuple[tuple[Any, ...], str]] = set()
        for system_key, record in self._records.items():
            for scene_name in record.live_scenes:
                keys.add((system_key, str(scene_name)))
        if self._active_session is not None and self._active_system_key is not None:
            for scene_name in self._active_session.live_scene_names():
                keys.add((self._active_system_key, str(scene_name)))
        return keys

    def _active_scene_key(self) -> tuple[tuple[Any, ...], str] | None:
        if self._active_session is None or self._active_system_key is None:
            return None
        scene_name = self._active_session.selected_scene
        if scene_name is None:
            return None
        return self._active_system_key, str(scene_name)

    def _reconcile_lru(self) -> None:
        actual = self._actual_live_scene_keys()
        for key in list(self._live_scene_lru.keys()):
            if key not in actual:
                self._live_scene_lru.pop(key, None)
        for key in actual:
            self._live_scene_lru.setdefault(key, None)

    def _mark_scene_used(self, system_key: tuple[Any, ...], scene_name: str) -> None:
        scene_key = (system_key, str(scene_name))
        self._reconcile_lru()
        self._live_scene_lru.setdefault(scene_key, None)
        self._live_scene_lru.move_to_end(scene_key)

    def _release_scene(self, scene: Any) -> None:
        release = getattr(scene, "release", None)
        if callable(release):
            release()

    def _enforce_live_scene_limit(self) -> None:
        self._reconcile_lru()
        pinned_key = self._active_scene_key()
        while len(self._live_scene_lru) > self._max_live_scenes:
            evicted = False
            for scene_key in list(self._live_scene_lru.keys()):
                if scene_key == pinned_key:
                    continue
                system_key, scene_name = scene_key
                if system_key == self._active_system_key and self._active_session is not None:
                    if self._active_session.evict_live_scene(scene_name):
                        logger.debug(
                            "Evicted live scene from active session system_key=%s scene=%s",
                            system_key,
                            scene_name,
                        )
                        self._live_scene_lru.pop(scene_key, None)
                        evicted = True
                        break
                    self._live_scene_lru.pop(scene_key, None)
                    continue
                record = self._records.get(system_key)
                if record is None:
                    self._live_scene_lru.pop(scene_key, None)
                    continue
                scene = record.live_scenes.pop(scene_name, None)
                if scene is None:
                    self._live_scene_lru.pop(scene_key, None)
                    continue
                logger.debug(
                    "Evicted live scene from inactive record system_key=%s scene=%s",
                    system_key,
                    scene_name,
                )
                self._release_scene(scene)
                self._live_scene_lru.pop(scene_key, None)
                evicted = True
                break
            if not evicted:
                break
        self._reconcile_lru()

    def _park_session(
        self,
        system_key: tuple[Any, ...],
        session: ViewerSession,
    ) -> None:
        record = self._records.setdefault(
            system_key,
            _SystemResidencyRecord(system=getattr(session.scene, "system", None)),
        )
        record.snapshot = session.export_residency_snapshot()
        record.live_scenes = session.export_live_scenes()
        logger.debug(
            "Parked viewer session system_key=%s selected_scene=%s live_scenes=%s",
            system_key,
            record.snapshot.selected_scene if record.snapshot is not None else None,
            sorted(record.live_scenes),
        )
        self._reconcile_lru()

    def activate(
        self,
        *,
        system: Any,
        state: np.ndarray | None,
        preferred_scene_name: str | None = None,
    ) -> ViewerSession:
        target_key = system_cache_key(system)
        if target_key is None:
            raise RuntimeError("Scene residency requires a non-null system.")

        if self._active_session is not None and self._active_system_key == target_key:
            self._active_session.set_render_state(state)
            active_scene = self._active_session.selected_scene
            if active_scene is not None:
                self._mark_scene_used(target_key, active_scene)
            self._last_activation_mode = "active_session"
            self._enforce_live_scene_limit()
            return self._active_session

        record = self._records.setdefault(
            target_key,
            _SystemResidencyRecord(system=system),
        )
        resolved_system = record.system
        scene_name = self._resolve_scene_name(
            system=resolved_system,
            record=record,
            preferred_scene_name=preferred_scene_name,
        )
        registry = self._scene_registry()
        scene = record.live_scenes.get(scene_name)
        activation_mode = "live_scene" if scene is not None else "new_scene"
        if scene is None:
            scene_cls = registry[scene_name]
            scene = self._host.create_scene(scene_cls, resolved_system)

        previous_key = self._active_system_key
        previous_session = self._active_session

        try:
            self._host.attach_scene(scene)
            session = self._session_factory(
                self._host,
                state=state,
                **self._session_kwargs,
            )
            session.seed_live_scenes(record.live_scenes)
            if record.snapshot is not None:
                session.restore_residency_snapshot(record.snapshot)
        except Exception:
            if previous_session is not None:
                self._host.attach_scene(previous_session.scene)
            raise

        if previous_key is not None and previous_session is not None:
            self._park_session(previous_key, previous_session)

        record.live_scenes = {}
        self._active_system_key = target_key
        self._active_session = session
        setattr(session, "_scene_residency_manager", self)
        self._last_activation_mode = activation_mode
        self._mark_scene_used(target_key, session.selected_scene or scene_name)
        self._enforce_live_scene_limit()
        logger.debug(
            "Activated residency-managed session mode=%s system_key=%s selected_scene=%s",
            activation_mode,
            target_key,
            session.selected_scene,
        )
        return session

    def switch_scene(self, scene_name: str) -> bool:
        if self._active_session is None or self._active_system_key is None:
            raise RuntimeError("No active viewer session to switch.")
        changed = self._active_session.switch_scene(scene_name)
        if not changed:
            return False
        self._mark_scene_used(self._active_system_key, scene_name)
        self._enforce_live_scene_limit()
        return True

    def deactivate(self) -> None:
        if self._active_session is None or self._active_system_key is None:
            return
        self._park_session(self._active_system_key, self._active_session)
        self._active_session = None
        self._active_system_key = None
        self._reconcile_lru()

    def release(self) -> None:
        if self._active_session is not None:
            self._active_session.release_live_scenes()
            self._active_session = None
            self._active_system_key = None
        for record in self._records.values():
            for scene in record.live_scenes.values():
                self._release_scene(scene)
            record.live_scenes = {}
            record.snapshot = None
        self._records.clear()
        self._live_scene_lru.clear()


__all__ = ["SceneResidencyManager", "system_cache_key"]
