"""Tests for abreka.runner."""

import json
import warnings
from pathlib import Path
from types import SimpleNamespace

import pytest
from rich.console import Console

from abreka.config import load_config, save_config
from abreka.models import Experiment, ExperimentStatus
from abreka.persistence import create_experiment_dir, load_experiment, save_experiment
from abreka.prompts import get_default_template, render_template
from abreka.run_executor import RunExecutionResult
from abreka.runner import AutonomousRunner


class TestShouldExplore:
    def _make_runner(self, explore_probability=None):
        return AutonomousRunner(
            abreka_dir=Path("/fake"),
            explore_probability=explore_probability,
        )

    def test_none_never_explores(self):
        runner = self._make_runner(explore_probability=None)
        assert runner._should_explore() is False

    def test_zero_never_explores(self):
        runner = self._make_runner(explore_probability=0.0)
        assert runner._should_explore() is False

    def test_one_always_explores(self):
        runner = self._make_runner(explore_probability=1.0)
        assert runner._should_explore() is True

    def test_stochastic_respects_probability(self):
        runner = self._make_runner(explore_probability=0.5)
        results = [runner._should_explore() for _ in range(1000)]
        ratio = sum(results) / len(results)
        assert 0.35 < ratio < 0.65


class TestFocusRoot:
    def test_focus_root_default_none(self):
        runner = AutonomousRunner(abreka_dir=Path("/fake"))
        assert runner.focus_root is None

    def test_focus_root_stored(self):
        runner = AutonomousRunner(abreka_dir=Path("/fake"), focus_root="0004")
        assert runner.focus_root == "0004"


def _seed_completed_experiment(abreka_dir: Path, exp_id: str, parents: list[str]) -> None:
    """Add a completed experiment with metrics to the project on disk + index."""
    exp_path = abreka_dir / "experiments" / exp_id
    (exp_path / "artifacts").mkdir(parents=True)
    import tomli_w

    toml_data = {
        "id": exp_id,
        "status": "completed",
        "created_at": "2026-04-15T00:00:00+00:00",
        "hypothesis": f"hypothesis for {exp_id}",
        "method": f"method for {exp_id}",
        "parents": parents,
        "tags": [],
        "metrics": {"val": {"accuracy": 0.9 + int(exp_id) * 0.001}},
    }
    with open(exp_path / "experiment.toml", "wb") as f:
        tomli_w.dump(toml_data, f)

    index_path = abreka_dir / "index.json"
    index = json.loads(index_path.read_text())
    index[exp_id] = toml_data
    index_path.write_text(json.dumps(index, indent=2))


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

        recorded: dict[str, object] = {}

        class FakePi:
            def __init__(self, *args, **kwargs):
                recorded["session_path"] = kwargs.get("session_path")

            def __enter__(self):
                return self

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

            def prompt(self, message, on_event=None):
                recorded["message"] = message
                exp = load_experiment(abreka_dir, "0001")
                exp.hypothesis = "reserved"
                exp.method = "method"
                exp.status = ExperimentStatus.PROPOSED
                save_experiment(exp, abreka_dir)
                from abreka.persistence import update_index_entry
                update_index_entry(abreka_dir, exp)
                return SimpleNamespace(text="", events=[])

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        monkeypatch.setattr(runner, "_get_pi_client", lambda model_override="", session_path=None: FakePi(session_path=session_path))

        exp_id = runner._propose()

        assert exp_id == "0001"
        assert recorded["session_path"] == (
            abreka_dir / "experiments" / "0001" / "artifacts" / "pi_session_propose.jsonl"
        )
        assert "Reserved 0001" in recorded["message"]

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


