# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""End-to-end: run_arm on 5 synthetic properties with the mock provider (no keys, no network)."""
from __future__ import annotations

import json
from pathlib import Path

import polars as pl
import pytest

from agentic_valuation.config import SYNTHETIC
from agentic_valuation.debate.aggregation import RULES
from agentic_valuation.experiments.freeze import build_freeze, check_freeze, write_freeze
from agentic_valuation.experiments.run_arm import PREDICTION_COLUMNS, RunSpec, prediction_rows, run_arm
from agentic_valuation.schemas import DebateTranscript

pytestmark = pytest.mark.skipif(not (SYNTHETIC / "synthetic_transactions.parquet").exists() and not (SYNTHETIC / "analytic.parquet").exists(), reason="synthetic data not generated")


def test_run_arm_end_to_end(tmp_path: Path, monkeypatch):
    monkeypatch.setenv("AV_PROVIDER_MODE", "mock")
    spec = RunSpec(arm="G_debate", rounds=1, committee="mixed_random", split="validation", n=5, synthetic=True, out_dir=tmp_path, run_id="t_g1")
    run_dir = run_arm(spec)
    preds = pl.read_parquet(run_dir / "predictions.parquet")
    assert preds.columns == list(PREDICTION_COLUMNS)
    assert preds.schema["point"] == pl.Float64 and preds.schema["tokens"] == pl.Int64
    assert preds["tx_id"].n_unique() == 5 and preds.height == 5 * (len(RULES) + 1)  # + mixture row
    assert set(preds["aggregation"].unique()) == set(RULES) | {"mixture"}
    assert set(preds["arm"].unique()) == {"G_debate_r1"} and set(preds["committee"].unique()) == {"mixed_random"} and set(preds["split"].unique()) == {"validation"}
    assert (preds["q05"] <= preds["q50"]).all() and (preds["q50"] <= preds["q95"]).all() and (preds["point"] > 0).all()
    assert len(list((run_dir / "transcripts").glob("*.json"))) == 5 and len(list((run_dir / "packets").glob("*.json"))) == 5
    m = json.loads((run_dir / "manifest.json").read_text())
    assert m["mode"] == "mock" and m["n_transcripts"] == 5 and m["n_failed_this_session"] == 0 and m["arm_label"] == "G_debate_r1"
    assert set(m["prompt_hashes"]) == set(m["roles"]) and len(m["models"]) == 2 and m["config_hashes"] and m["schema_hash"]
    tr = DebateTranscript.model_validate_json(next((run_dir / "transcripts").glob("*.json")).read_text())
    assert tr.rounds <= 1 and len(tr.final_estimates) == 12 and set(tr.assignment.values()) == set(m["models"])
    pk = json.loads(next((run_dir / "packets").glob("*.json")).read_text())
    assert "price" not in pk["subject"] and "street" not in json.dumps(pk).lower()
    # resumable: a second invocation skips everything and rebuilds the same predictions
    run_arm(spec)
    m2 = json.loads((run_dir / "manifest.json").read_text())
    assert m2["n_resumed_skipped"] == 5 and m2["n_transcripts"] == 5
    assert pl.read_parquet(run_dir / "predictions.parquet").height == preds.height


def test_run_arm_spend_cap_graceful(tmp_path: Path, monkeypatch):
    monkeypatch.setenv("AV_PROVIDER_MODE", "mock")
    monkeypatch.setenv("AV_MAX_SPEND_USD", "0.00001")
    spec = RunSpec(arm="F_independent", committee="all_anthropic", split="validation", n=3, synthetic=True, out_dir=tmp_path, run_id="t_cap")
    run_dir = run_arm(spec)
    m = json.loads((run_dir / "manifest.json").read_text())
    assert m["spend_cap_hit"] is True and m["n_transcripts"] == 0 and (run_dir / "predictions.parquet").exists()
    assert pl.read_parquet(run_dir / "predictions.parquet").height == 0


def test_prediction_rows_contract():
    tr = DebateTranscript(tx_id="x", arm="F_independent", committee="c", protocol="independent", rounds=0, packet_hash="h", turns=[], final_estimates={"a": 100.0, "b": 120.0}, final_quantiles={"a": [80, 90, 100, 110, 130], "b": [95, 110, 120, 130, 150]}, aggregated={"mean": 110.0})
    rows = prediction_rows(tr, "test", "rid")
    assert [r["aggregation"] for r in rows] == ["mean", "mixture"] and set(rows[0]) == set(PREDICTION_COLUMNS)


def test_freeze_roundtrip(tmp_path: Path):
    p = write_freeze(tmp_path / "freeze.json")
    doc = json.loads(p.read_text())
    assert len(doc["role_prompt_sha256"]) == 21 + 12 and "models.yaml" in doc["configs_sha256"] and "debate/protocols.py" in doc["code_sha256"]
    assert doc["model_ids"]["cost_matched_pairs"]["mid"]["anthropic"] == build_freeze()["model_ids"]["cost_matched_pairs"]["mid"]["anthropic"]
    assert check_freeze(p) == []