"""Autonomous experiment loop orchestration."""

from __future__ import annotations

import random
import re
import shutil
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from rich.console import Console

from abreka.config import load_config, load_goal, load_index, project_root
from abreka.display import stream_event
from abreka.metrics import extract_metrics_from_file, extract_metrics_from_output
from abreka.models import Config, ExperimentStatus
from abreka.persistence import (
    experiment_dir,
    load_experiment,
    reserve_experiment,
    save_experiment,
    update_index_entry,
)
from abreka.stage_sessions import stage_session_path
from abreka.pi_client import PiRpcClient
from abreka.prompts import (
    format_guidance,
    format_wiki_index,
    load_template,
    render_template,
)
from abreka.run_executor import RunExecutionResult, build_run_executor
from abreka.run_control import consume_stop_after_current_request
from abreka.runtime_costs import record_stage_runtime
from abreka.wiki import load_wiki_index

_DURATION_RE = re.compile(r"^(\d+)(h|m|s)$")


def parse_duration(spec: str) -> float:
    """Parse a duration string like '24h', '2h', '30m', '90s' into seconds."""
    m = _DURATION_RE.match(spec.strip().lower())
    if not m:
        raise ValueError(f"Invalid duration: {spec!r}. Expected format: 24h, 2h, 30m, 90s")
    value = int(m.group(1))
    unit = m.group(2)
    multipliers = {"h": 3600, "m": 60, "s": 1}
    return value * multipliers[unit]