class TestRunExecution:
    def _seed_runnable_experiment(self, abreka_dir: Path) -> Experiment:
        exp = Experiment(id="0001", status=ExperimentStatus.TESTING, hypothesis="Test run")
        create_experiment_dir(abreka_dir, exp.id)
        save_experiment(exp, abreka_dir)
        (abreka_dir / "experiments" / exp.id / "run.py").write_text("print('hello')\n")
        return exp

    def test_run_experiment_marks_failed_when_executor_raises(self, abreka_project: Path, monkeypatch):
        abreka_dir = abreka_project / "abreka"
        exp = self._seed_runnable_experiment(abreka_dir)

        class FailingExecutor:
            def execute(self, *, run_py, output_path, timeout_seconds):
                raise RuntimeError("rsync failed")

        monkeypatch.setattr(
            "abreka.runner.build_run_executor",
            lambda abreka_dir, config, console: FailingExecutor(),
        )

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        with pytest.raises(RuntimeError, match="rsync failed"):
            runner._run_experiment(exp.id)

        reloaded = load_experiment(abreka_dir, exp.id)
        assert reloaded.status == ExperimentStatus.FAILED
        assert reloaded.failure_reason == "rsync failed"

    def test_apply_run_result_reads_metrics_from_output_file_when_stdout_empty(self, abreka_project: Path):
        abreka_dir = abreka_project / "abreka"
        exp = Experiment(id="0001", status=ExperimentStatus.RUNNING, hypothesis="Test run")
        create_experiment_dir(abreka_dir, exp.id)
        save_experiment(exp, abreka_dir)

        output_path = abreka_dir / "experiments" / exp.id / "artifacts" / "run_output.txt"
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(
            "line\n"
            "METRIC:val:accuracy=0.91\n"
            "METRIC:train:loss=0.04\n"
        )

        result = RunExecutionResult(stdout="", stderr="", returncode=0, timed_out=False)
        runner = AutonomousRunner(abreka_dir=abreka_dir)
        metrics = runner._apply_run_result(
            exp.id,
            result,
            timeout_minutes=30,
            output_path=output_path,
        )

        assert metrics == {"val": {"accuracy": 0.91}, "train": {"loss": 0.04}}
        reloaded = load_experiment(abreka_dir, exp.id)
        assert reloaded.metrics == metrics

    def test_run_experiment_records_runtime_costs(self, abreka_project: Path, monkeypatch):
        abreka_dir = abreka_project / "abreka"
        exp = self._seed_runnable_experiment(abreka_dir)

        class FakeExecutor:
            def execute(self, *, run_py, output_path, timeout_seconds):
                output_path.write_text("METRIC:val:accuracy=0.99\n")
                return RunExecutionResult(stdout="", stderr="", returncode=0, timed_out=False)

        monkeypatch.setattr(
            "abreka.runner.build_run_executor",
            lambda abreka_dir, config, console: FakeExecutor(),
        )

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        runner._run_experiment(exp.id)

        reloaded = load_experiment(abreka_dir, exp.id)
        assert reloaded.runtime_costs["run"]["attempts"] == 1
        assert reloaded.runtime_costs["run"]["last_status"] == "ok"
        assert reloaded.runtime_costs["run"]["seconds_total"] >= 0.0

    def test_run_experiment_repairs_non_timeout_failure_and_reruns(
        self, abreka_project: Path, monkeypatch
    ):
        abreka_dir = abreka_project / "abreka"
        exp = self._seed_runnable_experiment(abreka_dir)
        (abreka_dir / "prompts" / "repair_run.md").write_text(get_default_template("repair_run"))

        class FlakyExecutor:
            calls = 0

            def execute(self, *, run_py, output_path, timeout_seconds):
                type(self).calls += 1
                if type(self).calls == 1:
                    output_path.write_text("Traceback\nValueError: bad mask\n")
                    return RunExecutionResult(stdout="", stderr="", returncode=1, timed_out=False)
                output_path.write_text("METRIC:val:accuracy=0.87\n")
                return RunExecutionResult(stdout="", stderr="", returncode=0, timed_out=False)

        monkeypatch.setattr(
            "abreka.runner.build_run_executor",
            lambda abreka_dir, config, console: FlakyExecutor(),
        )

        repair_sessions: list[Path] = []
        repair_prompts: list[str] = []

        class FakePi:
            def __init__(self, session_path):
                self.session_path = session_path

            def __enter__(self):
                return self

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

            def prompt(self, message, on_event=None):
                repair_sessions.append(self.session_path)
                repair_prompts.append(message)
                return SimpleNamespace(text="", events=[])

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        monkeypatch.setattr(
            runner,
            "_get_pi_client",
            lambda model_override="", session_path=None: FakePi(session_path),
        )

        runner._run_experiment(exp.id)

        reloaded = load_experiment(abreka_dir, exp.id)
        assert FlakyExecutor.calls == 2
        assert reloaded.status == ExperimentStatus.RUNNING
        assert reloaded.metrics == {"val": {"accuracy": 0.87}}
        assert len(reloaded.failures) == 1
        failure = reloaded.failures[0]
        assert failure["stage"] == "run"
        assert failure["attempt"] == 1
        assert failure["repair_attempt"] == 0
        assert failure["returncode"] == 1
        assert failure["timed_out"] is False
        assert "run.py exited with code 1" in failure["reason"]
        snapshot = abreka_dir / "experiments" / exp.id / failure["output_path"]
        assert snapshot.read_text() == "Traceback\nValueError: bad mask\n"
        assert repair_sessions == [
            abreka_dir / "experiments" / exp.id / "artifacts" / "pi_session_run_repair_1.jsonl"
        ]
        assert "ValueError: bad mask" in repair_prompts[0]
        assert reloaded.runtime_costs["run"]["attempts"] == 2
        assert reloaded.runtime_costs["run_repair"]["attempts"] == 1

    def test_run_experiment_exhausts_repair_attempts_then_fails(
        self, abreka_project: Path, monkeypatch
    ):
        abreka_dir = abreka_project / "abreka"
        exp = self._seed_runnable_experiment(abreka_dir)
        config = load_config(abreka_dir)
        config.run.max_repair_attempts = 2
        save_config(config, abreka_dir)

        class AlwaysFailingExecutor:
            calls = 0

            def execute(self, *, run_py, output_path, timeout_seconds):
                type(self).calls += 1
                output_path.write_text(f"Traceback attempt {type(self).calls}\n")
                return RunExecutionResult(stdout="", stderr="", returncode=1, timed_out=False)

        monkeypatch.setattr(
            "abreka.runner.build_run_executor",
            lambda abreka_dir, config, console: AlwaysFailingExecutor(),
        )

        class FakePi:
            calls = 0

            def __init__(self, session_path):
                self.session_path = session_path

            def __enter__(self):
                return self

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

            def prompt(self, message, on_event=None):
                type(self).calls += 1
                return SimpleNamespace(text="", events=[])

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        monkeypatch.setattr(
            runner,
            "_get_pi_client",
            lambda model_override="", session_path=None: FakePi(session_path),
        )

        with pytest.raises(RuntimeError, match="after 2 repair attempts"):
            runner._run_experiment(exp.id)

        reloaded = load_experiment(abreka_dir, exp.id)
        assert AlwaysFailingExecutor.calls == 3
        assert FakePi.calls == 2
        assert reloaded.status == ExperimentStatus.FAILED
        assert "after 2 repair attempts" in reloaded.failure_reason
        assert [f["attempt"] for f in reloaded.failures] == [1, 2, 3]
        assert reloaded.runtime_costs["run"]["attempts"] == 3
        assert reloaded.runtime_costs["run"]["last_status"] == "error"
        assert reloaded.runtime_costs["run_repair"]["attempts"] == 2

    def test_run_timeout_fails_without_repair(self, abreka_project: Path, monkeypatch):
        abreka_dir = abreka_project / "abreka"
        exp = self._seed_runnable_experiment(abreka_dir)

        class TimeoutExecutor:
            calls = 0

            def execute(self, *, run_py, output_path, timeout_seconds):
                type(self).calls += 1
                output_path.write_text("still running\n")
                return RunExecutionResult(stdout="", stderr="", returncode=143, timed_out=True)

        monkeypatch.setattr(
            "abreka.runner.build_run_executor",
            lambda abreka_dir, config, console: TimeoutExecutor(),
        )

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        monkeypatch.setattr(
            runner,
            "_get_pi_client",
            lambda *args, **kwargs: pytest.fail("timeouts should not invoke repair"),
        )

        with pytest.raises(RuntimeError, match="Timed out after 30m"):
            runner._run_experiment(exp.id)

        reloaded = load_experiment(abreka_dir, exp.id)
        assert TimeoutExecutor.calls == 1
        assert reloaded.status == ExperimentStatus.FAILED
        assert reloaded.runtime_costs["run"]["attempts"] == 1
        assert reloaded.runtime_costs["run"]["last_status"] == "timed_out"
        assert len(reloaded.failures) == 1
        assert reloaded.failures[0]["timed_out"] is True
        assert "run_repair" not in reloaded.runtime_costs


