"""DroidDetect loader — ModernBERT 4-class code-authorship detector.

The published checkpoints are NOT `AutoModelForSequenceClassification`. Their
config.json is a 4-line stub (`"model_type": "custom_model"`), so the wrapper
has to be rebuilt by hand. Everything it needs is inferable from the tensor
shapes, so this file hardcodes no dimensions:

    text_encoder.*        ModernBERT encoder      (768 hidden = Base, 1024 = Large)
    text_projection       hidden -> proj_dim      (128 on Base, 256 on Large)
    classifier            proj_dim -> 4 classes

DroidDetect-Large additionally ships a *second complete encoder* under
`additional_loss.sentence_embedder.*` — 395M parameters used only for the
training-time triplet loss. It is dropped on load, which halves resident memory
with no effect on predictions. That is what makes Large fit the free CPU tier.
"""

import os

# `spaces` must be imported BEFORE torch: on ZeroGPU it installs a CUDA
# emulation layer so `.to("cuda")` at module level works without a real GPU
# attached. Absent (a plain CPU Space, or local dev) the fallback below makes
# the decorator a no-op, so one codebase runs on both.
try:
    import spaces

    _gpu = spaces.GPU
except ImportError:  # not a ZeroGPU Space

    def _gpu(fn=None, **_kwargs):
        """No-op stand-in for @spaces.GPU."""
        return fn if fn is not None else (lambda f: f)


import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModel, AutoTokenizer

# HF sets this on ZeroGPU hardware. It is the switch for everything below.
ZERO_GPU = bool(os.environ.get("SPACES_ZERO_GPU"))
DEVICE = "cuda" if ZERO_GPU else "cpu"

# Seconds of GPU the decorated function may hold. Ours is one forward pass over
# at most 2048 tokens — well under a second on the Blackwell slice. Asking for
# less than the 60s default improves this Space's queue priority.
GPU_DURATION = int(os.environ.get("GPU_DURATION", "30"))

# Large is the default: +2 points of F1 over Base, and it still fits the free
# CPU tier once the training-only twin encoder is dropped (~1.7 GB resident of
# the 16 GB available). Set DROID_REPO=project-droid/DroidDetect-Base to trade
# accuracy for ~4x faster inference.
REPO = os.environ.get("DROID_REPO", "project-droid/DroidDetect-Large")

# ModernBERT handles 8192 tokens, but on 2 shared vCPUs every token costs real
# wall-clock. Contest submissions are a few hundred tokens; 2048 leaves large
# headroom and bounds the worst case to a few seconds. Raise it only if you
# actually see `truncated: true` in production.
MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "2048"))

# How the encoder's token states are collapsed into one vector before the head.
# The checkpoint does not record this and the model card does not say. Run
# smoke_test.py once — it tries both and tells you which one this checkpoint
# was trained with. See DEPLOY.md step 5.
POOLING = os.environ.get("POOLING", "mean")

# CPU Spaces get 2 vCPUs. Letting torch spawn a thread per core it *detects* on
# a shared host causes contention that makes inference slower, not faster.
# Irrelevant on ZeroGPU, where the forward pass runs on the GPU.
if not ZERO_GPU:
    torch.set_num_threads(int(os.environ.get("TORCH_THREADS", "2")))

# Label taxonomy + the coarse collapses live in labels.py so they can be tested
# without torch. Re-exported here so callers have one import site.
from labels import (  # noqa: E402  (kept with the other detector constants)
    BINARY_LABELS,
    BINARY_OF,
    LABELS,
    TERNARY_LABELS,
    TERNARY_MEANING,
    TERNARY_OF,
    collapse,
)

__all__ = [
    "BINARY_LABELS",
    "BINARY_OF",
    "LABELS",
    "MAX_LENGTH",
    "POOLING",
    "REPO",
    "TERNARY_LABELS",
    "TERNARY_MEANING",
    "TERNARY_OF",
    "classify",
    "collapse",
    "load_model",
]

# hidden size -> the backbone that produced it
BACKBONES = {
    768: "answerdotai/ModernBERT-base",
    1024: "answerdotai/ModernBERT-large",
}

# Checkpoint filenames to try, in order.
WEIGHT_FILES = ("model.safetensors", "pytorch_model.bin")


