"""Public document-state API and four-state answer semantics."""

import copy
import json
import math
from pathlib import Path

from ._vendor.contract import _state_parts, decision, parse_questions, present
from ._vendor.prompts import boolean_block, label_block, listwise_block
from ._vendor.semantics import listed_probs, p_yes
from .artifacts import digest, sha256

TASKS = ("boolean", "single", "ordered", "multilabel", "entity")


def branches(spec):
    task, req = spec["task"], spec["request"]
    if task == "boolean":
        return [(boolean_block(req["question"]), 4, "boolean/state4")]
    if task in ("single", "ordered"):
        block, width = listwise_block(
            req["question"], req["options"], ordered=task == "ordered", reserved=task == "single"
        )
        return [(block, width, task + ("/choiceR" if task == "single" else "/choiceS"))]
    if task == "entity":
        return [
            (boolean_block(req["template"].replace("{entity}", c)), 4, "entity/state4")
            for c in req["entities"]
        ]
    return [(*label_block(req["question"], c), "multilabel/state4") for c in req["labels"]]


def distributions(spec, rows, temperature):
    if spec["task"] in ("boolean", "entity", "multilabel"):
        values = [p_yes(r["letter_logits"], temperature) for r in rows]
        return [[p, 1 - p] for p in values]
    return [listed_probs(rows[0]["letter_logits"], len(spec["texts"]), temperature).tolist()]


def ordering_score(values):
    """Product of per-unit top probabilities; not a calibrated joint probability."""
    if not values:
        raise ValueError("At least one answer unit is required")
    for p in values:
        if len(p) < 2 or not all(math.isfinite(v) and v >= 0 for v in p) or abs(sum(p) - 1) > 1e-6:
            raise ValueError("Invalid answer distribution")
    return math.prod(max(p) for p in values)


class DocumentState:
    def __init__(self, owner, data):
        self._owner, self._data, self.closed = owner, data, False
        self.image_hashes = {p["image"]: sha256(p["image"]) for p in data["parts"] if "image" in p}

    @property
    def prefix_tokens(self):
        self._check()
        return len(self._data["prefix_ids"])

    def _check(self):
        if self.closed:
            raise ValueError("Document state is closed")
        if any(sha256(path) != value for path, value in self.image_hashes.items()):
            raise ValueError("Document image changed after prefill")

    def save(self, path):
        """Save a source-bound replay recipe, never pickle executable cache objects."""
        self._check()
        body = {
            "format": "solomon-mlx-replay-v1",
            "runtime": self._owner.identity["fingerprint"],
            "parts": self._data["parts"],
            "image_hashes": self.image_hashes,
            "prefix_ids_sha256": digest(self._data["prefix_ids"]),
        }
        Path(path).write_text(json.dumps({**body, "sha256": digest(body)}, indent=2))

    def close(self):
        with self._owner.engine.lock:
            self._data.clear()
            self.closed = True

    def __enter__(self):
        self._check()
        return self

    def __exit__(self, *args):
        self.close()