class TestGracefulStopAfterCurrent:
    def test_run_honors_existing_stop_request_before_new_iteration(
        self, abreka_project: Path, monkeypatch
    ):
        abreka_dir = abreka_project / "abreka"
        stop_path = abreka_dir / ".stop_after_current"
        stop_path.write_text("queued\n")

        console = Console(record=True, width=120)
        runner = AutonomousRunner(
            abreka_dir=abreka_dir,
            duration_seconds=3600,
            console=console,
        )

        calls = {"propose": 0}

        def fake_propose() -> str:
            calls["propose"] += 1
            raise AssertionError("_propose should not be called when stop is already queued")

        monkeypatch.setattr(runner, "_propose", fake_propose)

        runner.run()

        assert calls["propose"] == 0
        assert not stop_path.exists()
        assert "Stop requested" in console.export_text()

    def test_run_clears_stop_request_queued_during_current_experiment(
        self, abreka_project: Path, monkeypatch
    ):
        abreka_dir = abreka_project / "abreka"
        stop_path = abreka_dir / ".stop_after_current"

        console = Console(record=True, width=120)
        runner = AutonomousRunner(
            abreka_dir=abreka_dir,
            duration_seconds=None,
            console=console,
        )

        calls = {name: 0 for name in ("propose", "implement", "test", "run", "evaluate")}

        def fake_propose() -> str:
            calls["propose"] += 1
            return "0001"

        def fake_implement(exp_id: str) -> None:
            assert exp_id == "0001"
            calls["implement"] += 1

        def fake_test(exp_id: str) -> None:
            assert exp_id == "0001"
            calls["test"] += 1

        def fake_run_experiment(exp_id: str) -> None:
            assert exp_id == "0001"
            calls["run"] += 1

        def fake_evaluate(exp_id: str) -> None:
            assert exp_id == "0001"
            calls["evaluate"] += 1
            stop_path.write_text("queued\n")

        monkeypatch.setattr(runner, "_propose", fake_propose)
        monkeypatch.setattr(runner, "_implement", fake_implement)
        monkeypatch.setattr(runner, "_test", fake_test)
        monkeypatch.setattr(runner, "_run_experiment", fake_run_experiment)
        monkeypatch.setattr(runner, "_evaluate", fake_evaluate)

        runner.run()

        assert calls == {
            "propose": 1,
            "implement": 1,
            "test": 1,
            "run": 1,
            "evaluate": 1,
        }
        assert not stop_path.exists()
        assert "Stop requested" in console.export_text()


