"""ClefModel + optional BoxHead. Kept free of transformers/peft imports so it can be unit-tested with a mock backbone."""
from __future__ import annotations
from typing import Any
import torch
from .joint_schema_model import ClefModel, JointSchemaHead
from .box import BoxHead, BoxPrediction


class GlintModel(ClefModel):
    """`forward(batch)` returns the joint-schema logits exactly like ClefModel (list per record).
    `decide(batch)` returns `(logits, boxes)`; `boxes[i]` is a BoxPrediction or None."""

    def __init__(self, language_model: Any, head: JointSchemaHead, box_head: BoxHead | None = None) -> None:
        super().__init__(language_model, head)
        self.box_head = box_head

    def _base_model(self):
        lm = self.language_model
        return lm.get_base_model() if hasattr(lm, "get_base_model") else lm

    def _hidden(self, batch: dict[str, Any]):
        base_model = self._base_model()
        media = batch.get("media") or {}
        text_model = base_model.model
        if not media and hasattr(text_model, "language_model"):
            text_model = text_model.language_model
        outputs = text_model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"],
                             use_cache=False, return_dict=True, **media)
        return outputs.last_hidden_state, base_model

    def decide(self, batch: dict[str, Any]) -> tuple[list[list[torch.Tensor]], list[BoxPrediction | None]]:
        hidden, base_model = self._hidden(batch)
        records = batch["records"]
        logits: list[list[torch.Tensor]] = [[] for _ in records]
        keep = [i for i, r in enumerate(records) if r.questions]
        if keep:
            sel = keep if len(keep) < len(records) else slice(None)
            out = self.head(hidden[sel], batch["input_ids"][sel], batch["attention_mask"][sel], [records[i] for i in keep],
                            base_model.get_output_embeddings().weight)
            for i, z in zip(keep, out):
                logits[i] = z
        boxes: list[BoxPrediction | None] = [None] * len(records)
        if any(getattr(r, "box_questions", ()) for r in records):
            if self.box_head is None:
                raise ValueError("record has box questions but this checkpoint has no box head (train with --box-head on)")
            boxes = self.box_head(hidden, batch["attention_mask"], records)
        return logits, boxes

    def forward(self, batch: dict[str, Any]) -> list[list[torch.Tensor]]:
        return self.decide(batch)[0]