"""Registry of LoRA adapters for the Gradio demo.

Loads the pretrained base model once and wraps it with a `PeftModel` that
holds one or more named LoRA adapters. Each adapter also carries its own
classification head (saved separately as ``classifier.pt`` at training time),
so the demo can swap *both* the adapter and the head per request.

Inference modes:
    * Base (ImageNet-1k): ``model.disable_adapter()`` + original head.
    * Adapter N:          ``model.set_adapter(name)`` + the adapter's head.
"""

from __future__ import annotations

import copy
import json
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterable

import torch
from torch import nn
from transformers import AutoImageProcessor, AutoModelForImageClassification


BASE_MODEL_ID = "WinKawaks/vit-tiny-patch16-224"
BASE_TASK_NAME = "imagenet"
BASE_TASK_LABEL = "ImageNet-1k (base)"


@dataclass
class AdapterSpec:
    name: str
    repo_id: str
    display_name: str
    description: str = ""
    id2label: dict[int, str] = field(default_factory=dict)


ADAPTERS: list[AdapterSpec] = [
    AdapterSpec(
        name="food101",
        repo_id="turhancan97/vit-tiny-lora-food101",
        display_name="Food-101 (LoRA)",
        description="Fine-tuned LoRA adapter on the Food-101 dataset (101 classes).",
    ),
]


@dataclass
class TaskInfo:
    key: str
    display_name: str
    id2label: dict[int, str]
    is_base: bool


def _fetch_file(repo_id: str, filename: str) -> Path:
    local = Path(repo_id)
    if local.exists() and (local / filename).exists():
        return local / filename
    from huggingface_hub import hf_hub_download

    return Path(hf_hub_download(repo_id=repo_id, filename=filename))


class AdapterRegistry:
    """Loads the base model plus any registered adapters lazily and on demand."""

    def __init__(
        self,
        adapters: Iterable[AdapterSpec] | None = None,
        base_model_id: str = BASE_MODEL_ID,
        device: str | None = None,
    ) -> None:
        if adapters is None:
            adapters = ADAPTERS
        self.base_model_id = base_model_id
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        self.processor = AutoImageProcessor.from_pretrained(base_model_id, use_fast=True)
        self.base_model = AutoModelForImageClassification.from_pretrained(base_model_id)
        self.base_id2label: dict[int, str] = dict(self.base_model.config.id2label)

        self._original_classifier = copy.deepcopy(self.base_model.classifier)

        self._model = None
        self._loaded: dict[str, AdapterSpec] = {}
        self._classifiers: dict[str, nn.Module] = {}
        self._current_adapter: str | None = None

        for spec in adapters:
            self._try_load_adapter(spec)

        if self._model is None:
            self._model = self.base_model
        self._model.to(self.device).eval()
        self._original_classifier.to(self.device).eval()
        for head in self._classifiers.values():
            head.to(self.device).eval()

    @property
    def model(self):
        return self._model

    def available_tasks(self) -> list[TaskInfo]:
        tasks = [TaskInfo(
            key=BASE_TASK_NAME,
            display_name=BASE_TASK_LABEL,
            id2label=self.base_id2label,
            is_base=True,
        )]
        for spec in self._loaded.values():
            tasks.append(TaskInfo(
                key=spec.name,
                display_name=spec.display_name,
                id2label=spec.id2label,
                is_base=False,
            ))
        return tasks

    def adapter_task_keys(self) -> list[str]:
        return list(self._loaded.keys())

    def get_id2label(self, task_key: str) -> dict[int, str]:
        if task_key == BASE_TASK_NAME:
            return self.base_id2label
        return self._loaded[task_key].id2label

    def _classifier_host(self) -> nn.Module:
        """Return the module that actually *owns* the `classifier` submodule.

        `PeftModel` exposes `.classifier` via attribute proxying on top of
        `base_model.model.classifier`, but only the latter is what gets called
        during forward. Traverse until we find a module where `classifier` is
        a registered child (present in `_modules`).
        """
        visited: set[int] = set()
        m = self._model
        while id(m) not in visited:
            visited.add(id(m))
            if "classifier" in getattr(m, "_modules", {}):
                return m
            inner = getattr(m, "model", None) or getattr(m, "base_model", None)
            if inner is None or inner is m:
                break
            m = inner
        return self.base_model

    def _active_classifier(self) -> nn.Module:
        return self._classifier_host().classifier

    def _set_classifier(self, head: nn.Module) -> None:
        self._classifier_host().classifier = head

    @contextmanager
    def use_task(self, task_key: str):
        """Activate the right (adapter, head) pair for a single forward pass."""
        prev_head = self._active_classifier()
        if task_key == BASE_TASK_NAME:
            self._set_classifier(self._original_classifier)
            ctx = self._model.disable_adapter() if hasattr(self._model, "disable_adapter") else _null_ctx()
        else:
            spec = self._loaded[task_key]
            self._set_classifier(self._classifiers[spec.name])
            if hasattr(self._model, "set_adapter") and self._current_adapter != task_key:
                self._model.set_adapter(task_key)
                self._current_adapter = task_key
            ctx = _null_ctx()

        try:
            with ctx:
                yield
        finally:
            self._set_classifier(prev_head)

    def _try_load_adapter(self, spec: AdapterSpec) -> None:
        try:
            labels_path = _fetch_file(spec.repo_id, "labels.json")
            spec.id2label = {int(k): v for k, v in json.loads(labels_path.read_text()).items()}

            classifier_path = _fetch_file(spec.repo_id, "classifier.pt")
            head_state = torch.load(classifier_path, map_location="cpu", weights_only=True)
            hidden_size = self.base_model.config.hidden_size
            head = nn.Linear(hidden_size, len(spec.id2label))
            head.load_state_dict(head_state)

            from peft import PeftModel

            if self._model is None or not hasattr(self._model, "load_adapter"):
                self._model = PeftModel.from_pretrained(
                    self.base_model, spec.repo_id, adapter_name=spec.name,
                )
            else:
                self._model.load_adapter(spec.repo_id, adapter_name=spec.name)

            self._classifiers[spec.name] = head
            self._loaded[spec.name] = spec
            print(f"[adapters] loaded '{spec.name}' from {spec.repo_id} "
                  f"({len(spec.id2label)} classes)")
        except Exception as exc:
            print(f"[adapters] could not load '{spec.name}' from {spec.repo_id}: "
                  f"{type(exc).__name__}: {exc}")


@contextmanager
def _null_ctx():
    yield