"""Penny client: typed decisions from a local model. No dependencies.

Works with Ollama (default, http://localhost:11434) and any OpenAI-compatible
server that returns logprobs (llama.cpp's llama-server, LM Studio).

    from penny import Penny

    penny = Penny()                      # Ollama, model "penny"
    penny.choose(order, "What should we do with this refund request?",
                 ["Refund in full", "Offer store credit", "Decline"])
    # {'choice': 'Refund in full', 'index': 0, 'confidence': 0.91, 'probs': [0.91, 0.07, 0.02]}
    penny.yes_no(payout, "Should we release this payout now?")
    # {'yes': False, 'confidence': 0.88}
    penny.score(payment, "How likely is this payment to be fraud? 1 = very unlikely, 5 = very likely")
    # {'score': 4.3, 'confidence': 0.62, 'probs': [0.01, 0.02, 0.08, 0.45, 0.44]}

Each call is one forward pass: Penny never generates text. The confidence is
the model's calibrated probability for its answer; send low-confidence
answers to a person.
"""

from __future__ import annotations

import json
import math
import urllib.request
from typing import Any

LETTERS = "ABCDEFGH"
SCORES = "12345"
# Fitted on Penny-1's validation split (scripts/train.py, temperatures.json).
TEMPERATURES = {"choice": 1.1985, "yes_no": 1.0396, "score": 0.935}


def render(state: Any, question: str, qtype: str, options: list[str] | None = None) -> str:
    if not isinstance(state, str):
        state = json.dumps(state, indent=2, default=str)
    lines = ["<state>", state.strip(), "</state>", f'<question type="{qtype}">', question.strip()]
    if qtype == "choice":
        lines += [f"{LETTERS[i]}. {o.strip()}" for i, o in enumerate(options or [])]
    elif qtype == "yes_no":
        lines += ["A. Yes", "B. No"]
    else:
        lines.append("Answer from 1 (lowest) to 5 (highest).")
    return "\n".join(lines + ["</question>", "<answer>", ""])


class Penny:
    def __init__(self, model: str = "penny", host: str = "http://localhost:11434", api: str = "ollama",
                 timeout: float = 60.0):
        """api: "ollama" (native /api/generate) or "openai" (/v1/completions: llama-server, LM Studio)."""
        self.model, self.host, self.api, self.timeout = model, host.rstrip("/"), api, timeout

    # -- public -----------------------------------------------------------------

    def choose(self, state: Any, question: str, options: list[str]) -> dict:
        if not 2 <= len(options) <= len(LETTERS):
            raise ValueError(f"choose() takes 2-{len(LETTERS)} options")
        probs = self._ask(render(state, question, "choice", options), list(LETTERS[: len(options)]), "choice")
        i = max(range(len(probs)), key=probs.__getitem__)
        return {"choice": options[i], "index": i, "confidence": round(probs[i], 4), "probs": [round(p, 4) for p in probs]}

    def yes_no(self, state: Any, question: str) -> dict:
        p_yes, p_no = self._ask(render(state, question, "yes_no"), ["A", "B"], "yes_no")
        return {"yes": p_yes >= p_no, "confidence": round(max(p_yes, p_no), 4)}

    def score(self, state: Any, question: str) -> dict:
        probs = self._ask(render(state, question, "score"), list(SCORES), "score")
        ev = sum(p * (i + 1) for i, p in enumerate(probs))
        k = round(ev) - 1
        # confidence: probability that the true score is within one point of the call
        return {"score": round(ev, 2), "confidence": round(sum(probs[max(0, k - 1): k + 2]), 4),
                "probs": [round(p, 4) for p in probs]}

    # -- internals --------------------------------------------------------------

    def _ask(self, prompt: str, allowed: list[str], qtype: str) -> list[float]:
        top = self._top_logprobs(prompt)
        t = TEMPERATURES.get(qtype, 1.0)
        logits = [top.get(a, -1e9) / t for a in allowed]
        m = max(logits)
        exps = [math.exp(x - m) for x in logits]
        z = sum(exps)
        return [e / z for e in exps]

    def _top_logprobs(self, prompt: str) -> dict[str, float]:
        if self.api == "ollama":
            res = self._post("/api/generate", {
                "model": self.model, "prompt": prompt, "raw": True, "stream": False,
                "logprobs": True, "top_logprobs": 20, "options": {"num_predict": 1, "temperature": 0}})
            entries = res["logprobs"][0]["top_logprobs"]
            pairs = [(e["token"], e["logprob"]) for e in entries]
        else:
            res = self._post("/v1/completions", {
                "model": self.model, "prompt": prompt, "max_tokens": 1, "temperature": 0, "logprobs": 20})
            lp = res["choices"][0]["logprobs"]
            if "top_logprobs" in lp and isinstance(lp["top_logprobs"], list) and lp["top_logprobs"]:
                pairs = list(lp["top_logprobs"][0].items())
            else:  # newer servers: content[0].top_logprobs
                pairs = [(e["token"], e["logprob"]) for e in lp["content"][0]["top_logprobs"]]
        out: dict[str, float] = {}
        for tok, lp in pairs:
            key = tok.strip()
            if key and (key not in out or lp > out[key]):
                out[key] = lp
        return out

    def _post(self, path: str, body: dict) -> dict:
        req = urllib.request.Request(self.host + path, data=json.dumps(body).encode(),
                                     headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=self.timeout) as r:
            return json.load(r)