from __future__ import annotations

import asyncio
import json

from gap_pipeline.clients import ScriptedClient
from gap_pipeline.kernel_models import MethodPlan, ProofDAG, ReplacementPlan
from gap_pipeline.kernel_prompts import diffusion_user
from gap_pipeline.paper_pipeline import PaperKernelPipeline, PaperPipelineConfig
from gap_pipeline.prompts import JUDGE_SYSTEM_PROMPT
from gap_pipeline.store import RunStore


def _dag() -> dict:
    return {
        "nodes": [
            {
                "node_id": "n1",
                "claim": "(a-1)^2 >= 0",
                "dependencies": [],
            },
            {
                "node_id": "n2",
                "claim": "a+1/a >= 2",
                "dependencies": ["n1"],
            },
        ],
        "terminal_node_id": "n2",
    }


def _methods() -> dict:
    return {
        "nodes": [
            {
                "node_id": "n1",
                "method_label": "use nonnegativity of a square",
            },
            {
                "node_id": "n2",
                "method_label": "expand and divide by a positive quantity",
            },
        ]
    }


def _replacement() -> dict:
    return {
        "changes": [
            {
                "slot_id": "slot1",
                "source_node_id": "n1",
                "description": "positive equality value",
                "original_value": "1",
                "replacement_value": "2",
                "guard_condition": "replacement is positive",
                "guard_justification": "2 is positive",
            }
        ],
        "closure_statement": "No undeclared changes.",
    }


def _diffused() -> dict:
    return {
        "nodes": [
            {
                "node_id": "n1",
                "dependencies": [],
                "method_label": "use nonnegativity of a square",
                "instantiated_claim": "(x-2)^2 >= 0",
                "justification": "squares are nonnegative",
            },
            {
                "node_id": "n2",
                "dependencies": ["n1"],
                "method_label": "expand and divide by a positive quantity",
                "instantiated_claim": "x+4/x >= 4",
                "justification": "expand and divide by x>0",
            },
        ],
        "terminal_node_id": "n2",
        "terminal_answer": "4",
    }


def _variant(suffix: str = "") -> dict:
    return {
        "question": f"Let x>0. Prove that x+4/x >= 4.{suffix}",
        "solution": (
            "[n1] Since (x-2)^2 >= 0. "
            "[n2] Expand and divide by x>0 to obtain x+4/x >= 4."
        ),
        "node_order": ["n1", "n2"],
        "terminal_answer": "4",
    }


def _accept() -> dict:
    return {
        "verdict": "accept",
        "step_by_step_check": "n1 is valid; n2 is valid",
        "blocking_issues": "",
        "patch_suggestion": "",
    }


def _reject() -> dict:
    return {
        "verdict": "reject",
        "step_by_step_check": "n1 is valid; n2 needs a wording correction",
        "blocking_issues": "the terminal wording is ambiguous",
        "patch_suggestion": "clarify the terminal wording",
    }


def test_five_stage_repair_uses_prior_bundle_and_appendix_judge(
    tmp_path,
    item,
) -> None:
    proposer = ScriptedClient(
        {
            f"{item.item_id}.stage1.dag": _dag(),
            f"{item.item_id}.stage2.methods": _methods(),
            f"{item.item_id}.stage3.replacement": [
                _replacement(),
                _replacement(),
            ],
            f"{item.item_id}.stage4.diffusion": [
                _diffused(),
                _diffused(),
            ],
            f"{item.item_id}.stage5.render": [
                _variant(),
                _variant(" The requested bound is explicit."),
            ],
        }
    )
    judges = []
    for judge_id in range(1, 6):
        first = _reject() if judge_id == 1 else _accept()
        judges.append(
            ScriptedClient(
                {f"{item.item_id}.verify": [first, _accept(), _accept()]}
            )
        )

    run_root = tmp_path / "run"
    result = asyncio.run(
        PaperKernelPipeline(
            proposer=proposer,
            judges=judges,
            store=RunStore(run_root, item.item_id),
            config=PaperPipelineConfig(
                proposer_model="scripted",
                judge_model="scripted",
            ),
        ).run(item)
    )

    assert result.status == "accepted"
    assert [row.pass_streak_after for row in result.iterations] == [0, 1, 2]
    assert result.iterations[0].bundle_sha256 != result.iterations[1].bundle_sha256
    assert result.iterations[1].bundle_sha256 == result.iterations[2].bundle_sha256

    calls_dir = run_root / "items" / item.item_id / "calls"
    repair_stage_calls = [
        json.loads(
            (
                calls_dir
                / f"{item.item_id}.stage{stage}.{name}.v02.json"
            ).read_text()
        )
        for stage, name in [
            (3, "replacement"),
            (4, "diffusion"),
            (5, "render"),
        ]
    ]
    assert "PREVIOUS REPLACEMENT PLAN" in repair_stage_calls[0]["user_prompt"]
    assert "PREVIOUS DIFFUSED PROOF" in repair_stage_calls[1]["user_prompt"]
    assert "PREVIOUS RENDERED VARIANT" in repair_stage_calls[2]["user_prompt"]
    assert all(
        "the terminal wording is ambiguous" in call["user_prompt"]
        for call in repair_stage_calls
    )

    judge_calls = sorted(calls_dir.glob(f"{item.item_id}.verify.*.json"))
    assert len(judge_calls) == 15
    assert not any(".a1." in path.name for path in judge_calls)
    first_judge_call = json.loads(judge_calls[0].read_text())
    assert first_judge_call["system_prompt"] == JUDGE_SYSTEM_PROMPT
    assert "METHOD-LABEL SEQUENCE (abstract plan):" in first_judge_call["user_prompt"]
    assert "SOURCE PROOF DAG:" not in first_judge_call["user_prompt"]


def test_stage4_retries_identical_prompt_when_required_field_is_missing(
    tmp_path,
    item,
) -> None:
    invalid = _diffused()
    for node in invalid["nodes"]:
        node.pop("justification")

    request_prefix = f"{item.item_id}.stage4.diffusion.v01"
    proposer = ScriptedClient({request_prefix: [invalid, _diffused()]})
    pipeline = PaperKernelPipeline(
        proposer=proposer,
        judges=[ScriptedClient({}) for _ in range(5)],
        store=RunStore(tmp_path / "run", item.item_id),
        config=PaperPipelineConfig(
            proposer_model="scripted",
            judge_model="scripted",
        ),
    )
    dag = ProofDAG.model_validate(_dag())
    methods = MethodPlan.model_validate(_methods())
    replacements = ReplacementPlan.model_validate(_replacement())

    result = asyncio.run(
        pipeline.diffuse_dag(
            item,
            dag,
            methods,
            replacements,
            version=1,
        )
    )

    assert all(node.justification for node in result.nodes)
    assert proposer.calls == [
        request_prefix,
        f"{request_prefix}.validation_retry01",
    ]

    calls_dir = tmp_path / "run" / "items" / item.item_id / "calls"
    first_call = json.loads((calls_dir / f"{request_prefix}.json").read_text())
    retry_call = json.loads(
        (
            calls_dir / f"{request_prefix}.validation_retry01.json"
        ).read_text()
    )
    expected_prompt = diffusion_user(item, dag, methods, replacements)
    assert first_call["user_prompt"] == expected_prompt
    assert retry_call["user_prompt"] == expected_prompt

    stage = json.loads(
        (
            tmp_path
            / "run"
            / "items"
            / item.item_id
            / "stages"
            / "04_diffused_proof_v01.json"
        ).read_text()
    )
    assert stage["request_id"] == f"{request_prefix}.validation_retry01"