"""Tests for native pi session handling in agent stage commands."""

from __future__ import annotations

import os
from types import SimpleNamespace

import pytest
from click.testing import CliRunner

from abreka.cli import main
from abreka.models import Experiment, ExperimentStatus
from abreka.persistence import (
    create_experiment_dir,
    load_experiment,
    reserve_experiment,
    save_experiment,
    update_index_entry,
)


class RecordingPiClient:
    instances: list["RecordingPiClient"] = []
    on_prompt = None

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
        self.prompts: list[str] = []
        type(self).instances.append(self)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False

    def prompt(self, message: str, on_event=None):
        self.prompts.append(message)
        if type(self).on_prompt is not None:
            type(self).on_prompt(message, self.kwargs)
        return SimpleNamespace(text="", events=[])


@pytest.fixture(autouse=True)
def _reset_fake_pi():
    RecordingPiClient.instances = []
    RecordingPiClient.on_prompt = None
    yield
    RecordingPiClient.instances = []
    RecordingPiClient.on_prompt = None


def _seed_experiment(abreka_project, exp_id: str, status: ExperimentStatus) -> None:
    abreka_dir = abreka_project / "abreka"
    exp = Experiment(id=exp_id, status=status, hypothesis=f"hypothesis {exp_id}")
    create_experiment_dir(abreka_dir, exp_id)
    save_experiment(exp, abreka_dir)
    update_index_entry(abreka_dir, exp)


class TestProposeSessions:
    def test_propose_reserves_id_and_uses_native_session_file(self, abreka_project, monkeypatch):
        abreka_dir = abreka_project / "abreka"
        (abreka_dir / "prompts" / "propose_exploit.md").write_text(
            "Reserved {{reserved_exp_id}}\n{{goal}}\n{{guidance}}\n"
        )

        def finalize_reserved_experiment(_message, _kwargs):
            exp = load_experiment(abreka_dir, "0001")
            exp.hypothesis = "reserved"
            exp.method = "method"
            exp.status = ExperimentStatus.PROPOSED
            save_experiment(exp, abreka_dir)
            update_index_entry(abreka_dir, exp)

        RecordingPiClient.on_prompt = finalize_reserved_experiment
        monkeypatch.setattr("abreka.cli.agent.propose.PiRpcClient", RecordingPiClient)

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["propose"])

        assert result.exit_code == 0, result.output
        exp = load_experiment(abreka_dir, "0001")
        assert exp.status == ExperimentStatus.PROPOSED
        assert exp.runtime_costs["propose"]["attempts"] == 1
        assert exp.runtime_costs["propose"]["last_status"] == "ok"

        client = RecordingPiClient.instances[-1]
        assert client.kwargs["session_path"] == (
            abreka_dir / "experiments" / "0001" / "artifacts" / "pi_session_propose.jsonl"
        )
        assert "Reserved 0001" in client.prompts[0]

    def test_propose_resume_reuses_reserved_experiment_session(self, abreka_project, monkeypatch):
        abreka_dir = abreka_project / "abreka"
        reserved = reserve_experiment(abreka_dir)
        session_path = abreka_dir / "experiments" / reserved.id / "artifacts" / "pi_session_propose.jsonl"
        session_path.write_text('{"type":"session","version":3}\n')

        def finalize_reserved_experiment(_message, _kwargs):
            exp = load_experiment(abreka_dir, reserved.id)
            exp.hypothesis = "done"
            exp.method = "method"
            exp.status = ExperimentStatus.PROPOSED
            save_experiment(exp, abreka_dir)
            update_index_entry(abreka_dir, exp)

        RecordingPiClient.on_prompt = finalize_reserved_experiment
        monkeypatch.setattr("abreka.cli.agent.propose.PiRpcClient", RecordingPiClient)

        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["propose", "--resume", "1"])

        assert result.exit_code == 0, result.output
        client = RecordingPiClient.instances[-1]
        assert client.kwargs["session_path"] == session_path
        assert client.prompts[0].startswith("Resume the interrupted propose step for experiment 0001")

        exp = load_experiment(abreka_dir, "0001")
        assert exp.runtime_costs["propose"]["attempts"] == 1
        assert exp.runtime_costs["propose"]["last_status"] == "ok"