@dataclass
class AutonomousRunner:
    """Runs the propose -> implement -> test -> run -> evaluate loop."""

    abreka_dir: Path
    duration_seconds: float | None = None  # None = single step
    explore_probability: float | None = None  # None = never explore
    focus_root: str | None = None  # None = no focus constraint
    guidance: str = ""  # User-provided guidance injected into propose prompt
    debug: bool = False
    console: Console = field(default_factory=Console)
    consecutive_failures: int = field(default=0, init=False)
    MAX_CONSECUTIVE_FAILURES: int = 5

    def _honor_stop_after_current_request(self) -> bool:
        """Stop cleanly before starting a new experiment when requested."""
        if not consume_stop_after_current_request(self.abreka_dir):
            return False
        self.console.print(
            "\n[bold yellow]Stop requested; exiting before starting a new experiment.[/bold yellow]"
        )
        return True

    def run(self) -> None:
        start = time.monotonic()
        iteration = 0

        while True:
            if self._honor_stop_after_current_request():
                break

            iteration += 1
            self.console.print(f"\n[bold]{'='*60}[/bold]")
            self.console.print(f"[bold]Iteration {iteration}[/bold]")
            self.console.print(f"[bold]{'='*60}[/bold]\n")

            stop_for_failures = False
            try:
                exp_id = self._propose()
                self._implement(exp_id)
                self._test(exp_id)
                self._run_experiment(exp_id)
                self._evaluate(exp_id)
                self.consecutive_failures = 0
            except Exception as e:
                self.consecutive_failures += 1
                self.console.print(f"\n[red]Error: {e}[/red]")
                if self.consecutive_failures >= self.MAX_CONSECUTIVE_FAILURES:
                    self.console.print(
                        f"[bold red]{self.consecutive_failures} consecutive failures. Stopping.[/bold red]"
                    )
                    stop_for_failures = True

            if self._honor_stop_after_current_request():
                break
            if stop_for_failures:
                break

            # Check if we should stop
            if self.duration_seconds is None:
                break  # --step mode
            elapsed = time.monotonic() - start
            if elapsed >= self.duration_seconds:
                self.console.print(f"\n[bold]Duration reached ({elapsed:.0f}s). Stopping.[/bold]")
                break

    def _log_prompt(self, template_name: str, prompt_text: str, model: str) -> None:
        """Print a compact summary of the prompt being sent."""
        lines = prompt_text.strip().splitlines()
        # Count non-empty lines
        line_count = len([l for l in lines if l.strip()])
        char_count = len(prompt_text)
        self.console.print(
            f"  [dim]prompt: {template_name}.md ({line_count} lines, {char_count} chars) "
            f"→ {model}[/dim]"
        )

    def _should_explore(self) -> bool:
        """Decide whether this iteration should use the explore prompt."""
        if self.explore_probability is None or self.explore_probability <= 0.0:
            return False
        if self.explore_probability >= 1.0:
            return True
        return random.random() < self.explore_probability

    def _get_pi_client(
        self,
        model_override: str = "",
        session_path: Path | None = None,
    ) -> PiRpcClient:
        config = load_config(self.abreka_dir)
        model = model_override or config.agent.model
        return PiRpcClient(
            provider=config.agent.provider,
            model=model,
            cwd=project_root(self.abreka_dir),
            thinking_level=config.agent.thinking_level,
            extra_args=config.agent.pi_extra_args,
            session_path=session_path,
            session_capture_mode=config.agent.session_capture_mode,
            debug=self.debug,
        )

    def _build_propose_prompt_vars(
        self,
        template_name: str,
        *,
        reserved_exp_id: str = "",
        exclude_exp_id: str | None = None,
    ) -> dict[str, str]:
        """Assemble the template-variable dict for the propose prompt.

        Kept as a pure helper (no pi/IO beyond reads) so tests can verify that
        every placeholder in the default propose templates is supplied.
        """
        config = load_config(self.abreka_dir)
        goal = load_goal(self.abreka_dir)
        index = load_index(self.abreka_dir)
        if exclude_exp_id is not None:
            index = {eid: entry for eid, entry in index.items() if eid != exclude_exp_id}

        from abreka.cli.agent.propose import (
            _build_backlog_context,
            _build_branch_overview,
            _build_experiment_context,
            _build_focus_constraint,
            _build_focused_branch_summary,
        )
        context = _build_experiment_context(
            index, self.abreka_dir,
            primary_metric=config.experiment.primary_metric,
            metric_direction=config.experiment.metric_direction,
            focus_root=self.focus_root,
        )
        focus_constraint = _build_focus_constraint(index, self.focus_root) if self.focus_root else ""
        backlog_context = _build_backlog_context(
            self.abreka_dir, index=index, focus_root=self.focus_root,
        )
        branch_overview = _build_branch_overview(
            index, self.abreka_dir,
            config.experiment.primary_metric, config.experiment.metric_direction,
            focus_root=self.focus_root,
        )
        focused_branch_summary = _build_focused_branch_summary(
            index, self.abreka_dir,
            config.experiment.primary_metric, config.experiment.metric_direction,
            self.focus_root,
        )
        return {
            "goal": goal,
            "primary_metric": config.experiment.primary_metric,
            "metric_direction": config.experiment.metric_direction,
            "experiment_context": context,
            "focus_constraint": focus_constraint,
            "backlog_context": backlog_context,
            "guidance": format_guidance(self.guidance),
            "branch_overview": branch_overview,
            "focused_branch_summary": focused_branch_summary,
            "wiki_index": format_wiki_index(load_wiki_index(self.abreka_dir)),
            "reserved_exp_id": reserved_exp_id,
        }

    def _propose(self) -> str:
        """Run the propose step. Returns the reserved experiment ID."""
        exploring = self._should_explore()
        template_name = "propose_explore" if exploring else "propose_exploit"
        mode_label = "[magenta]explore[/magenta]" if exploring else "exploit"
        if self.focus_root:
            mode_label += f" [bold yellow]focus={self.focus_root}[/bold yellow]"
        self.console.print(f"[bold cyan]Step 1: Propose[/bold cyan] ({mode_label})")
        config = load_config(self.abreka_dir)
        reserved = reserve_experiment(self.abreka_dir)
        exp_id = reserved.id

        started_at = datetime.now(timezone.utc)
        started_monotonic = time.monotonic()
        stage_status = "ok"

        try:
            template = load_template(self.abreka_dir, template_name)
            prompt_text = render_template(
                template,
                self._build_propose_prompt_vars(
                    template_name,
                    reserved_exp_id=exp_id,
                    exclude_exp_id=exp_id,
                ),
            )

            model = config.agent.propose_model or config.agent.model
            self._log_prompt(template_name, prompt_text, model)
            session_path = stage_session_path(self.abreka_dir, exp_id, "propose")
            with self._get_pi_client(model, session_path=session_path) as pi:
                pi.prompt(prompt_text, on_event=stream_event)

            updated = load_experiment(self.abreka_dir, exp_id)
            if updated.status != ExperimentStatus.PROPOSED:
                stage_status = "error"
                if updated.status == ExperimentStatus.FAILED and updated.failure_reason:
                    raise RuntimeError(updated.failure_reason)
                raise RuntimeError(
                    f"Propose session for experiment {exp_id} did not finalize the reserved experiment."
                )

            self.console.print(f"  Created experiment [bold]{exp_id}[/bold]")
            return exp_id
        except Exception:
            if stage_status == "ok":
                stage_status = "error"
            raise
        finally:
            record_stage_runtime(
                self.abreka_dir,
                exp_id,
                stage="propose",
                elapsed_seconds=time.monotonic() - started_monotonic,
                status=stage_status,
                started_at=started_at,
                finished_at=datetime.now(timezone.utc),
            )

    def _implement(self, exp_id: str) -> None:
        """Run the implement step."""
        self.console.print(f"[bold cyan]Step 2: Implement {exp_id}[/bold cyan]")
        config = load_config(self.abreka_dir)
        exp = load_experiment(self.abreka_dir, exp_id)

        started_at = datetime.now(timezone.utc)
        started_monotonic = time.monotonic()
        stage_status = "ok"

        try:
            exp.status = ExperimentStatus.IMPLEMENTING
            save_experiment(exp, self.abreka_dir)
            update_index_entry(self.abreka_dir, exp)

            goal = load_goal(self.abreka_dir)
            index = load_index(self.abreka_dir)
            from abreka.cli.agent.propose import _build_experiment_context
            context = _build_experiment_context(
                index, self.abreka_dir,
                primary_metric=config.experiment.primary_metric,
                metric_direction=config.experiment.metric_direction,
            )

            template = load_template(self.abreka_dir, "implement")
            prompt_text = render_template(template, {
                "goal": goal,
                "exp_id": exp_id,
                "hypothesis": exp.hypothesis,
                "method": exp.method,
                "parents": ", ".join(exp.parents) or "none",
                "lib_name": config.project.lib,
                "experiment_context": context,
                "primary_metric": config.experiment.primary_metric,
                "metric_direction": config.experiment.metric_direction,
                "wiki_index": format_wiki_index(load_wiki_index(self.abreka_dir)),
            })

            model = config.agent.implement_model or config.agent.model
            self._log_prompt("implement", prompt_text, model)
            session_path = stage_session_path(self.abreka_dir, exp_id, "implement")
            if session_path.exists():
                stage_status = "error"
                raise RuntimeError(
                    f"Implement session already exists for experiment {exp_id}. Resume it manually instead of restarting."
                )
            with self._get_pi_client(model, session_path=session_path) as pi:
                pi.prompt(prompt_text, on_event=stream_event)
        except Exception:
            if stage_status == "ok":
                stage_status = "error"
            raise
        finally:
            record_stage_runtime(
                self.abreka_dir,
                exp_id,
                stage="implement",
                elapsed_seconds=time.monotonic() - started_monotonic,
                status=stage_status,
                started_at=started_at,
                finished_at=datetime.now(timezone.utc),
            )

    def _test(self, exp_id: str) -> None:
        """Run the test step."""
        self.console.print(f"[bold cyan]Step 3: Test {exp_id}[/bold cyan]")
        config = load_config(self.abreka_dir)
        exp = load_experiment(self.abreka_dir, exp_id)

        started_at = datetime.now(timezone.utc)
        started_monotonic = time.monotonic()
        stage_status = "ok"

        try:
            exp.status = ExperimentStatus.TESTING
            save_experiment(exp, self.abreka_dir)
            update_index_entry(self.abreka_dir, exp)

            template = load_template(self.abreka_dir, "test")
            prompt_text = render_template(template, {
                "exp_id": exp_id,
                "hypothesis": exp.hypothesis,
                "method": exp.method,
                "lib_name": config.project.lib,
                "test_command": config.project.test_command,
                "wiki_index": format_wiki_index(load_wiki_index(self.abreka_dir)),
            })

            self._log_prompt("test", prompt_text, config.agent.model)
            session_path = stage_session_path(self.abreka_dir, exp_id, "test")
            if session_path.exists():
                stage_status = "error"
                raise RuntimeError(
                    f"Test session already exists for experiment {exp_id}. Resume it manually instead of restarting."
                )
            with self._get_pi_client(session_path=session_path) as pi:
                pi.prompt(prompt_text, on_event=stream_event)
        except Exception:
            if stage_status == "ok":
                stage_status = "error"
            raise
        finally:
            record_stage_runtime(
                self.abreka_dir,
                exp_id,
                stage="test",
                elapsed_seconds=time.monotonic() - started_monotonic,
                status=stage_status,
                started_at=started_at,
                finished_at=datetime.now(timezone.utc),
            )

    def _mark_experiment_running(self, exp_id: str) -> None:
        """Set an experiment to running and clear any stale failure reason."""
        exp = load_experiment(self.abreka_dir, exp_id)
        exp.status = ExperimentStatus.RUNNING
        exp.failure_reason = ""
        save_experiment(exp, self.abreka_dir)
        update_index_entry(self.abreka_dir, exp)

    def _mark_experiment_failed(self, exp_id: str, reason: str) -> None:
        """Set an experiment to failed with a reason."""
        exp = load_experiment(self.abreka_dir, exp_id)
        exp.status = ExperimentStatus.FAILED
        exp.failure_reason = reason
        save_experiment(exp, self.abreka_dir)
        update_index_entry(self.abreka_dir, exp)

    def _extract_and_store_run_metrics(
        self,
        exp_id: str,
        result: RunExecutionResult,
        output_path: Path,
    ) -> dict[str, dict[str, float]]:
        """Extract run metrics from executor stdout or run_output.txt and persist them."""
        if result.stdout:
            metrics = extract_metrics_from_output(result.stdout)
        elif output_path.is_file():
            metrics = extract_metrics_from_file(output_path)
        else:
            metrics = {}

        if metrics:
            exp = load_experiment(self.abreka_dir, exp_id)
            exp.metrics = metrics
            save_experiment(exp, self.abreka_dir)
            update_index_entry(self.abreka_dir, exp)
        return metrics

    def _snapshot_run_failure_output(
        self,
        exp_id: str,
        run_attempt: int,
        output_path: Path,
    ) -> str:
        """Copy the current run output to a per-attempt failure artifact."""
        if not output_path.is_file():
            return ""
        exp_path = experiment_dir(self.abreka_dir, exp_id)
        failure_dir = exp_path / "artifacts" / "run_failures"
        failure_dir.mkdir(parents=True, exist_ok=True)
        snapshot_path = failure_dir / f"run_attempt_{run_attempt}.txt"
        shutil.copyfile(output_path, snapshot_path)
        return snapshot_path.relative_to(exp_path).as_posix()

    def _append_run_failure(
        self,
        exp_id: str,
        *,
        run_attempt: int,
        repair_attempt: int,
        reason: str,
        result: RunExecutionResult,
        output_path: Path,
    ) -> dict[str, object]:
        """Append a run failure record and preserve the corresponding output."""
        exp = load_experiment(self.abreka_dir, exp_id)
        entry: dict[str, object] = {
            "stage": "run",
            "attempt": run_attempt,
            "repair_attempt": repair_attempt,
            "reason": reason,
            "returncode": result.returncode,
            "timed_out": result.timed_out,
            "occurred_at": datetime.now(timezone.utc).isoformat(),
        }
        output_rel = self._snapshot_run_failure_output(exp_id, run_attempt, output_path)
        if output_rel:
            entry["output_path"] = output_rel
        exp.failures = [*list(exp.failures or []), entry]
        save_experiment(exp, self.abreka_dir)
        update_index_entry(self.abreka_dir, exp)
        return entry

    def _read_text_tail(self, path: Path, max_chars: int = 12000) -> str:
        """Read a bounded tail from a text file for prompts."""
        if not path.is_file():
            return "(no output captured)"
        text = path.read_text(errors="replace")
        if len(text) <= max_chars:
            return text
        return f"... (truncated to last {max_chars} chars) ...\n" + text[-max_chars:]

    def _format_failure_history(self, exp_id: str) -> str:
        """Render saved failure records compactly for repair prompts."""
        exp = load_experiment(self.abreka_dir, exp_id)
        if not exp.failures:
            return "No prior run failures recorded."
        lines: list[str] = []
        for failure in exp.failures:
            if failure.get("stage") != "run":
                continue
            repair_attempt = failure.get("repair_attempt", 0)
            repair_label = (
                "initial run" if repair_attempt in (0, "0", None)
                else f"after repair {repair_attempt}"
            )
            bits = [
                f"- run attempt {failure.get('attempt', '?')}",
                repair_label,
                f"{failure.get('reason', '')}",
            ]
            if failure.get("output_path"):
                bits.append(f"output: `{failure['output_path']}`")
            lines.append("; ".join(bits))
        return "\n".join(lines) if lines else "No prior run failures recorded."

    def _build_run_repair_prompt_vars(
        self,
        exp_id: str,
        *,
        repair_attempt: int,
        max_repair_attempts: int,
        run_attempt: int,
        failure_reason: str,
        output_path: Path,
    ) -> dict[str, str]:
        """Assemble variables for the run-repair prompt."""
        config = load_config(self.abreka_dir)
        exp = load_experiment(self.abreka_dir, exp_id)
        run_py = experiment_dir(self.abreka_dir, exp_id) / "run.py"
        try:
            run_rel = run_py.relative_to(project_root(self.abreka_dir)).as_posix()
        except ValueError:
            run_rel = str(run_py)
        return {
            "goal": load_goal(self.abreka_dir),
            "exp_id": exp_id,
            "hypothesis": exp.hypothesis,
            "method": exp.method,
            "primary_metric": config.experiment.primary_metric,
            "metric_direction": config.experiment.metric_direction,
            "lib_name": config.project.lib,
            "test_command": config.project.test_command,
            "run_command": f"uv run {run_rel}",
            "failure_reason": failure_reason,
            "run_attempt": str(run_attempt),
            "repair_attempt": str(repair_attempt),
            "max_repair_attempts": str(max_repair_attempts),
            "run_output": self._read_text_tail(output_path),
            "failure_history": self._format_failure_history(exp_id),
            "wiki_index": format_wiki_index(load_wiki_index(self.abreka_dir)),
        }

    def _repair_run_failure(
        self,
        exp_id: str,
        *,
        repair_attempt: int,
        max_repair_attempts: int,
        run_attempt: int,
        failure_reason: str,
        output_path: Path,
    ) -> None:
        """Ask an agent to repair a failed run before retrying run.py."""
        self.console.print(
            f"[bold cyan]Step 4a: Repair run {exp_id}[/bold cyan] "
            f"(attempt {repair_attempt}/{max_repair_attempts})"
        )
        config = load_config(self.abreka_dir)
        started_at = datetime.now(timezone.utc)
        started_monotonic = time.monotonic()
        stage_status = "ok"

        try:
            template = load_template(self.abreka_dir, "repair_run")
            prompt_text = render_template(
                template,
                self._build_run_repair_prompt_vars(
                    exp_id,
                    repair_attempt=repair_attempt,
                    max_repair_attempts=max_repair_attempts,
                    run_attempt=run_attempt,
                    failure_reason=failure_reason,
                    output_path=output_path,
                ),
            )
            model = config.agent.implement_model or config.agent.model
            self._log_prompt("repair_run", prompt_text, model)
            session_path = stage_session_path(
                self.abreka_dir,
                exp_id,
                f"run_repair_{repair_attempt}",
            )
            if session_path.exists():
                stage_status = "error"
                raise RuntimeError(
                    f"Run repair session already exists for experiment {exp_id} "
                    f"attempt {repair_attempt}."
                )
            with self._get_pi_client(model, session_path=session_path) as pi:
                pi.prompt(prompt_text, on_event=stream_event)
        except Exception:
            if stage_status == "ok":
                stage_status = "error"
            raise
        finally:
            record_stage_runtime(
                self.abreka_dir,
                exp_id,
                stage="run_repair",
                elapsed_seconds=time.monotonic() - started_monotonic,
                status=stage_status,
                started_at=started_at,
                finished_at=datetime.now(timezone.utc),
            )

    def _apply_run_result(
        self,
        exp_id: str,
        result: RunExecutionResult,
        timeout_minutes: int,
        output_path: Path,
    ) -> dict[str, dict[str, float]]:
        """Persist metrics/failure state from a completed run executor result."""
        metrics = self._extract_and_store_run_metrics(exp_id, result, output_path)

        if result.timed_out:
            reason = f"Timed out after {timeout_minutes}m"
            if metrics:
                self.console.print(f"  Partial metrics: {metrics}")
                reason += " (partial metrics captured)"
            else:
                reason += " with no metrics"
            self._mark_experiment_failed(exp_id, reason)
            raise RuntimeError(reason)
        if result.returncode != 0:
            reason = f"run.py exited with code {result.returncode}"
            self._mark_experiment_failed(exp_id, reason)
            raise RuntimeError(reason)

        return metrics

    def _execute_run_attempt(
        self,
        exp_id: str,
        *,
        executor,
        run_py: Path,
        output_path: Path,
        timeout_seconds: float,
    ) -> RunExecutionResult:
        """Execute run.py once and record runtime for that attempt."""
        self._mark_experiment_running(exp_id)
        started_at = datetime.now(timezone.utc)
        started_monotonic = time.monotonic()
        stage_status = "ok"
        try:
            result = executor.execute(
                run_py=run_py,
                output_path=output_path,
                timeout_seconds=timeout_seconds,
            )
            if result.timed_out:
                stage_status = "timed_out"
            elif result.returncode != 0:
                stage_status = "error"
            return result
        except Exception as e:
            stage_status = "error"
            self._mark_experiment_failed(exp_id, str(e))
            raise RuntimeError(str(e)) from e
        finally:
            record_stage_runtime(
                self.abreka_dir,
                exp_id,
                stage="run",
                elapsed_seconds=time.monotonic() - started_monotonic,
                status=stage_status,
                started_at=started_at,
                finished_at=datetime.now(timezone.utc),
            )

    def _run_experiment(self, exp_id: str) -> None:
        """Execute run.py, repairing non-timeout failures before giving up."""
        self.console.print(f"[bold cyan]Step 4: Run {exp_id}[/bold cyan]")
        config = load_config(self.abreka_dir)
        max_repair_attempts = max(0, int(config.run.max_repair_attempts))
        run_py = experiment_dir(self.abreka_dir, exp_id) / "run.py"
        if not run_py.is_file():
            reason = f"run.py not found at {run_py}"
            self._mark_experiment_failed(exp_id, reason)
            raise RuntimeError(reason)

        timeout = config.run.timeout_minutes * 60
        output_path = experiment_dir(self.abreka_dir, exp_id) / "artifacts" / "run_output.txt"
        executor = build_run_executor(self.abreka_dir, config, self.console)

        run_attempt = 0
        repair_attempts_used = 0
        while True:
            run_attempt += 1
            result = self._execute_run_attempt(
                exp_id,
                executor=executor,
                run_py=run_py,
                output_path=output_path,
                timeout_seconds=timeout,
            )
            metrics = self._extract_and_store_run_metrics(exp_id, result, output_path)

            if result.timed_out:
                reason = f"Timed out after {config.run.timeout_minutes}m"
                if metrics:
                    self.console.print(f"  Partial metrics: {metrics}")
                    reason += " (partial metrics captured)"
                else:
                    reason += " with no metrics"
                self._append_run_failure(
                    exp_id,
                    run_attempt=run_attempt,
                    repair_attempt=repair_attempts_used,
                    reason=reason,
                    result=result,
                    output_path=output_path,
                )
                self._mark_experiment_failed(exp_id, reason)
                raise RuntimeError(reason)

            if result.returncode == 0:
                self.console.print(f"  Run completed. Metrics: {metrics}")
                return

            reason = f"run.py exited with code {result.returncode}"
            self._append_run_failure(
                exp_id,
                run_attempt=run_attempt,
                repair_attempt=repair_attempts_used,
                reason=reason,
                result=result,
                output_path=output_path,
            )

            if repair_attempts_used >= max_repair_attempts:
                final_reason = (
                    f"{reason} after {max_repair_attempts} repair attempt"
                    f"{'s' if max_repair_attempts != 1 else ''}"
                )
                self._mark_experiment_failed(exp_id, final_reason)
                raise RuntimeError(final_reason)

            repair_attempts_used += 1
            self.console.print(
                f"  [yellow]Run failed ({reason}); invoking repair "
                f"{repair_attempts_used}/{max_repair_attempts}[/yellow]"
            )
            try:
                self._repair_run_failure(
                    exp_id,
                    repair_attempt=repair_attempts_used,
                    max_repair_attempts=max_repair_attempts,
                    run_attempt=run_attempt,
                    failure_reason=reason,
                    output_path=output_path,
                )
            except Exception as e:
                repair_reason = f"run repair attempt {repair_attempts_used} failed: {e}"
                self._mark_experiment_failed(exp_id, repair_reason)
                raise RuntimeError(repair_reason) from e

    def _build_evaluate_prompt_vars(self, exp_id: str) -> dict[str, str]:
        """Assemble the template-variable dict for the evaluate prompt.

        Pure helper (no pi/IO beyond reads) so tests can verify full coverage
        of the default evaluate template.
        """
        config = load_config(self.abreka_dir)
        exp = load_experiment(self.abreka_dir, exp_id)
        exp_path = experiment_dir(self.abreka_dir, exp_id)
        index = load_index(self.abreka_dir)

        goal = load_goal(self.abreka_dir)

        run_output_path = exp_path / "artifacts" / "run_output.txt"
        run_output = run_output_path.read_text() if run_output_path.is_file() else "(no output)"

        from abreka.cli.agent.evaluate import (
            _build_branch_context,
            _build_global_context,
            _build_parent_context,
            _list_artifacts,
        )
        artifacts_list = _list_artifacts(exp_path)
        parent_context = _build_parent_context(
            index, self.abreka_dir, exp_id,
            config.experiment.primary_metric, config.experiment.metric_direction,
        )
        branch_summary = _build_branch_context(
            index, self.abreka_dir, exp_id,
            config.experiment.primary_metric, config.experiment.metric_direction,
        )
        global_context = _build_global_context(
            index, exp_id,
            config.experiment.primary_metric, config.experiment.metric_direction,
        )
        return {
            "goal": goal,
            "exp_id": exp_id,
            "hypothesis": exp.hypothesis,
            "method": exp.method,
            "primary_metric": config.experiment.primary_metric,
            "metric_direction": config.experiment.metric_direction,
            "run_output": run_output,
            "artifacts_list": artifacts_list,
            "parent_context": parent_context,
            "branch_summary": branch_summary,
            "global_context": global_context,
            "wiki_index": format_wiki_index(load_wiki_index(self.abreka_dir)),
        }

    def _evaluate(self, exp_id: str) -> None:
        """Run the evaluate step."""
        self.console.print(f"[bold cyan]Step 5: Evaluate {exp_id}[/bold cyan]")
        config = load_config(self.abreka_dir)
        exp_path = experiment_dir(self.abreka_dir, exp_id)

        started_at = datetime.now(timezone.utc)
        started_monotonic = time.monotonic()
        stage_status = "ok"

        try:
            template = load_template(self.abreka_dir, "evaluate")
            prompt_text = render_template(
                template, self._build_evaluate_prompt_vars(exp_id)
            )

            model = config.agent.evaluate_model or config.agent.model
            self._log_prompt("evaluate", prompt_text, model)
            session_path = stage_session_path(self.abreka_dir, exp_id, "evaluate")
            if session_path.exists():
                stage_status = "error"
                raise RuntimeError(
                    f"Evaluate session already exists for experiment {exp_id}. Resume it manually instead of restarting."
                )
            with self._get_pi_client(model, session_path=session_path) as pi:
                pi.prompt(prompt_text, on_event=stream_event)

            self.console.print(f"  Evaluation complete for [bold]{exp_id}[/bold]")
        except Exception:
            if stage_status == "ok":
                stage_status = "error"
            raise
        finally:
            record_stage_runtime(
                self.abreka_dir,
                exp_id,
                stage="evaluate",
                elapsed_seconds=time.monotonic() - started_monotonic,
                status=stage_status,
                started_at=started_at,
                finished_at=datetime.now(timezone.utc),
            )
