"""Lazy document/state loading for the desktop viewer."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

import numpy as np

from ..core import System

logger = logging.getLogger(__name__)


@dataclass(slots=True)
class StateEntry:
    """One lazily loadable state inside a viewer document."""

    label: str
    index: tuple[int, ...]
    load_state: Callable[[], np.ndarray]
    metadata: dict[str, Any]


@dataclass(slots=True)
class ViewerDocument:
    """One opened viewer source file."""

    path: Path
    label: str
    format_name: str
    system: System
    entries: list[StateEntry]
    metadata: dict[str, Any]


def _npz_state_shape(images: np.ndarray) -> tuple[tuple[int, ...], int]:
    if images.ndim == 5 and images.shape[-1] == 3:
        return tuple(int(v) for v in images.shape[:3]), 1
    if images.ndim == 6 and images.shape[-1] == 3:
        return tuple(int(v) for v in images.shape[1:4]), int(images.shape[0])
    raise ValueError(f"Shape {images.shape!r} of 'PATH' is not supported.")


def _load_npz_state(path: Path, array_name: str, index: int) -> np.ndarray:
    with np.load(path, allow_pickle=True) as datafile:
        images = np.asarray(datafile[array_name])
        if images.ndim == 5:
            if index != 0:
                raise IndexError(index)
            state = np.asarray(images)
        else:
            state = np.asarray(images[index])
    logger.debug(
        "Loaded NPZ state path=%s array=%s index=%s shape=%s dtype=%s",
        path,
        array_name,
        index,
        tuple(int(v) for v in state.shape),
        getattr(state, "dtype", None),
    )
    return state


def load_npz_document(path: str | Path, array_name: str = "PATH") -> ViewerDocument:
    """Open the current LCSIM-style NPZ format as a lazy viewer document."""
    resolved = Path(path).expanduser().resolve()
    with np.load(resolved, allow_pickle=True) as datafile:
        if array_name not in datafile:
            raise KeyError(f"NPZ document has no array {array_name!r}.")
        images = np.asarray(datafile[array_name])
        state_shape, nstates = _npz_state_shape(images)
        settings = (
            json.loads(datafile["settings"][()])
            if "settings" in datafile
            else None
        )

    system = System.cubic(size=state_shape)
    logger.info(
        "Opened NPZ document path=%s array=%s states=%s system_size=%s",
        resolved,
        array_name,
        nstates,
        tuple(int(v) for v in system.size),
    )
    logger.debug(
        "NPZ metadata path=%s array_shape=%s settings_present=%s",
        resolved,
        tuple(int(v) for v in images.shape),
        settings is not None,
    )
    entries = [
        StateEntry(
            label=f"{array_name}[{index}]",
            index=(int(index),),
            load_state=(lambda i=index: _load_npz_state(resolved, array_name, i)),
            metadata={"array_name": array_name, "index": int(index)},
        )
        for index in range(nstates)
    ]
    return ViewerDocument(
        path=resolved,
        label=resolved.name,
        format_name="NPZ",
        system=system,
        entries=entries,
        metadata={"array_name": array_name, "settings": settings},
    )


def _load_sad_module():
    try:
        import sadcompressor as sad
    except ModuleNotFoundError as exc:
        raise ImportError(
            "SAD support requires the optional `sadcompressor` dependency. "
            "Install TopoVec with `uv tool install 'topovec[view]'`, "
            "`uv sync --extra view`, or "
            "`uv sync --all-extras`."
        ) from exc
    return sad


def _load_sad_state(path: Path, index: int) -> np.ndarray:
    sad = _load_sad_module()
    from ..sad.io import load_lcsim_sad_single

    with sad.SADRandomReader(str(path)) as reader:
        reader.seek(index)
        director, _system, _settings = load_lcsim_sad_single(reader)
    state = np.asarray(director)
    logger.debug(
        "Loaded SAD state path=%s index=%s shape=%s dtype=%s",
        path,
        index,
        tuple(int(v) for v in state.shape),
        getattr(state, "dtype", None),
    )
    return state


def load_sad_document(path: str | Path) -> ViewerDocument:
    """Open a SAD file as a lazy viewer document."""
    sad = _load_sad_module()
    from ..sad.io import load_lcsim_sad_single

    resolved = Path(path).expanduser().resolve()
    with sad.SADRandomReader(str(resolved)) as reader:
        reader.seek(0)
        director, system, settings = load_lcsim_sad_single(reader)
        nstates = int(reader.nkeys)
    logger.info(
        "Opened SAD document path=%s states=%s system_size=%s",
        resolved,
        nstates,
        tuple(int(v) for v in system.size),
    )
    logger.debug(
        "SAD metadata path=%s preview_shape=%s settings_present=%s",
        resolved,
        tuple(int(v) for v in np.asarray(director).shape),
        settings is not None,
    )

    entries = [
        StateEntry(
            label=f"PATH[{index}]",
            index=(int(index),),
            load_state=(lambda i=index: _load_sad_state(resolved, i)),
            metadata={"index": int(index)},
        )
        for index in range(nstates)
    ]
    return ViewerDocument(
        path=resolved,
        label=resolved.name,
        format_name="SAD",
        system=system,
        entries=entries,
        metadata={"settings": settings, "preview_shape": tuple(director.shape)},
    )


_DOCUMENT_LOADERS: dict[str, Callable[[str | Path], ViewerDocument]] = {
    ".npz": load_npz_document,
    ".sad": load_sad_document,
}


def load_viewer_document(path: str | Path) -> ViewerDocument:
    """Open one viewer document based on filename suffix."""
    resolved = Path(path).expanduser().resolve()
    suffix = resolved.suffix.lower()
    if suffix not in _DOCUMENT_LOADERS:
        raise ValueError(
            f"Unsupported viewer document type {suffix!r}. Supported: "
            f"{', '.join(sorted(_DOCUMENT_LOADERS))}."
        )
    return _DOCUMENT_LOADERS[suffix](resolved)


__all__ = [
    "StateEntry",
    "ViewerDocument",
    "load_npz_document",
    "load_sad_document",
    "load_viewer_document",
]
