"""abreka evaluate — agent interprets results and writes findings."""

from __future__ import annotations

import time
from datetime import datetime, timezone

import click
from rich.console import Console

from abreka.branch import (
    compute_branch_summary,
    render_branch_summary,
)
from abreka.config import find_abreka_dir, load_config, load_goal, load_index
from abreka.display import stream_event
from abreka.metrics import parse_metric_ref
from abreka.models import ExperimentStatus
from abreka.persistence import (
    experiment_dir,
    load_experiment,
    normalize_exp_id,
    save_experiment,
    update_index_entry,
)
from abreka.pi_client import PiRpcClient
from abreka.prompts import (
    format_wiki_index,
    load_template,
    render_template,
)
from abreka.wiki import load_wiki_index
from abreka.stage_sessions import stage_resume_message, stage_session_path
from abreka.runtime_costs import record_stage_runtime


_FINDINGS_TRUNCATE = 400


def _exp_primary(entry: dict, split: str, key: str) -> float | None:
    val = entry.get("metrics", {}).get(split, {}).get(key)
    if val is None:
        return None
    try:
        return float(val)
    except (TypeError, ValueError):
        return None


def _signed_improvement(child: float | None, parent: float | None, maximize: bool) -> float | None:
    if child is None or parent is None:
        return None
    raw = child - parent
    return raw if maximize else -raw


def _build_parent_context(
    index: dict,
    abreka_dir,
    exp_id: str,
    primary_metric: str,
    metric_direction: str,
) -> str:
    """Render parent hypothesis/method/metric/findings and delta vs parent(s)."""
    entry = index.get(exp_id)
    if not entry:
        return "No parent information (experiment not in index)."
    parents = entry.get("parents", [])
    if not parents:
        return "This experiment has no parent — it is a branch root or new direction."

    split, key = parse_metric_ref(primary_metric)
    maximize = metric_direction == "maximize"
    child_val = _exp_primary(entry, split, key)

    sections: list[str] = []
    for pid in parents:
        pentry = index.get(pid)
        if not pentry:
            sections.append(f"### Parent {pid}\n- (not found in index)")
            continue
        parent_val = _exp_primary(pentry, split, key)
        block = [f"### Parent {pid}"]
        block.append(f"- Hypothesis: {pentry.get('hypothesis', '') or '(none)'}")
        method = pentry.get("method", "")
        if method:
            block.append(f"- Method: {method}")
        if parent_val is not None:
            block.append(f"- {primary_metric}: {parent_val:.4g}")
        findings_path = experiment_dir(abreka_dir, pid) / "findings.md"
        if findings_path.is_file():
            content = findings_path.read_text().strip()
            if content:
                snippet = content[:_FINDINGS_TRUNCATE]
                if len(content) > _FINDINGS_TRUNCATE:
                    snippet += "..."
                block.append(f"- Findings: {snippet}")
        delta = _signed_improvement(child_val, parent_val, maximize)
        if delta is not None:
            block.append(f"- **Delta vs current ({exp_id})**: {delta:+.4g}")
        sections.append("\n".join(block))
    return "\n\n".join(sections)


def _find_branch_root(index: dict, exp_id: str) -> str:
    """Walk up parents to the topmost ancestor and return its id.

    When there are multiple parents, picks the first. Falls back to ``exp_id``.
    """
    visited: set[str] = set()
    current = exp_id
    while current in index and current not in visited:
        visited.add(current)
        parents = index[current].get("parents", [])
        if not parents:
            return current
        current = parents[0]
    return exp_id


def _build_branch_context(
    index: dict,
    abreka_dir,
    exp_id: str,
    primary_metric: str,
    metric_direction: str,
) -> str:
    """Render the BranchSummary for the subtree the experiment belongs to."""
    root_id = _find_branch_root(index, exp_id)
    if root_id not in index:
        return "Branch summary unavailable."
    summary = compute_branch_summary(
        index, abreka_dir, root_id, primary_metric, metric_direction,
    )
    return render_branch_summary(summary, mode="default")


