"""
Model registry system inspired by torchvision and timm.

This module provides a centralized registry for model constructors,
enabling dynamic model creation and discovery.
"""

from __future__ import annotations

from dataclasses import dataclass
from importlib import import_module
from typing import Any, Callable, Dict, Iterable, List, Optional


@dataclass
class ModelEntry:
    name: str
    constructor: Callable[..., Any]
    tags: tuple[str, ...]
    metadata: Dict[str, Any]


class ModelRegistry:
    """Registry for storing model constructors with metadata."""

    def __init__(self) -> None:
        self._registry: Dict[str, ModelEntry] = {}

    # ------------------------------------------------------------------
    def register(
        self,
        name: str,
        constructor: Callable[..., Any],
        *,
        tags: Optional[Iterable[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        overwrite: bool = False,
    ) -> None:
        if not overwrite and name in self._registry:
            existing = self._registry[name]
            # Allow real constructors to replace lazy placeholders without
            # requiring every `@register_model(..., overwrite=True)` callsite.
            if not bool(existing.metadata.get("_lazy_placeholder", False)):
                raise KeyError(f"Model {name!r} already exists. Set overwrite=True to replace it.")
        entry = ModelEntry(
            name=name,
            constructor=constructor,
            tags=tuple(tags or ()),
            metadata=_apply_default_metadata_fields(tags=tuple(tags or ()), metadata=metadata),
        )
        self._registry[name] = entry

    def get(self, name: str) -> Callable[..., Any]:
        try:
            return self._registry[name].constructor
        except KeyError as exc:
            available = ", ".join(sorted(self._registry)) or "<empty>"
            raise KeyError(f"Model {name!r} not found. Available models: {available}") from exc

    def available(self, *, tags: Optional[Iterable[str]] = None) -> List[str]:
        if tags is None:
            return sorted(self._registry)
        tag_set = set(tags)
        return sorted(
            entry.name for entry in self._registry.values() if tag_set.issubset(entry.tags)
        )

    def info(self, name: str) -> ModelEntry:
        if name not in self._registry:
            raise KeyError(f"Model {name!r} not found in registry")
        return self._registry[name]

    def model_info(self, name: str) -> Dict[str, Any]:
        from pyimgano.models.introspection import model_entry_info

        entry = self.info(name)
        return model_entry_info(entry)


MODEL_REGISTRY = ModelRegistry()


def _apply_default_metadata_fields(
    *,
    tags: Iterable[str],
    metadata: Optional[Dict[str, Any]],
) -> Dict[str, Any]:
    meta = dict(metadata or {})
    tag_set = {str(tag).strip().lower() for tag in tags}
    requires_checkpoint = bool(meta.get("requires_checkpoint", False))

    if requires_checkpoint and not str(meta.get("weights_source", "")).strip():
        backend = str(meta.get("backend", "")).strip().lower()
        if backend == "anomalib" or "anomalib" in tag_set:
            meta["weights_source"] = "upstream-anomalib-checkpoint"
        elif backend == "patchcore_inspection" or "patchcore_inspection" in tag_set:
            meta["weights_source"] = "upstream-patchcore-inspection-checkpoint"
        elif "onnx" in tag_set:
            meta["weights_source"] = "local-exported-onnx"
        elif "torchscript" in tag_set:
            meta["weights_source"] = "local-exported-torchscript"

    return meta


def materialize_model_constructor(name: str) -> Callable[..., Any]:
    """Return the real model constructor, importing its module if needed.

    Notes
    -----
    Under the v7 lazy registry, `pyimgano.models` registers lightweight
    placeholder constructors and records the owning module name in
    `metadata._lazy_module`. This helper imports that module to replace the
    placeholder with the real constructor, without instantiating the detector.
    """

    entry = MODEL_REGISTRY.info(str(name))
    if bool(entry.metadata.get("_lazy_placeholder", False)):
        module_name = str(entry.metadata.get("_lazy_module", "")).strip()
        if module_name:
            try:
                import_module(f"pyimgano.models.{module_name}")
            except ModuleNotFoundError as exc:
                from pyimgano.utils.extras import extra_for_root_module

                missing = getattr(exc, "name", None)
                root = str(missing).split(".", 1)[0].strip() if missing else ""
                if extra_for_root_module(root) is None:
                    raise
                return entry.constructor
            entry = MODEL_REGISTRY.info(str(name))
    return entry.constructor


def register_model(
    name: str,
    *,
    tags: Optional[Iterable[str]] = None,
    metadata: Optional[Dict[str, Any]] = None,
    overwrite: bool = False,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """
    Decorator for automatically registering models at import time.

    Parameters
    ----------
    name : str
        Unique name for the model
    tags : Iterable[str], optional
        Tags for categorizing the model (e.g., 'vision', 'classical', 'deep')
    metadata : Dict[str, Any], optional
        Additional metadata (e.g., description, paper, year)
    overwrite : bool, default=False
        Whether to overwrite existing registration

    Returns
    -------
    Callable
        Decorator function

    Examples
    --------
    >>> @register_model("my_model", tags=["vision", "ml"])
    ... class MyModel:
    ...     pass
    """

    def decorator(constructor: Callable[..., Any]) -> Callable[..., Any]:
        tag_list = list(tags or ())

        # Best-effort industrial default: if a model exposes pixel-map methods,
        # tag it as `pixel_map` for discovery (CLI, workbench) even if the author
        # forgot to add the tag explicitly.
        has_map_method = bool(
            hasattr(constructor, "predict_anomaly_map") or hasattr(constructor, "get_anomaly_map")
        )
        if has_map_method and "pixel_map" not in tag_list:
            tag_list.append("pixel_map")

        metadata_payload = _apply_default_metadata_fields(tags=tag_list, metadata=metadata)

        MODEL_REGISTRY.register(
            name,
            constructor,
            tags=tag_list,
            metadata=metadata_payload,
            overwrite=overwrite,
        )
        return constructor

    return decorator


def create_model(name: str, *args, **kwargs):
    """
    Create a model instance from its registered name.

    Parameters
    ----------
    name : str
        Registered model name
    *args
        Positional arguments to pass to the model constructor
    **kwargs
        Keyword arguments to pass to the model constructor

    Returns
    -------
    model
        Instantiated model object

    Examples
    --------
    >>> model = create_model("vision_ecod", contamination=0.1)
    """
    constructor = MODEL_REGISTRY.get(name)
    return constructor(*args, **kwargs)


def list_models(*, tags: Optional[Iterable[str]] = None) -> List[str]:
    """
    List available model names, optionally filtered by tags.

    Parameters
    ----------
    tags : Iterable[str], optional
        Filter models by tags (only models with all specified tags are returned)

    Returns
    -------
    List[str]
        Sorted list of model names

    Examples
    --------
    >>> all_models = list_models()
    >>> classical_models = list_models(tags=["classical"])
    >>> vision_ml_models = list_models(tags=["vision", "ml"])
    """
    return MODEL_REGISTRY.available(tags=tags)


def model_info(name: str) -> Dict[str, Any]:
    """Return a stable, JSON-friendly model info payload.

    Notes
    -----
    `ModelRegistry.model_info()` is intentionally a minimal, generic
    introspection payload. This module-level helper enriches the payload with
    computed capabilities for convenience in CLI/workbench discovery.
    """

    from pyimgano.models.capabilities import (
        compute_model_capabilities,
        compute_model_deployment_profile,
    )
    from pyimgano.models.introspection import model_entry_info

    materialize_model_constructor(name)
    entry = MODEL_REGISTRY.info(name)
    payload = model_entry_info(entry)

    caps = compute_model_capabilities(entry)
    caps_payload = {
        "input_modes": list(caps.input_modes),
        "supports_pixel_map": bool(caps.supports_pixel_map),
        "supports_checkpoint": bool(caps.supports_checkpoint),
        "requires_checkpoint": bool(caps.requires_checkpoint),
        "supports_save_load": bool(caps.supports_save_load),
        "supports_confidence": bool(caps.supports_confidence),
    }
    payload["capabilities"] = caps_payload
    payload["deployment_profile"] = compute_model_deployment_profile(entry)
    # Convenience top-level aliases (stable for CLI output).
    payload["input_modes"] = list(caps.input_modes)
    payload["supports_pixel_map"] = bool(caps.supports_pixel_map)
    payload["supports_checkpoint"] = bool(caps.supports_checkpoint)
    payload["requires_checkpoint"] = bool(caps.requires_checkpoint)
    payload["supports_save_load"] = bool(caps.supports_save_load)
    payload["supports_confidence"] = bool(caps.supports_confidence)
    return payload


def model_metadata_contract() -> list[dict[str, Any]]:
    """Return the structured metadata contract for registry models."""

    from pyimgano.models.metadata_contract import metadata_contract_fields

    return metadata_contract_fields()


def audit_model_metadata(
    *,
    names: Optional[Iterable[str]] = None,
    limit: int | None = None,
) -> dict[str, Any]:
    """Audit registry entries against the model metadata contract."""

    from pyimgano.models.metadata_contract import audit_metadata_contract

    selected = None if names is None else list(names)
    return audit_metadata_contract(MODEL_REGISTRY, names=selected, limit=limit)
