from __future__ import annotations

import asyncio
import json

from gap_pipeline.clients import ScriptedClient
from gap_pipeline.pipeline import KernelPipeline, PipelineConfig
from gap_pipeline.store import RunStore

from conftest import changed_candidate


def accept_review() -> dict[str, str]:
    return {
        "verdict": "accept",
        "step_by_step_check": "every method label is instantiated",
        "blocking_issues": "",
        "patch_suggestion": "",
    }


def reject_review(iteration: int) -> dict[str, str]:
    return {
        "verdict": "reject",
        "step_by_step_check": f"terminal step fails at iteration {iteration}",
        "blocking_issues": "terminal algebra is incorrect",
        "patch_suggestion": "correct the terminal algebra",
    }


def proposer_responses(
    item_id: str,
    plan_dict: dict,
    candidate_dict: dict,
) -> dict:
    return {
        f"{item_id}.plan": plan_dict,
        f"{item_id}.candidate": candidate_dict,
    }


def repair_response(candidate: dict) -> dict[str, str]:
    return {
        "corrected_question": candidate["question"],
        "corrected_solution": candidate["solution"],
        "changes_made": "corrected the terminal algebra",
    }


def test_two_consecutive_unanimous_passes_accept(
    tmp_path,
    item,
    plan_dict,
    candidate_dict,
) -> None:
    proposer = ScriptedClient(
        proposer_responses(item.item_id, plan_dict, candidate_dict)
    )
    judges = [
        ScriptedClient(
            {f"{item.item_id}.verify": [accept_review(), accept_review()]}
        )
        for _ in range(5)
    ]
    pipeline = KernelPipeline(
        proposer=proposer,
        judges=judges,
        store=RunStore(tmp_path, item.item_id),
        config=PipelineConfig(proposer_model="scripted", judge_model="scripted"),
    )
    result = asyncio.run(pipeline.run(item))
    assert result.status == "accepted"
    assert len(result.iterations) == 2
    assert [row.pass_streak_after for row in result.iterations] == [1, 2]

    final_path = tmp_path / "items" / item.item_id / "final.json"
    final = json.loads(final_path.read_text())
    assert final["accepted_candidate_sha256"]
    assert len(list((final_path.parent / "calls").glob("*.json"))) == 12


def test_rejection_resets_streak_and_repairs(
    tmp_path,
    item,
    plan_dict,
    candidate_dict,
) -> None:
    repaired = changed_candidate(candidate_dict, "fixed")
    responses = proposer_responses(item.item_id, plan_dict, candidate_dict)
    responses[f"{item.item_id}.repair"] = repair_response(repaired)
    proposer = ScriptedClient(responses)
    judges = []
    for judge_id in range(1, 6):
        first = reject_review(1) if judge_id == 1 else accept_review()
        judges.append(
            ScriptedClient(
                {
                    f"{item.item_id}.verify": [
                        first,
                        accept_review(),
                        accept_review(),
                    ]
                }
            )
        )
    result = asyncio.run(
        KernelPipeline(
            proposer=proposer,
            judges=judges,
            store=RunStore(tmp_path, item.item_id),
            config=PipelineConfig(
                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].candidate_sha256 != result.iterations[1].candidate_sha256
    assert result.iterations[1].candidate_sha256 == result.iterations[2].candidate_sha256


def test_rejection_after_one_pass_resets_streak(
    tmp_path,
    item,
    plan_dict,
    candidate_dict,
) -> None:
    repaired = changed_candidate(candidate_dict, "after-broken-streak")
    responses = proposer_responses(item.item_id, plan_dict, candidate_dict)
    responses[f"{item.item_id}.repair"] = repair_response(repaired)
    proposer = ScriptedClient(responses)
    judges = []
    for judge_id in range(1, 6):
        second = reject_review(2) if judge_id == 3 else accept_review()
        judges.append(
            ScriptedClient(
                {
                    f"{item.item_id}.verify": [
                        accept_review(),
                        second,
                        accept_review(),
                        accept_review(),
                    ]
                }
            )
        )
    result = asyncio.run(
        KernelPipeline(
            proposer=proposer,
            judges=judges,
            store=RunStore(tmp_path, item.item_id),
            config=PipelineConfig(
                proposer_model="scripted",
                judge_model="scripted",
            ),
        ).run(item)
    )
    assert result.status == "accepted"
    assert [row.pass_streak_after for row in result.iterations] == [1, 0, 1, 2]
    assert result.iterations[0].candidate_sha256 == result.iterations[1].candidate_sha256
    assert result.iterations[1].candidate_sha256 != result.iterations[2].candidate_sha256


def test_fifteen_failed_rounds_reject(
    tmp_path,
    item,
    plan_dict,
    candidate_dict,
) -> None:
    responses = proposer_responses(item.item_id, plan_dict, candidate_dict)
    responses[f"{item.item_id}.repair"] = [
        repair_response(changed_candidate(candidate_dict, f"repair-{iteration}"))
        for iteration in range(1, 15)
    ]
    proposer = ScriptedClient(responses)
    judges = [
        ScriptedClient(
            {
                f"{item.item_id}.verify": [
                    reject_review(iteration) for iteration in range(1, 16)
                ]
            }
        )
        for _ in range(5)
    ]
    result = asyncio.run(
        KernelPipeline(
            proposer=proposer,
            judges=judges,
            store=RunStore(tmp_path, item.item_id),
            config=PipelineConfig(
                proposer_model="scripted",
                judge_model="scripted",
            ),
        ).run(item)
    )
    assert result.status == "rejected"
    assert len(result.iterations) == 15
    assert all(row.pass_streak_after == 0 for row in result.iterations)