class TestPromptVariablesCoverage:
    """Ensure the runner passes every placeholder referenced by default prompts.

    The regression this guards against: new template variables were added to the
    default propose/evaluate prompts (branch_overview, focused_branch_summary,
    parent_context, branch_summary, global_context) but only the standalone CLI
    path was updated to supply them. The runner path silently rendered them as
    literal ``{{placeholder}}`` text in production, with only a UserWarning.
    """

    def _install_default_prompts(self, abreka_dir: Path) -> None:
        for name in ("propose_exploit", "propose_explore", "evaluate", "repair_run"):
            (abreka_dir / "prompts" / f"{name}.md").write_text(
                get_default_template(name)
            )

    def _assert_no_missing_vars(self, template: str, variables: dict) -> None:
        with warnings.catch_warnings(record=True) as captured:
            warnings.simplefilter("always")
            rendered = render_template(template, variables)
        missing = [
            str(w.message)
            for w in captured
            if "Template variable" in str(w.message)
        ]
        assert not missing, f"Missing template variables: {missing}"
        assert "{{" not in rendered, (
            f"Unresolved placeholders in rendered template: "
            f"{[line for line in rendered.splitlines() if '{{' in line]}"
        )

    @pytest.mark.parametrize("template_name", ["propose_exploit", "propose_explore"])
    def test_propose_prompt_vars_cover_default_template(
        self, abreka_project: Path, template_name: str
    ):
        abreka_dir = abreka_project / "abreka"
        self._install_default_prompts(abreka_dir)
        _seed_completed_experiment(abreka_dir, "0001", parents=[])
        _seed_completed_experiment(abreka_dir, "0002", parents=["0001"])

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        variables = runner._build_propose_prompt_vars(template_name)

        template = get_default_template(template_name)
        self._assert_no_missing_vars(template, variables)

    def test_evaluate_prompt_vars_cover_default_template(self, abreka_project: Path):
        abreka_dir = abreka_project / "abreka"
        self._install_default_prompts(abreka_dir)
        _seed_completed_experiment(abreka_dir, "0001", parents=[])
        _seed_completed_experiment(abreka_dir, "0002", parents=["0001"])

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        variables = runner._build_evaluate_prompt_vars("0002")

        template = get_default_template("evaluate")
        self._assert_no_missing_vars(template, variables)

    def test_repair_run_prompt_vars_cover_default_template(self, abreka_project: Path):
        abreka_dir = abreka_project / "abreka"
        self._install_default_prompts(abreka_dir)
        exp = Experiment(
            id="0001",
            status=ExperimentStatus.RUNNING,
            hypothesis="hypothesis",
            method="method",
        )
        create_experiment_dir(abreka_dir, exp.id)
        save_experiment(exp, abreka_dir)
        output_path = abreka_dir / "experiments" / exp.id / "artifacts" / "run_output.txt"
        output_path.write_text("Traceback\nValueError: bad\n")

        runner = AutonomousRunner(abreka_dir=abreka_dir)
        variables = runner._build_run_repair_prompt_vars(
            "0001",
            repair_attempt=1,
            max_repair_attempts=3,
            run_attempt=1,
            failure_reason="run.py exited with code 1",
            output_path=output_path,
        )

        template = get_default_template("repair_run")
        self._assert_no_missing_vars(template, variables)
