"""Standalone loader for the SporeLabs Conversation Cues (Cue) text model.

Depends only on torch, transformers and safetensors.

    from cue import CueModel
    m = CueModel.from_pretrained(".")            # the folder that holds config.json
    m.score("I'll send the deck by Friday.",
            context=[{"speaker": "A", "text": "Can someone own the budget deck?"}], speaker="B")
    # {'commitment': 0.97, 'open_question': 0.01, 'delegate': 0.02, ...}

Folder layout: config.json, thresholds.json, backbone/ (ModernBERT encoder with its pooling head, tokenizer files),
heads.safetensors (one Linear(hidden, 1) per head). The forward pass is: encoder -> mean pool -> dense/GELU/LayerNorm
head -> one logit per question -> sigmoid.
"""

from __future__ import annotations

import json
import os
import re
from pathlib import Path

import torch

SEP = " [SEP] "


def format_input(text: str, context: list[dict] | None = None, speaker: str | None = None, n_ctx: int = 3) -> str:
    """One string: `<target> [SEP] <rel>: <previous turn> [SEP] <rel>: <turn before> ...`, most recent turn first.

    <rel> is `same` (same speaker as the target), `other`, or `?` when the target speaker is unknown.
    """
    ctx = list(context or [])[-n_ctx:] if n_ctx > 0 else []
    parts = [text or "."]
    for c in reversed(ctx):
        if speaker is None or c.get("speaker") is None:
            rel = "?"
        else:
            rel = "same" if c.get("speaker") == speaker else "other"
        parts.append(f"{rel}: {c.get('text', '')}")
    return SEP.join(parts)


def parse_turn(line: str) -> dict:
    """'A: hello there' -> {'speaker': 'A', 'text': 'hello there'}; a line without 'X:' has no speaker."""
    m = re.match(r"^\s*([^:\s][^:]{0,30}?)\s*:\s*(.*)$", line)
    if m:
        return {"speaker": m.group(1).strip(), "text": m.group(2).strip()}
    return {"speaker": None, "text": line.strip()}


class CueModel:
    def __init__(self, backbone, tokenizer, heads: dict, config: dict, thresholds: dict, device: str):
        self.backbone, self.tok, self.heads = backbone, tokenizer, heads
        self.config, self.thresholds, self.device = config, thresholds, device
        self.head_names = list(heads)

    # ------------------------------------------------------------------ loading
    @classmethod
    def from_pretrained(cls, path_or_repo: str | os.PathLike, device: str = "cpu") -> "CueModel":
        from safetensors.torch import load_file
        from transformers import AutoModelForSequenceClassification, AutoTokenizer

        root = Path(path_or_repo)
        config = json.loads((root / "config.json").read_text())
        tp = root / "thresholds.json"
        thresholds = json.loads(tp.read_text()) if tp.exists() else {}
        bb_dir = root / config.get("backbone_dir", "backbone")
        backbone = AutoModelForSequenceClassification.from_pretrained(bb_dir)  # ModernBertForSequenceClassification
        backbone.to(device).eval()
        tok = AutoTokenizer.from_pretrained(bb_dir)
        sd = load_file(str(root / config.get("heads_file", "heads.safetensors")))
        d = backbone.config.hidden_size
        heads = {}
        for q in config["heads"]:
            lin = torch.nn.Linear(d, 1)
            lin.load_state_dict({"weight": sd[f"{q}.weight"], "bias": sd[f"{q}.bias"]})
            heads[q] = lin.to(device).eval()
        return cls(backbone, tok, heads, config, thresholds, device)

    # ------------------------------------------------------------------ forward
    @torch.no_grad()
    def _logits(self, input_ids, attention_mask) -> dict[str, torch.Tensor]:
        b = self.backbone
        h = b.model(input_ids=input_ids, attention_mask=attention_mask)[0]
        if b.config.classifier_pooling == "cls":
            h = h[:, 0]
        else:
            m = attention_mask.unsqueeze(-1).to(h.dtype)
            h = (h * m).sum(1) / m.sum(1)
        v = b.drop(b.head(h))
        return {q: head(v)[:, 0] for q, head in self.heads.items()}

    def probs_from_strings(self, texts: list[str], batch: int = 64) -> dict[str, list[float]]:
        """Sigmoid scores per head for already formatted input strings."""
        max_len = int(self.config.get("max_tokens", 160))
        out = {q: [0.0] * len(texts) for q in self.head_names}
        order = sorted(range(len(texts)), key=lambda i: len(texts[i]))
        for i in range(0, len(order), batch):
            ix = order[i:i + batch]
            enc = self.tok([texts[j] or "." for j in ix], padding=True, truncation=True, max_length=max_len,
                           return_tensors="pt").to(self.device)
            for q, z in self._logits(enc["input_ids"], enc["attention_mask"]).items():
                for j, v in zip(ix, torch.sigmoid(z.float()).cpu().tolist()):
                    out[q][j] = 0.0 if not texts[j] else float(v)
        return out

    def score_many(self, rows: list[dict], batch: int = 64) -> list[dict[str, float]]:
        """rows: [{text, context?: [{speaker, text}], speaker?}] -> one {head: probability} dict per row."""
        n_ctx = int(self.config.get("n_ctx", 3))
        xs = [format_input(r["text"], r.get("context"), r.get("speaker"), n_ctx) for r in rows]
        p = self.probs_from_strings(xs, batch)
        return [{q: p[q][i] for q in self.head_names} for i in range(len(rows))]

    def score(self, text: str, context: list[dict] | None = None, speaker: str | None = None) -> dict[str, float]:
        return self.score_many([{"text": text, "context": context, "speaker": speaker}])[0]

    # ------------------------------------------------------------------ decisions
    def threshold(self, head: str) -> float | None:
        t = (self.thresholds.get("heads") or {}).get(head) or {}
        return t.get("threshold")

    def decide(self, scores: dict[str, float]) -> dict[str, bool | None]:
        """Fired / not fired at the dev thresholds (None when a head has no threshold yet)."""
        return {q: (None if self.threshold(q) is None else s >= self.threshold(q)) for q, s in scores.items()}

    @property
    def product_heads(self) -> list[str]:
        return list(self.config.get("product_heads") or self.head_names)