class Solomon:
    @classmethod
    def load(
        cls,
        model_dir,
        profile="quality",
        *,
        chunk_size=2048,
        max_tokens=40960,
        page_selector=None,
        calibration=None,
    ):
        if profile != "quality":
            raise ValueError("Only full BF16 quality is implemented; quantization is secondary")
        from .engine import Engine

        return cls(
            Engine(model_dir, chunk_size=chunk_size, max_tokens=max_tokens),
            page_selector=page_selector,
            calibration=calibration,
        )

    def __init__(self, engine, *, page_selector=None, calibration=None):
        self.engine, self.identity, self.page_selector = engine, engine.identity, page_selector
        self.temperatures = dict.fromkeys(TASKS, 1.0)
        self.calibration_status = "uncalibrated"
        if calibration is not None:
            artifact = json.loads(Path(calibration).read_text())
            payload = {k: v for k, v in artifact.items() if k != "sha256"}
            if (
                artifact.get("sha256") != digest(payload)
                or artifact["runtime"] != self.identity["fingerprint"]
            ):
                raise ValueError("Calibration checksum or MLX runtime identity mismatch")
            temps = artifact["temperatures"]
            if set(temps) != set(TASKS) or any(
                isinstance(v, bool)
                or not isinstance(v, (int, float))
                or not math.isfinite(v)
                or not 0 < v <= 20
                for v in temps.values()
            ):
                raise ValueError("Invalid temperatures")
            self.temperatures, self.calibration_status = temps, "profile_fitted"

    def prefill(self, document):
        parts = copy.deepcopy(_state_parts(document))
        if not parts:
            parts = [{"text": ""}]
        for p in parts:
            if not isinstance(p, dict) or set(p) not in ({"text"}, {"image"}):
                raise ValueError("Each document part must contain only text or image")
            if "text" in p and not isinstance(p["text"], str):
                raise ValueError("Text parts must be strings")
            if "image" in p:
                p["image"] = str(Path(p["image"]).resolve(strict=True))
        hashes = {p["image"]: sha256(p["image"]) for p in parts if "image" in p}
        state = DocumentState(self, self.engine.prefill(parts))
        if state.image_hashes != hashes:
            state.close()
            raise ValueError("Image changed while document was being prefilled")
        return state

    def replay(self, path):
        body = json.loads(Path(path).read_text())
        expected = body.pop("sha256")
        if (
            digest(body) != expected
            or body["format"] != "solomon-mlx-replay-v1"
            or body["runtime"] != self.identity["fingerprint"]
        ):
            raise ValueError("Replay checksum or runtime mismatch")
        if any(sha256(p) != h for p, h in body["image_hashes"].items()):
            raise ValueError("Replay image changed")
        state = self.prefill(body["parts"])
        if digest(state._data["prefix_ids"]) != body["prefix_ids_sha256"]:
            state.close()
            raise ValueError("Replay tokenization differs")
        return state

    def _answer(self, state, spec, execution="cached"):
        rows = [self.engine.ask(state._data, b, n, h, execution=execution) for b, n, h in branches(spec)]
        dists = distributions(spec, rows, self.temperatures[spec["task"]])
        return {
            **present(spec, dists),
            "ordering_score": ordering_score(dists),
            "temperature": self.temperatures[spec["task"]],
        }, rows

    def decide(
        self,
        *,
        state,
        questions,
        evidence="support",
        evidence_max_calls=64,
        execution="cached",
        diagnostics=False,
    ):
        if not isinstance(state, DocumentState) or state._owner is not self:
            raise ValueError("State belongs to a different model instance")
        if evidence not in ("none", "support", "sufficiency", "removal"):
            raise ValueError("Invalid evidence level")
        if type(evidence_max_calls) is not int or not 0 <= evidence_max_calls <= 512:
            raise ValueError("Invalid evidence call budget")
        specs = parse_questions(questions)
        with self.engine.lock:
            state._check()
            answers, usage = {}, {"branches": 0, "input_tokens": 0, "evidence_calls": 0}
            for spec in specs:
                answer, rows = self._answer(state, spec, execution)
                body = self._evidence(state, spec, answer, evidence, evidence_max_calls)
                answer.update(
                    evidence=body["references"], evidence_status=body["status"], evidence_detail=body
                )
                if diagnostics:
                    answer["branches"] = rows
                answers[spec["id"]] = answer
                usage["branches"] += len(rows)
                usage["input_tokens"] += sum(
                    r["branch_tokens"] if execution == "cached" else r["prompt_tokens"] for r in rows
                )
                usage["evidence_calls"] += body.get("calls", 0)
            return {
                "answers": answers,
                "usage": usage,
                "runtime": self.identity,
                "calibration_status": self.calibration_status,
                "answer_policy": "always_answers",
            }

    def _fresh(self, document, spec):
        with self.prefill(document) as state:
            answer, _ = self._answer(state, spec)
            return decision(spec, answer)

    def _evidence(self, state, spec, answer, level, budget):
        from ._vendor import evidence_v3 as v3
        from ._vendor.evidence import image_pages, validate_pages, validate_spans
        from ._vendor.retrieval import lexical_select, remove

        body = {
            "references": [],
            "status": "not_requested",
            "calls": 0,
            "verification": "none",
            "faithfulness_established": False,
        }
        if level == "none":
            return body
        parts, req = state._data["parts"], spec["request"]
        task = spec["task"]
        if task == "entity":
            questions = [req["template"].replace("{entity}", c) for c in req["entities"]]
        elif task == "multilabel":
            questions = [req["question"] + " Label: " + c for c in req["labels"]]
        else:
            questions = [req["question"] + (" " + " ".join(req["options"]) if "options" in req else "")]
        images = [p["image"] for p in parts if "image" in p]
        needed = (1 if images else len(questions)) if level in ("sufficiency", "removal") else 0
        needed += int(level == "removal")
        if needed > budget:
            return {**body, "status": "budget_exhausted", "required_calls": needed}
        baseline = decision(spec, answer)
        if images:
            if self.page_selector is None:
                return {**body, "status": "unsupported_page_selector", "pages_available": len(images)}
            pages = image_pages(images)
            selector = self.page_selector
            plan = None
            if hasattr(selector, "plan"):
                plan = selector.plan(
                    pages, questions, **({"task": task} if getattr(selector, "task_aware", False) else {})
                )
                if type(plan.get("calls")) is not int or plan["calls"] < 0:
                    raise ValueError("Invalid page selector call estimate")
                if needed + plan["calls"] > budget:
                    return {**body, "status": "budget_exhausted", "required_calls": needed + plan["calls"]}
            selection = (
                selector.execute(plan) if plan is not None else selector(copy.deepcopy(pages), questions)
            )
            calls = selection.get("cost", {}).get("calls", 0)
            if calls != (plan["calls"] if plan is not None else 0):
                raise ValueError("Page selector exceeded its declared call budget")
            refs = validate_pages(pages, selection["evidence"])
            body["calls"] = calls
            selected = {r["page"] for r in refs}
            page, remainder = 0, []
            for part in parts:
                if "image" in part:
                    page += 1
                    if page in selected:
                        continue
                remainder.append(part)
            subsets = [([{"image": r["path"]} for r in refs], spec)]
        else:
            text = "".join(p["text"] for p in parts)
            selection = lexical_select(text, questions)
            refs = validate_spans(text, selection["evidence"])
            structure = v3.Structure(text)
            packages, subsets = [], []
            for i, q in enumerate(questions):
                subject = req["entities"][i] if task == "entity" else None
                package = v3.build(text, q, refs, subject=subject, structure=structure)
                unit = copy.deepcopy(spec)
                if "candidates" in spec:
                    candidate = spec["candidates"][i]
                    unit["candidates"] = [candidate]
                    unit["request"]["entities" if task == "entity" else "labels"] = [candidate]
                subsets.append((package["text"], unit))
                packages.append({k: v for k, v in package.items() if k != "text"})
            body["packages"] = packages
            remainder = remove(text, refs)
        body.update(
            references=refs, status="found" if refs else "no_support_found", verification="retrieval_only"
        )
        if level in ("sufficiency", "removal"):
            predictions = [self._fresh(doc, unit) for doc, unit in subsets]
            assembled = (
                {k: v for d in predictions for k, v in d.items()}
                if "candidates" in spec and not images
                else predictions[0]
            )
            body["evidence_only"] = {"prediction": assembled, "agrees_with_full": assembled == baseline}
            body["calls"] += len(subsets)
            body["verification"] = "fresh_source_reencoding"
        if level == "removal":
            removed = self._fresh(remainder, spec)
            body["evidence_removed"] = {"prediction": removed, "agrees_with_full": removed == baseline}
            body["calls"] += 1
        return body