"""Focused public API checks using the real BF16 model, without calibration."""

import argparse
import json
from pathlib import Path

import numpy as np

from solomon_mlx import Solomon


def validate(output):
    output = Path(output)
    if output.exists():
        raise FileExistsError("Validation outputs are immutable")
    model = Solomon.load("models/quality")
    document = "Alice is certified. Bob is not certified. The current priority is high."
    questions = {
        "boolean": "Is Alice certified?",
        "single": {"type": "choice", "instructions": "Who is certified?", "options": ["Alice", "Bob"]},
        "ordered": {"type": "score", "instructions": "What is the priority?", "levels": ["low", "high"]},
        "entity": {"instructions": "Is {candidate} certified?", "candidates": ["Alice", "Bob"]},
        "multilabel": {
            "instructions": "Which facts apply?",
            "candidates": ["Alice is certified", "Bob is certified"],
        },
    }
    checks = {}
    with model.prefill(document) as state:
        first = model.decide(state=state, questions=questions, evidence="none", diagnostics=True)
        assert set(first["answers"]) == set(questions)
        checks["all_five_answer_types"] = True
        repeated = model.decide(
            state=state, questions={"boolean": questions["boolean"]}, evidence="none", diagnostics=True
        )
        a = first["answers"]["boolean"]["branches"][0]["letter_logits"]
        b = repeated["answers"]["boolean"]["branches"][0]["letter_logits"]
        np.testing.assert_array_equal(a, b)
        checks["repeated_question_logits_exact"] = True
        full = model.decide(
            state=state,
            questions={"boolean": questions["boolean"]},
            evidence="none",
            execution="full",
            diagnostics=True,
        )
        c = full["answers"]["boolean"]["branches"][0]["letter_logits"]
        checks["cached_full_max_logit_drift"] = float(np.max(np.abs(np.array(a) - c)))
        assert (first["answers"]["boolean"]["noul"] >= 0.5) == (full["answers"]["boolean"]["noul"] >= 0.5)
        reverse = model.decide(
            state=state,
            questions={"entity": {**questions["entity"], "candidates": ["Bob", "Alice"]}},
            evidence="none",
        )
        assert reverse["answers"]["entity"]["candidates"] == first["answers"]["entity"]["candidates"]
        assert list(reverse["answers"]["entity"]["candidates"]) == ["Bob", "Alice"]
        checks["candidate_order_and_cache_isolation"] = True
        evidence = model.decide(state=state, questions={"boolean": questions["boolean"]}, evidence="removal")
        answer = evidence["answers"]["boolean"]
        assert answer["evidence"]
        for span in answer["evidence"]:
            assert document[span["start"] : span["end"]] == span["text"]
        assert answer["evidence_detail"]["verification"] == "fresh_source_reencoding"
        assert answer["evidence_detail"]["calls"] == 2
        checks["evidence_spans_and_fresh_verification"] = True
        exhausted = model.decide(
            state=state, questions={"boolean": questions["boolean"]}, evidence="removal", evidence_max_calls=0
        )
        assert exhausted["answers"]["boolean"]["evidence_status"] == "budget_exhausted"
        assert exhausted["usage"]["evidence_calls"] == 0
        checks["evidence_budget_enforced"] = True
        replay = output.with_suffix(".replay.json")
        state.save(replay)
    try:
        model.decide(state=state, questions={"q": "Fact?"})
    except ValueError:
        checks["closed_state_rejected"] = True
    else:
        raise AssertionError("Closed state accepted")
    with model.replay(replay) as restored:
        result = model.decide(
            state=restored, questions={"boolean": questions["boolean"]}, evidence="none", diagnostics=True
        )
        np.testing.assert_array_equal(a, result["answers"]["boolean"]["branches"][0]["letter_logits"])
        checks["public_api_replay_exact"] = True
    corrupt = json.loads(replay.read_text())
    corrupt["parts"][0]["text"] += " changed"
    bad_path = output.with_suffix(".corrupt-replay.json")
    bad_path.write_text(json.dumps(corrupt))
    try:
        model.replay(bad_path)
    except ValueError:
        checks["corrupt_replay_rejected"] = True
    else:
        raise AssertionError("Corrupt replay accepted")
    image_parts = json.loads(Path("evaluations/image-jobs.json").read_text())[0]["parts"]
    with model.prefill(image_parts) as images:
        result = model.decide(state=images, questions={"q": "Is Alice certified?"}, evidence="support")
        assert result["answers"]["q"]["evidence_status"] == "unsupported_page_selector"
        checks["missing_page_selector_reported"] = True
    with model.prefill({"subject": "Alice", "certified": True}) as structured:
        assert structured.prefix_tokens > 0
        checks["structured_document_accepted"] = True
    try:
        model.engine.admit(40961)
    except ValueError:
        checks["context_ceiling_enforced"] = True
    else:
        raise AssertionError("Context limit not enforced")
    assert model.engine.context["start"] is None
    assert len(model.engine.heads) == 10
    checks["adapter_state_reset_and_ten_heads_loaded"] = True
    report = {
        "runtime": model.identity,
        "checks": checks,
        "passed": True,
        "scope": "real-weight API behavior; these checks do not establish held-out CUDA parity",
        "answers": first["answers"],
        "evidence": answer,
    }
    output.write_text(json.dumps(report, indent=2))
    print(json.dumps({"passed": True, "checks": checks}, indent=2))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", required=True)
    validate(parser.parse_args().output)