"""Standalone ``trust_remote_code`` wrapper for the RateBV continuous scorer.

This module is shipped verbatim into an exported HuggingFace repo as ``model.py``
(next to ``sapiens2.py``) by ``export_to_hf`` in ``modal_app.py``. It reconstructs
a trained run for inference and is loadable with::

    from transformers import AutoModel
    m = AutoModel.from_pretrained("<repo>", trust_remote_code=True)
    m.predict("photo.jpg")          # -> float in [score_min, score_max]
    m.predict(["a.jpg", "b.jpg"])   # -> list[float]

The pipeline is: a frozen backbone (the vendored Sapiens2 ViT, or an HF ``Auto*``
model) with a LoRA adapter applied at load time, then a small ``RegressionHead``
whose sigmoid output is mapped back to a human score.

Only the LoRA adapter and the head are stored in the exported repo; the base
backbone weights (``config.base_model_id``) are downloaded at load time. That base
is gated on the HF Hub, so the caller must have accepted its license and have a
valid ``HF_TOKEN`` in the environment.

The helper classes/functions below are deliberately COPIED from the training
library (``ratebvdeps/train_lib.py``) so this file stands alone with no dependency
on the training package — the only sibling import is the vendored backbone. The
single-dot ``from .sapiens2 import Sapiens2`` resolves both in-repo (to
``ratebvdeps.sapiens2``) and inside the published repo, where transformers' dynamic
module loader copies the relatively-imported sibling next to ``model.py``.
"""

from __future__ import annotations

import io
import os
from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn
import torchvision.transforms as T
import torchvision.transforms.functional as TF
from PIL import Image
from transformers import PretrainedConfig, PreTrainedModel

from .sapiens2 import Sapiens2

# ---------------------------------------------------------------------------
# image coercion (subset of train_lib.to_pil for inference inputs)
# ---------------------------------------------------------------------------


def _to_pil(x: Any) -> Image.Image:
    """Coerce a PIL image, filesystem path, or raw bytes into an RGB PIL image."""
    if isinstance(x, Image.Image):
        return x.convert("RGB")
    if isinstance(x, (bytes, bytearray)):
        return Image.open(io.BytesIO(bytes(x))).convert("RGB")
    if isinstance(x, str):
        return Image.open(x).convert("RGB")
    raise TypeError(f"Unsupported image input type: {type(x)}")


# ---------------------------------------------------------------------------
# backbone adapter (copied from train_lib.Sapiens2Backbone / _PooledOutput)
# ---------------------------------------------------------------------------


class _PooledOutput:
    """Tiny HF-style model output exposing ``pooler_output``."""

    def __init__(self, pooler_output, last_hidden_state=None):
        self.pooler_output = pooler_output
        self.last_hidden_state = last_hidden_state


class Sapiens2Backbone(nn.Module):
    """Adapt the vendored Sapiens2 ViT to ``forward(pixel_values=...) -> out.pooler_output``.

    Pooling is the CLS token (``out_type='cls_token'`` -> ``outs[-1]`` of shape
    ``(B, C)``), matching how the DINOv3 path consumes ``pooler_output``. LoRA wraps
    this transparently because its attention linears are ordinary modules.
    """

    def __init__(self, model: nn.Module):
        super().__init__()
        self.model = model

    def forward(self, pixel_values=None, **kwargs):
        outs = self.model(pixel_values)
        cls = outs[-1]  # out_type='cls_token' -> (B, C)
        return _PooledOutput(pooler_output=cls)


# ---------------------------------------------------------------------------
# regression head (copied verbatim from train_lib.RegressionHead)
# ---------------------------------------------------------------------------


class RegressionHead(nn.Module):
    """Maps a backbone embedding to a scalar in [0,1] via a sigmoid.

    ``hidden_dims == []`` gives the minimal head ``ReLU -> Dropout -> Linear(in,1)``;
    otherwise each hidden layer is ``Linear -> ReLU -> Dropout`` followed by a final
    ``Linear(_,1)``. A sigmoid is applied to the output.
    """

    def __init__(self, in_dim: int, hidden_dims: List[int], dropout: float):
        super().__init__()
        layers: List[nn.Module] = []
        prev = in_dim
        if len(hidden_dims) == 0:
            layers.append(nn.ReLU(inplace=True))
            layers.append(nn.Dropout(dropout))
        else:
            for h in hidden_dims:
                layers.append(nn.Linear(prev, h))
                layers.append(nn.ReLU(inplace=True))
                layers.append(nn.Dropout(dropout))
                prev = h
        layers.append(nn.Linear(prev, 1))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.net(x)
        out = torch.sigmoid(out)
        return out.squeeze(-1)


# ---------------------------------------------------------------------------
# preprocessing (copied from train_lib build_transforms eval path)
# ---------------------------------------------------------------------------