class TLModel(nn.Module):
    """Encoder -> projection -> classifier, matching the checkpoint's layout."""

    def __init__(self, encoder, hidden, proj_dim, num_classes):
        super().__init__()
        self.text_encoder = encoder
        self.text_projection = nn.Linear(hidden, proj_dim)
        self.classifier = nn.Linear(proj_dim, num_classes)

    def forward(self, input_ids, attention_mask, pooling=POOLING):
        states = self.text_encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
        ).last_hidden_state

        if pooling == "cls":
            pooled = states[:, 0]
        else:
            # Mean over real tokens only — padding must not drag the vector.
            mask = attention_mask.unsqueeze(-1).to(states.dtype)
            pooled = (states * mask).sum(1) / mask.sum(1).clamp(min=1e-9)

        embedding = self.text_projection(pooled)
        return self.classifier(embedding), embedding


tokenizer = None
model = None


def _download_state():
    """Fetch the checkpoint, preferring safetensors. Returns a state dict."""
    last = None
    for name in WEIGHT_FILES:
        try:
            path = hf_hub_download(repo_id=REPO, filename=name)
        except Exception as err:  # not published in this format — try the next
            last = err
            continue
        if name.endswith(".safetensors"):
            from safetensors.torch import load_file

            return load_file(path, device="cpu")
        return torch.load(path, map_location="cpu", weights_only=True)
    raise RuntimeError(f"No usable checkpoint found in {REPO}: {last}")


def load_model():
    """Download and build the detector. Idempotent; safe to call repeatedly."""
    global tokenizer, model
    if model is not None:
        return

    state = _download_state()

    # Training-only twin encoder — never used at inference. Dropping it takes
    # Large from ~790M resident parameters to ~395M.
    state = {k: v for k, v in state.items() if not k.startswith("additional_loss.")}

    proj_dim, hidden = state["text_projection.weight"].shape
    num_classes = state["classifier.weight"].shape[0]

    backbone = BACKBONES.get(hidden)
    if backbone is None:
        raise RuntimeError(f"Unrecognized encoder hidden size {hidden} in {REPO}")

    # from_config, not from_pretrained: the checkpoint carries every encoder
    # weight already, so pulling the base model would be a wasted download.
    config = AutoConfig.from_pretrained(backbone)
    # ModernBERT torch.compile()s part of itself when this is on. Must stay off
    # on BOTH targets: ZeroGPU does not support torch.compile at all, and on a
    # CPU Space it's a long first-request stall (or an outright failure on some
    # torch builds) for no throughput gain.
    config.reference_compile = False

    try:
        # Ask for SDPA explicitly so transformers never tries to import
        # flash-attn, which isn't installed and is CUDA-only anyway.
        encoder = AutoModel.from_config(config, attn_implementation="sdpa")
    except (TypeError, ValueError):
        # Older transformers without the kwarg — the CPU default is fine.
        encoder = AutoModel.from_config(config)

    net = TLModel(encoder, hidden, proj_dim, num_classes)
    # strict: a silently mismatched head would return confident nonsense.
    net.load_state_dict(state, strict=True)
    net.float()
    net.eval()
    # On ZeroGPU this must happen at startup, not inside the @spaces.GPU
    # function: CUDA transfers are optimised for placements done during
    # startup, and lazy moves are explicitly discouraged. Outside a GPU
    # function this is absorbed by the CUDA emulation layer; inside one, the
    # weights are already really there.
    net.to(DEVICE)

    tokenizer = AutoTokenizer.from_pretrained(REPO)
    model = net


@_gpu(duration=GPU_DURATION)
@torch.no_grad()
def classify(code_text, pooling=None):
    """Return (probabilities dict, embedding list, token count, truncated).

    On ZeroGPU this whole body runs in a forked worker with a real GPU
    attached. Everything it returns is plain Python (dict / list / int / bool)
    so it survives the process boundary.
    """
    load_model()

    encoded = tokenizer(
        [code_text],
        truncation=True,
        max_length=MAX_LENGTH,
        return_tensors="pt",
    )
    input_ids = encoded["input_ids"].to(DEVICE)
    attention_mask = encoded["attention_mask"].to(DEVICE)

    logits, embedding = model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        pooling=pooling or POOLING,
    )
    probs = torch.softmax(logits[0], dim=-1)

    return (
        {label: round(float(p), 6) for label, p in zip(LABELS, probs)},
        [round(float(v), 4) for v in embedding[0]],
        int(attention_mask.sum().item()),
        bool(input_ids.shape[1] >= MAX_LENGTH),
    )


# ZeroGPU needs the weights resident and placed on cuda at ROOT MODULE LEVEL,
# before the first @spaces.GPU call — lazy placement inside the decorated
# function is explicitly discouraged and much slower. On CPU we stay lazy:
# app.py warm-loads in __main__ instead.
if ZERO_GPU:
    load_model()