"""`acaua.Model` — unified factory + abstract base for every vision adapter.

             from_pretrained dispatch
             ────────────────────────

 Model.from_pretrained(name, allow_non_apache=False, device='auto')
             │
             ▼
     ┌─────────────────────┐
     │ 1. license check    │  ──── _data/licenses.py (shipped in wheel)
     │    repo_id → license│       static manifest, pre-download
     │    LicenseError     │
     └──────────┬──────────┘
                │
                ▼
     ┌─────────────────────┐
     │ 2. task routing     │
     │  a. config.json     │ ──── hf_hub_download(revision=acaua_sha)
     │     acaua_task field│
     │  b. fallback dict   │ <─── _data/task_fallback.py
     │  TaskResolutionError│
     └──────────┬──────────┘
                │ "detection" | "segmentation" | ...
                ▼
     ┌─────────────────────┐
     │ 3. adapter dispatch │
     │  importlib lazy-load│ <─── _data/adapters.py (task → module.Class)
     └──────────┬──────────┘
                │
      ┌─────────┴─────────┐
      ▼                   ▼
RTDETRAdapter       Mask2FormerAdapter
(transformers)      (transformers)
      │                   │
      ▼                   ▼
.predict(img)       .predict(img)
      │                   │
      ▼                   ▼
DetectionResult     SegmentationResult
      │                   │
      └── .to_supervision() ──▶ supervision.Detections
"""

from __future__ import annotations

import importlib
import json
from abc import ABC, abstractmethod
from collections.abc import Sequence
from typing import TYPE_CHECKING

import torch

from acaua._data.adapters import ADAPTERS
from acaua._data.licenses import LICENSES, LicenseEntry
from acaua._data.task_fallback import TASK_FALLBACK
from acaua.device import resolve as resolve_device
from acaua.exceptions import LicenseError, TaskResolutionError, WeightDownloadError
from acaua.io import ImageInput

if TYPE_CHECKING:
    from acaua.results import Result


class Model(ABC):
    """Abstract base class for all acaua model adapters.

    Users call `Model.from_pretrained(name)` and get back a concrete subclass
    instance (RTDETRAdapter, Mask2FormerAdapter, ...). Subclassing this
    directly is possible for advanced users but not part of the public API
    in v0.1.
    """

    repo_id: str
    device: torch.device

    @classmethod
    def from_pretrained(
        cls,
        name: str,
        *,
        allow_non_apache: bool = False,
        device: str = "auto",
    ) -> Model:
        """Load a pretrained acaua model.

        Args:
            name: HF repo id (e.g. `"CondadosAI/rtdetr_r50vd"`). Must appear
                in `acaua._data.licenses.LICENSES`.
            allow_non_apache: bypass the `LicenseError` raise for non-Apache
                weights. Default `False` — the whole point of the library.
            device: target device. `"auto"` picks `cuda > mps > cpu`; explicit
                strings (`"cuda"`, `"mps"`, `"cpu"`, `"cuda:1"`, ...) are
                validated.

        Returns:
            A concrete `Model` subclass instance.

        Raises:
            LicenseError: repo is non-Apache and `allow_non_apache=False`.
            TaskResolutionError: repo not in the manifest, or its task cannot
                be resolved from `config.json` or the fallback dict.
            WeightDownloadError: `huggingface_hub` download failed.
        """
        entry = _lookup(name)
        _enforce_license(name, entry, allow_non_apache=allow_non_apache)
        task = _resolve_task(name, entry.acaua_sha)
        adapter_cls = _load_adapter(task)
        resolved_device = resolve_device(device)
        return adapter_cls._load(
            name=name,
            revision=entry.acaua_sha,
            device=resolved_device,
        )

    @classmethod
    @abstractmethod
    def _load(cls, *, name: str, revision: str, device: torch.device) -> Model:
        """Adapter-specific construction.

        Called by `Model.from_pretrained` after license + task routing have
        resolved. Concrete adapters fetch weights with the pinned `revision`,
        move to `device`, and return an instance.
        """
        raise NotImplementedError

    @abstractmethod
    def predict(
        self,
        inputs: ImageInput | Sequence[ImageInput],
        *,
        batch_size: int = 8,
        strict_rgb: bool = True,
    ) -> Result | list[Result]:
        """Run inference on one input or a sequence.

        All adapters must expose this exact signature — enforced by
        `tests/test_predict_parity.py`. Dropping a kwarg or adding a
        positional arg breaks the contract test.

        Args:
            inputs: single image or sequence of images. Accepted types:
                `str` (file path or http/https URL), `pathlib.Path`,
                `PIL.Image.Image`, `numpy.ndarray` (HWC uint8 RGB),
                `torch.Tensor` (CHW float32 in `[0, 1]`).
            batch_size: how many inputs to feed the backend at once. Default
                `8` — safe for ~8 GB VRAM. Tune per hardware.
            strict_rgb: if `True`, numpy inputs whose channel stats look
                BGR are rejected with a helpful error. Default `True`.

        Returns:
            Single input → single `Result` subclass instance.
            Sequence input → `list[Result]`.
        """
        raise NotImplementedError