def _bound(image_size):
    """Parse image_size into a resize bound: an int or a ``(max_h, max_w)`` box."""
    if isinstance(image_size, (list, tuple)):
        return (int(image_size[0]), int(image_size[1]))
    return int(image_size)


class AspectResizeDown:
    """Aspect-preserving downscale to fit a bound; never upscales, never crops."""

    def __init__(self, bound):
        self.bound = bound

    def __call__(self, img):
        w, h = img.size  # PIL size is (W, H)
        if isinstance(self.bound, (list, tuple)):
            max_h, max_w = int(self.bound[0]), int(self.bound[1])
            scale = min(max_h / h, max_w / w, 1.0)
        else:
            s = int(self.bound)
            scale = min(s / float(max(h, w)), 1.0)
        if scale >= 1.0:
            return img  # already within bound -> never upscale
        new_h = max(1, int(round(h * scale)))
        new_w = max(1, int(round(w * scale)))
        return TF.resize(img, [new_h, new_w], antialias=True)


def _pad_to_multiple(img: torch.Tensor, patch_size: int) -> torch.Tensor:
    """Zero-pad a CHW tensor (bottom/right) so H and W are multiples of patch_size."""
    c, h, w = img.shape
    H = ((h + patch_size - 1) // patch_size) * patch_size
    W = ((w + patch_size - 1) // patch_size) * patch_size
    if H == h and W == w:
        return img
    out = img.new_zeros(c, H, W)
    out[:, :h, :w] = img
    return out


# ---------------------------------------------------------------------------
# gated base-weight download (compact copy of train_lib state-dict helpers)
# ---------------------------------------------------------------------------


def _read_state_dict(path: str) -> Dict[str, Any]:
    """Load a state dict from a .safetensors or torch checkpoint, unwrapped."""
    if path.endswith(".safetensors"):
        from safetensors.torch import load_file

        sd = load_file(path)
    else:
        sd = torch.load(path, map_location="cpu", weights_only=False)
    for key in ("state_dict", "model", "module", "weights"):
        if isinstance(sd, dict) and key in sd and isinstance(sd[key], dict):
            sd = sd[key]
            break
    return sd  # pyrefly: ignore [bad-return]


def _align_state_dict(sd: Dict[str, Any], model: nn.Module) -> Dict[str, Any]:
    """Strip a common leading prefix (e.g. ``backbone.``) so keys line up."""
    model_keys = set(model.state_dict().keys())
    if any(k in model_keys for k in sd):
        return sd
    prefixes = {k.split(".", 1)[0] + "." for k in sd if "." in k}
    for p in prefixes:
        stripped = {k[len(p) :]: v for k, v in sd.items() if k.startswith(p)}
        if any(k in model_keys for k in stripped):
            return stripped
    return sd


def _download_base_state_dict(
    model_id: str, token: Optional[str], log
) -> Optional[Dict[str, Any]]:
    """Fetch a consolidated weight file for ``model_id`` from a local dir or the Hub."""
    if os.path.isdir(model_id):
        files = os.listdir(model_id)

        def _resolve(f: str) -> str:
            return os.path.join(model_id, f)
    else:
        from huggingface_hub import hf_hub_download, list_repo_files

        files = list_repo_files(model_id, token=token)

        def _resolve(f: str) -> str:
            return hf_hub_download(model_id, f, token=token)

    weights = [f for f in files if f.endswith((".safetensors", ".pth", ".pt", ".bin"))]
    # Prefer safetensors, then the shortest name (usually the consolidated file).
    weights.sort(key=lambda f: (not f.endswith(".safetensors"), len(f)))
    if not weights:
        log(f"no weight files found for {model_id}")
        return None
    fname = weights[0]
    log(f"fetching base weights: {model_id}:{fname}")
    return _read_state_dict(_resolve(fname))


# ---------------------------------------------------------------------------
# HF config + model
# ---------------------------------------------------------------------------


class ScorerConfig(PretrainedConfig):
    """Configuration for the RateBV scorer (a frozen backbone + LoRA + a head).

    All fields are baked at export time from the training run's ``meta.json`` so
    inference needs no extra metadata fetch. ``image_size`` may be an int (square
    side) or a ``[H, W]`` bound (aspect-preserving mode).
    """

    model_type = "ratebv_scorer"

    def __init__(
        self,
        base_model_id: str = "",
        is_sapiens2: bool = False,
        sapiens2_arch: Optional[str] = None,
        patch_size: int = 16,
        image_size=224,
        preserve_aspect: bool = False,
        hidden_size: int = 768,
        head_hidden_dims: Optional[List[int]] = None,
        dropout: float = 0.0,
        score_min: float = 0.0,
        score_max: float = 1000.0,
        image_mean: Optional[List[float]] = None,
        image_std: Optional[List[float]] = None,
        sapiens2_out_type: str = "cls_token",
        sapiens2_with_cls_token: bool = True,
        sapiens2_use_tokenizer: bool = False,
        sapiens2_final_norm: bool = True,
        **kwargs,
    ):
        self.base_model_id = base_model_id
        self.is_sapiens2 = bool(is_sapiens2)  # pyrefly: ignore [unnecessary-type-conversion]
        self.sapiens2_arch = sapiens2_arch
        self.patch_size = int(patch_size)  # pyrefly: ignore [unnecessary-type-conversion]
        self.image_size = image_size
        self.preserve_aspect = bool(preserve_aspect)  # pyrefly: ignore [unnecessary-type-conversion]
        self.hidden_size = int(hidden_size)  # pyrefly: ignore [unnecessary-type-conversion]
        self.head_hidden_dims = list(head_hidden_dims) if head_hidden_dims else []
        self.dropout = float(dropout)  # pyrefly: ignore [unnecessary-type-conversion]
        self.score_min = float(score_min)  # pyrefly: ignore [unnecessary-type-conversion]
        self.score_max = float(score_max)  # pyrefly: ignore [unnecessary-type-conversion]
        self.image_mean = list(image_mean) if image_mean else [0.485, 0.456, 0.406]
        self.image_std = list(image_std) if image_std else [0.229, 0.224, 0.225]
        self.sapiens2_out_type = sapiens2_out_type
        self.sapiens2_with_cls_token = bool(sapiens2_with_cls_token)  # pyrefly: ignore [unnecessary-type-conversion]
        self.sapiens2_use_tokenizer = bool(sapiens2_use_tokenizer)  # pyrefly: ignore [unnecessary-type-conversion]
        self.sapiens2_final_norm = bool(sapiens2_final_norm)  # pyrefly: ignore [unnecessary-type-conversion]
        super().__init__(**kwargs)


class ScorerModel(PreTrainedModel):
    """A frozen backbone + LoRA adapter + regression head that scores an image.

    Loaded via ``AutoModel.from_pretrained(repo, trust_remote_code=True)``. The
    backbone and adapter are assembled in :meth:`from_pretrained` (the head skeleton
    is built in ``__init__``). Use :meth:`predict` for scores in
    ``[score_min, score_max]``; :meth:`forward` returns the raw sigmoid in ``[0, 1]``.
    """

    config_class = ScorerConfig
    base_model_prefix = "backbone"
    main_input_name = "pixel_values"
    _no_split_modules = ["TransformerEncoderLayer2"]

    def __init__(self, config: ScorerConfig):
        super().__init__(config)
        self.head = RegressionHead(
            int(config.hidden_size),  # pyrefly: ignore [unnecessary-type-conversion]
            list(config.head_hidden_dims or []),
            float(config.dropout),  # pyrefly: ignore [unnecessary-type-conversion]
        )
        self.backbone = None  # attached in from_pretrained

    # -- preprocessing ------------------------------------------------------

    def _eval_transform(self):
        mean, std = self.config.image_mean, self.config.image_std
        if self.config.preserve_aspect:
            front: List[Any] = [AspectResizeDown(_bound(self.config.image_size))]
        else:
            size = int(self.config.image_size)
            front = [T.Resize(size, antialias=True), T.CenterCrop(size)]
        return T.Compose([*front, T.ToTensor(), T.Normalize(mean=mean, std=std)])

    def _preprocess_one(self, image) -> torch.Tensor:
        """PIL/path/bytes -> a normalized CHW tensor (padded in aspect mode)."""
        x = self._eval_transform()(_to_pil(image))
        if self.config.preserve_aspect:
            x = _pad_to_multiple(x, int(self.config.patch_size))
        return x

    # -- inference ----------------------------------------------------------

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Return the head's sigmoid output in [0, 1] for a ``(B, 3, H, W)`` batch."""
        out = self.backbone(pixel_values=pixel_values)  # pyrefly: ignore [not-callable]
        return self.head(out.pooler_output)

    @torch.no_grad()
    def predict(self, images):
        """Score one image or a list of images.

        Accepts a PIL image, a filesystem path, or raw bytes (or a list thereof).
        Returns a single ``float`` for a single input, else a ``list[float]``, each
        in ``[score_min, score_max]``.
        """
        single = not isinstance(images, (list, tuple))
        items = [images] if single else list(images)
        param = next(self.parameters())
        device, dtype = param.device, param.dtype

        if self.config.preserve_aspect:
            # variable-size inputs can't be stacked -> score one at a time
            outs = []
            for it in items:
                x = self._preprocess_one(it).unsqueeze(0).to(device=device, dtype=dtype)
                outs.append(self.forward(x).reshape(-1).float())
            s01 = torch.cat(outs)
        else:
            batch = torch.stack([self._preprocess_one(it) for it in items])
            batch = batch.to(device=device, dtype=dtype)
            s01 = self.forward(batch).reshape(-1).float()

        lo, hi = float(self.config.score_min), float(self.config.score_max)
        scores = (s01 * (hi - lo) + lo).tolist()
        return scores[0] if single else scores

    # -- assembly -----------------------------------------------------------

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        """Assemble the model: base backbone (gated download) + LoRA adapter + head.

        Overrides the standard machinery because the repo stores only the adapter
        and the head — the base weights live in ``config.base_model_id``. Tolerates
        the kwargs ``AutoModel.from_pretrained`` injects (notably ``config=``).
        """
        token = kwargs.pop("token", None)
        if token is None:
            token = kwargs.pop("use_auth_token", None)
        else:
            kwargs.pop("use_auth_token", None)
        cache_dir = kwargs.pop("cache_dir", None)
        revision = kwargs.pop("revision", None)
        local_files_only = bool(kwargs.pop("local_files_only", False))
        torch_dtype = kwargs.pop("torch_dtype", None)
        device_map = kwargs.pop("device_map", None)
        injected_config = kwargs.pop("config", None)
        # Drop any remaining HF-injected kwargs we don't act on.
        kwargs.clear()

        log = print
        name = str(pretrained_model_name_or_path)

        # 1) resolve to a local directory
        if os.path.isdir(name):
            local = name
        else:
            from huggingface_hub import snapshot_download

            local = snapshot_download(
                name,
                token=token,
                cache_dir=cache_dir,
                revision=revision,
                local_files_only=local_files_only,
            )

        # 2) config (prefer the one AutoModel already built)
        if isinstance(injected_config, ScorerConfig):
            config = injected_config
        else:
            config = ScorerConfig.from_pretrained(local)

        # 3) head skeleton
        self = cls(config)

        # 4) base backbone + gated weights
        if config.is_sapiens2:
            net = Sapiens2(
                arch=config.sapiens2_arch,
                img_size=_bound(config.image_size),
                patch_size=int(config.patch_size),  # pyrefly: ignore [unnecessary-type-conversion]
                out_type=config.sapiens2_out_type,
                with_cls_token=config.sapiens2_with_cls_token,
                use_tokenizer=config.sapiens2_use_tokenizer,
                final_norm=config.sapiens2_final_norm,
            )
            sd = _download_base_state_dict(config.base_model_id, token, log)
            if sd:
                sd = _align_state_dict(sd, net)
                missing, unexpected = net.load_state_dict(sd, strict=False)
                n_total = len(net.state_dict())
                matched = n_total - len(missing)
                log(
                    f"base weights: matched={matched}/{n_total} "
                    f"missing={len(missing)} unexpected={len(unexpected)}"
                )
                if matched == 0:
                    log(
                        "WARNING: no base checkpoint keys matched -> backbone is "
                        "RANDOM (check base_model_id and your HF license/token)."
                    )
            else:
                log(
                    "WARNING: no base weights downloaded -> backbone is RANDOM "
                    f"(base_model_id={config.base_model_id!r})."
                )
            backbone = Sapiens2Backbone(net)
        else:
            from transformers import AutoModel

            backbone = AutoModel.from_pretrained(
                config.base_model_id,
                attn_implementation="sdpa",
                token=token,
                cache_dir=cache_dir,
            )

        # 5) apply the LoRA adapter (target modules come from the saved adapter config)
        #
        # The adapter lives in an ``adapter/`` subfolder rather than the repo root.
        # If ``adapter_config.json`` sat at the root, transformers' PEFT auto-detection
        # (``find_adapter_config_file``) would intercept ``AutoModel.from_pretrained``
        # *before* this custom ``from_pretrained`` ran, reassign the load path to the
        # adapter's (null) ``base_model_name_or_path``, and fail. Keeping it in a
        # subfolder dodges that; the fallback keeps older flat exports loadable.
        from peft import PeftModel

        adapter_dir = os.path.join(local, "adapter")
        if not os.path.isdir(adapter_dir):
            adapter_dir = local  # backward-compat: pre-subfolder flat exports
        self.backbone = PeftModel.from_pretrained(backbone, adapter_dir)

        # 6) head weights
        head_sd = torch.load(os.path.join(local, "head.pt"), map_location="cpu")
        self.head.load_state_dict(head_sd)

        # 7) dtype / device (best effort; full device_map sharding is unsupported)
        if torch_dtype is not None and torch_dtype != "auto":
            self = self.to(torch_dtype)
        if isinstance(device_map, str) and device_map in ("auto", "cuda"):
            if torch.cuda.is_available():
                self = self.to("cuda")

        self.eval()
        return self