"""
jev_toy/serve.py

A mirror of the System One request/response API, backed by our toy model.
Shows the proven interface in action:

    POST: { state: "...", questions: { <name>: {type, instructions, ...} } }
    ->     { <name>: { <typed result with probabilities> } }

All questions are answered from ONE encoder pass over the state (parallel).

Usage:
  source .venv/bin/activate
  python -m jev_toy.serve --ckpt checkpoints/model.pt
"""

from __future__ import annotations

import argparse

import torch
import torch.nn.functional as F

from jev_toy.data import tokenize
from jev_toy.model import SystemOneConfig, SystemOneModel


def build_stoi(vocab_dict):
    return vocab_dict


class ToyServer:
    def __init__(self, ckpt, device="cpu"):
        self.device = device
        ck = torch.load(ckpt, map_location=device)
        self.cfg = SystemOneConfig(**ck["config"])
        self.stoi = ck["vocab"]
        # __oov__ key stores the oov index; default 1
        self.oov = ck["vocab"].get("__oov__", 1)
        self.pad = self.cfg.pad_token_id
        self.model = SystemOneModel(self.cfg).to(device)
        self.model.load_state_dict(ck["state_dict"])
        self.model.eval()

    def _enc(self, text, keep=None):
        ids, mask = tokenize(text, self.stoi, self.cfg.max_seq_len, self.oov, self.pad)
        return (
            torch.tensor([ids], device=self.device),
            torch.tensor([mask], device=self.device),
        )

    @torch.no_grad()
    def answer(self, state: str, questions: dict):
        """questions: {name: {type, instructions, options?}}. Returns {name: decision}."""
        s_ids, s_mask = self._enc(state)
        h_state = self.model.encode_state(s_ids, s_mask)  # ONE pass over state
        results = {}
        for name, q in questions.items():
            q_ids, q_mask = self._enc(q.get("instructions", name))  # type: ignore
            # h_state is [1, D]; one question => one row, so pass it directly.
            logits, _ = self.model.answer(h_state, q_ids, q_mask, q["type"])
            t = q["type"]
            if t == "noul":
                p = torch.sigmoid(logits["noul"][0]).item()
                results[name] = {"type": "noul", "noul": round(p, 4), "is_true": p >= 0.5}
            elif t == "choice":
                opts = q.get("options") or []
                probs = logits["choice"][0]
                k = len(opts)
                if k == 0:
                    k = probs.shape[0]
                    opts = [f"option_{i}" for i in range(k)]
                probs = probs[:k]
                probs = F.softmax(probs / probs.sum().clamp_min(1e-9), dim=-1)  # renormalize over given options
                dist = {str(opts[i]): float(f"{probs[i].item():.4f}") for i in range(k)}
                # confidence = margin above uniform (community formula for choice confidence)
                u = 1.0 / k
                conf = float((probs.max() - u) / (1 - u))
                results[name] = {
                    "type": "choice", "distribution": dist,
                    "argmax": str(opts[int(probs.argmax())]),
                    "confidence": round(conf, 4),
                }
            elif t == "score":
                lo, hi = self.cfg.score_range
                p = (hi - lo) * torch.sigmoid(logits["score"][0]) + lo
                results[name] = {"type": "score", "score": round(float(p), 3)}
            else:
                raise ValueError(t)
        return results


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="checkpoints/model.pt")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    server = ToyServer(args.ckpt, device)

    # --- demo 1: support-ticket-style multi-question (the Jev flagship use case) ---
    demo_state = "I've been trying to connect my Stripe account for 3 days, it keeps failing with a 403. I'm losing sales and my manager is anxious. Please help ASAP."
    demo_qs = {
        "is_urgent": {"type": "noul", "instructions": "The message conveys urgency or time-sensitivity."},
        "department": {
            "type": "choice",
            "instructions": "Which team should handle this?",
            "options": ["billing", "technical", "sales"],
        },
        "frustration": {
            "type": "score",
            "instructions": "How frustrated does the customer appear? (0=civil, 1=angry)",
        },
    }
    print("=== demo: one state, three parallel typed questions ===")
    out = server.answer(demo_state, demo_qs)
    for k, v in out.items():
        print(f"  {k}: {v}")

    # --- demo 2: routing a news headline ---
    print("\n=== demo: route a news article ===")
    news = "News article: NVIDIA stock hits all-time high after beating earnings estimates, analysts raise price targets."
    ro = server.answer(news, {
        "topic": {
            "type": "choice", "instructions": "Which topic does this article belong to?",
            "options": ["World", "Sports", "Business", "Sci/Tech"],
        }
    })
    print("  topic:", ro["topic"])


if __name__ == "__main__":
    main()