def _build_global_context(
    index: dict,
    exp_id: str,
    primary_metric: str,
    metric_direction: str,
) -> str:
    """Render global-best comparison and branch-best/global-best flags."""
    split, key = parse_metric_ref(primary_metric)
    maximize = metric_direction == "maximize"
    entry = index.get(exp_id, {})
    child_val = _exp_primary(entry, split, key)

    completed = [
        (eid, _exp_primary(e, split, key))
        for eid, e in index.items()
        if e.get("status") == "completed"
    ]
    completed = [(e, v) for e, v in completed if v is not None]
    if not completed:
        return "No completed experiments to compare against."

    g_eid, g_val = max(
        completed, key=lambda pair: pair[1] if maximize else -pair[1]
    )

    # Branch-best: best in this experiment's subtree (using branch root lookup)
    root_id = _find_branch_root(index, exp_id)
    subtree_best_val: float | None = None
    subtree_best_id: str | None = None
    try:
        from abreka.branch import get_subtree_ids

        subtree_ids = get_subtree_ids(index, root_id)
    except Exception:
        subtree_ids = {exp_id}
    for eid, v in completed:
        if eid in subtree_ids:
            if (
                subtree_best_val is None
                or (v > subtree_best_val if maximize else v < subtree_best_val)
            ):
                subtree_best_val = v
                subtree_best_id = eid

    lines: list[str] = []
    lines.append(f"- Global best: {g_eid} ({primary_metric}={g_val:.4g})")
    if subtree_best_id is not None and subtree_best_val is not None:
        lines.append(
            f"- Branch best ({root_id}): {subtree_best_id} "
            f"({primary_metric}={subtree_best_val:.4g})"
        )

    if child_val is not None:
        new_global = g_eid == exp_id
        new_branch_best = subtree_best_id == exp_id
        gap_global = _signed_improvement(g_val, child_val, maximize)
        if gap_global is not None:
            lines.append(f"- Gap to global best: {abs(gap_global):.4g}")
        if subtree_best_val is not None:
            gap_branch = _signed_improvement(subtree_best_val, child_val, maximize)
            if gap_branch is not None:
                lines.append(f"- Gap to branch best: {abs(gap_branch):.4g}")
        if new_global:
            lines.append("- **This experiment is the new global best.**")
        elif new_branch_best:
            lines.append("- **This experiment is the new branch best.**")
    else:
        lines.append(f"- Current experiment {exp_id} has no {primary_metric} recorded yet.")

    return "\n".join(lines)


def _list_artifacts(exp_dir) -> str:
    """List artifact files for the experiment."""
    artifacts_dir = exp_dir / "artifacts"
    if not artifacts_dir.is_dir():
        return "No artifacts."
    files = []
    for f in sorted(artifacts_dir.rglob("*")):
        if f.is_file() and not f.name.startswith("pi_session_"):
            files.append(str(f.relative_to(exp_dir)))
    return "\n".join(files) if files else "No artifacts."


@click.command("evaluate")
@click.argument("exp_id")
@click.option("--resume", is_flag=True, help="Resume the existing evaluate session.")
def evaluate_cmd(exp_id: str, resume: bool):
    """Agent reads outputs, writes findings, and extracts metrics."""
    abreka_dir = find_abreka_dir()
    config = load_config(abreka_dir)
    exp_id = normalize_exp_id(exp_id)
    exp = load_experiment(abreka_dir, exp_id)
    session_path = stage_session_path(abreka_dir, exp_id, "evaluate")

    if exp.status not in (ExperimentStatus.RUNNING, ExperimentStatus.COMPLETED):
        raise click.ClickException(
            f"Experiment {exp_id} is '{exp.status.value}', expected 'running' or 'completed'."
        )
    if resume:
        if not session_path.is_file():
            raise click.ClickException(
                f"No existing evaluate session found for experiment {exp_id}."
            )
    elif session_path.exists():
        raise click.ClickException(
            f"Evaluate session already exists for experiment {exp_id}. Use --resume to continue it."
        )

    console = Console()
    console.print(f"[bold]Evaluating experiment {exp_id}...[/bold]\n")

    exp_path = experiment_dir(abreka_dir, exp_id)

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

    goal = load_goal(abreka_dir)
    index = load_index(abreka_dir)
    template = load_template(abreka_dir, "evaluate")
    parent_context = _build_parent_context(
        index, abreka_dir, exp_id,
        config.experiment.primary_metric, config.experiment.metric_direction,
    )
    branch_context = _build_branch_context(
        index, 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,
    )
    prompt_text = render_template(template, {
        "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": _list_artifacts(exp_path),
        "parent_context": parent_context,
        "branch_summary": branch_context,
        "global_context": global_context,
        "wiki_index": format_wiki_index(load_wiki_index(abreka_dir)),
    })
    message = stage_resume_message("evaluate", exp_id) if resume else prompt_text

    model = config.agent.evaluate_model or config.agent.model
    started_at = datetime.now(timezone.utc)
    started_monotonic = time.monotonic()
    stage_status = "ok"

    try:
        with PiRpcClient(
            provider=config.agent.provider,
            model=model,
            cwd=abreka_dir.parent,
            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,
        ) as pi:
            pi.prompt(message, on_event=stream_event)

        console.print(f"\n[bold green]Evaluation complete for {exp_id}[/bold green]")
    except Exception:
        if stage_status == "ok":
            stage_status = "error"
        raise
    finally:
        record_stage_runtime(
            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),
        )