# ---------------------------------------------------------------------------
# internal helpers (module-level, unit-testable without subclassing Model)
# ---------------------------------------------------------------------------


def _lookup(name: str) -> LicenseEntry:
    entry = LICENSES.get(name)
    if entry is None:
        raise TaskResolutionError(
            f"{name!r} is not in acaua's model manifest "
            "(src/acaua/_data/licenses.py). Only manifest-listed repos are "
            "supported. See https://github.com/CondadosAI/acaua for the "
            "current list."
        )
    return entry


def _enforce_license(name: str, entry: LicenseEntry, *, allow_non_apache: bool) -> None:
    if entry.license == "apache-2.0":
        return
    if allow_non_apache:
        return
    raise LicenseError(
        f"{name!r} is declared {entry.license!r} (upstream: "
        f"{entry.upstream_src!r}), not Apache-2.0. Pass "
        "allow_non_apache=True to load anyway."
    )


def _resolve_task(repo_id: str, revision: str) -> str:
    """Read `config.json` at the pinned revision; fall back to static dict.

    The `acaua_task` field wins when present. This matters for self-authored
    mirrors (v0.2+) and for the `CondadosAI/acaua-task-test` fixture; every
    v0.1 upstream-provenanced mirror hits the fallback because upstreams
    don't know about acaua.
    """
    from huggingface_hub import hf_hub_download
    from huggingface_hub.errors import HfHubHTTPError

    try:
        config_path = hf_hub_download(repo_id=repo_id, filename="config.json", revision=revision)
    except HfHubHTTPError as e:
        raise WeightDownloadError(
            f"failed to fetch config.json for {repo_id!r} at revision {revision!r}: {e}"
        ) from e
    except Exception as e:
        raise WeightDownloadError(
            f"unexpected error fetching config.json for {repo_id!r}: {e}"
        ) from e

    try:
        with open(config_path) as f:
            cfg = json.load(f)
    except (OSError, json.JSONDecodeError) as e:
        raise WeightDownloadError(f"failed to parse config.json for {repo_id!r}: {e}") from e

    task = cfg.get("acaua_task")
    if isinstance(task, str):
        return task
    if repo_id in TASK_FALLBACK:
        return TASK_FALLBACK[repo_id]
    raise TaskResolutionError(
        f"{repo_id!r} has no 'acaua_task' field in config.json and is not "
        "in the task fallback dict (src/acaua/_data/task_fallback.py). "
        "Cannot determine task."
    )


def _load_adapter(task: str) -> type[Model]:
    try:
        module_path, class_name = ADAPTERS[task]
    except KeyError as e:
        raise TaskResolutionError(
            f"task {task!r} has no registered adapter in "
            "src/acaua/_data/adapters.py. Supported tasks: "
            f"{sorted(ADAPTERS)}."
        ) from e
    module = importlib.import_module(module_path)
    adapter_cls: type = getattr(module, class_name)
    if not issubclass(adapter_cls, Model):
        raise TypeError(f"{module_path}.{class_name} is not a subclass of acaua.Model")
    return adapter_cls