@pytest.mark.parametrize(
    ("command", "patch_target", "fresh_status", "resume_status", "expected_status", "filename", "resume_prefix"),
    [
        (
            "implement",
            "abreka.cli.agent.implement.PiRpcClient",
            ExperimentStatus.PROPOSED,
            ExperimentStatus.IMPLEMENTING,
            ExperimentStatus.IMPLEMENTING,
            "pi_session_implement.jsonl",
            "Resume the interrupted implement step for experiment 0001",
        ),
        (
            "test",
            "abreka.cli.agent.test_cmd.PiRpcClient",
            ExperimentStatus.IMPLEMENTING,
            ExperimentStatus.TESTING,
            ExperimentStatus.TESTING,
            "pi_session_test.jsonl",
            "Resume the interrupted test step for experiment 0001",
        ),
        (
            "evaluate",
            "abreka.cli.agent.evaluate.PiRpcClient",
            ExperimentStatus.RUNNING,
            ExperimentStatus.RUNNING,
            ExperimentStatus.RUNNING,
            "pi_session_evaluate.jsonl",
            "Resume the interrupted evaluate step for experiment 0001",
        ),
    ],
)
def test_stage_command_uses_native_session_file_on_fresh_run(
    abreka_project,
    monkeypatch,
    command,
    patch_target,
    fresh_status,
    resume_status,
    expected_status,
    filename,
    resume_prefix,
):
    _seed_experiment(abreka_project, "0001", fresh_status)
    if command == "test":
        (abreka_project / "abreka" / "prompts" / "test.md").write_text("# test\n{{test_command}}\n")
    monkeypatch.setattr(patch_target, RecordingPiClient)

    runner = CliRunner()
    os.chdir(abreka_project)
    result = runner.invoke(main, [command, "1"])

    assert result.exit_code == 0, result.output
    exp = load_experiment(abreka_project / "abreka", "0001")
    assert exp.status == expected_status
    assert exp.runtime_costs[command]["attempts"] == 1
    assert exp.runtime_costs[command]["last_status"] == "ok"
    assert RecordingPiClient.instances[-1].kwargs["session_path"] == (
        abreka_project / "abreka" / "experiments" / "0001" / "artifacts" / filename
    )


@pytest.mark.parametrize(
    ("command", "patch_target", "status", "filename", "resume_prefix"),
    [
        (
            "implement",
            "abreka.cli.agent.implement.PiRpcClient",
            ExperimentStatus.IMPLEMENTING,
            "pi_session_implement.jsonl",
            "Resume the interrupted implement step for experiment 0001",
        ),
        (
            "test",
            "abreka.cli.agent.test_cmd.PiRpcClient",
            ExperimentStatus.TESTING,
            "pi_session_test.jsonl",
            "Resume the interrupted test step for experiment 0001",
        ),
        (
            "evaluate",
            "abreka.cli.agent.evaluate.PiRpcClient",
            ExperimentStatus.RUNNING,
            "pi_session_evaluate.jsonl",
            "Resume the interrupted evaluate step for experiment 0001",
        ),
    ],
)
def test_stage_command_resume_uses_existing_session(
    abreka_project,
    monkeypatch,
    command,
    patch_target,
    status,
    filename,
    resume_prefix,
):
    _seed_experiment(abreka_project, "0001", status)
    if command == "test":
        (abreka_project / "abreka" / "prompts" / "test.md").write_text("# test\n{{test_command}}\n")
    session_path = abreka_project / "abreka" / "experiments" / "0001" / "artifacts" / filename
    session_path.write_text('{"type":"session","version":3}\n')
    monkeypatch.setattr(patch_target, RecordingPiClient)

    runner = CliRunner()
    os.chdir(abreka_project)
    result = runner.invoke(main, [command, "1", "--resume"])

    assert result.exit_code == 0, result.output
    client = RecordingPiClient.instances[-1]
    assert client.kwargs["session_path"] == session_path
    assert client.prompts[0].startswith(resume_prefix)

    exp = load_experiment(abreka_project / "abreka", "0001")
    assert exp.runtime_costs[command]["attempts"] == 1
    assert exp.runtime_costs[command]["last_status"] == "ok"


def test_implement_refuses_to_reuse_existing_session_without_resume(abreka_project):
    _seed_experiment(abreka_project, "0001", ExperimentStatus.PROPOSED)
    session_path = (
        abreka_project
        / "abreka"
        / "experiments"
        / "0001"
        / "artifacts"
        / "pi_session_implement.jsonl"
    )
    session_path.write_text('{"type":"session","version":3}\n')

    runner = CliRunner()
    os.chdir(abreka_project)
    result = runner.invoke(main, ["implement", "1"])

    assert result.exit_code != 0
    assert "Use --resume" in result.output
