"""abreka propose — agent proposes the next experiment."""

from __future__ import annotations

import sys
import time
from datetime import datetime, timezone

import click
from rich.console import Console

from abreka.branch import (
    compute_branch_summary,
    render_branch_overview,
    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.pi_client import PiRpcClient
from abreka.persistence import (
    experiment_dir,
    load_backlog_index,
    load_experiment,
    normalize_exp_id,
    reserve_experiment,
)
from abreka.stage_sessions import stage_resume_message, stage_session_path
from abreka.runtime_costs import record_stage_runtime
from abreka.prompts import (
    format_guidance,
    format_wiki_index,
    load_template,
    render_template,
)
from abreka.wiki import load_wiki_index

_FINDINGS_TRUNCATE = 500
_FRONTIER_LIMIT = 5
_OUTSIDE_TOP_N = 5


def get_subtree_ids(index: dict, root_id: str) -> set[str]:
    """Return *root_id* plus all transitive descendants in the experiment DAG."""
    if root_id not in index:
        raise ValueError(f"Experiment {root_id!r} not found in index.")
    # Build parent -> children mapping
    children: dict[str, list[str]] = {}
    for eid, entry in index.items():
        for pid in entry.get("parents", []):
            children.setdefault(pid, []).append(eid)
    # BFS
    result: set[str] = set()
    queue = [root_id]
    while queue:
        eid = queue.pop()
        if eid in result:
            continue
        result.add(eid)
        queue.extend(children.get(eid, []))
    return result


def _build_branch_overview(
    index: dict,
    abreka_dir,
    primary_metric: str,
    metric_direction: str,
    *,
    focus_root: str | None = None,
    top_k: int = 5,
) -> str:
    """Build a compact branch overview for the propose prompt.

    Summarizes every parentless root with a non-trivial subtree. When ``focus_root``
    is set, that subtree's summary is excluded (it is rendered separately as the
    focused branch summary).
    """
    roots = [eid for eid, entry in index.items() if not entry.get("parents")]
    summaries = []
    for rid in roots:
        if rid == focus_root:
            continue
        summary = compute_branch_summary(
            index, abreka_dir, rid, primary_metric, metric_direction,
        )
        if summary.subtree_size < 2:
            continue
        summaries.append(summary)
    if not summaries:
        return "No multi-experiment branches yet."
    # Cap to top_k by attention
    from abreka.branch import _sort_summaries  # internal helper

    ranked = _sort_summaries(summaries, "attention")[:top_k]
    return render_branch_overview(ranked, "attention")


def _build_focused_branch_summary(
    index: dict,
    abreka_dir,
    primary_metric: str,
    metric_direction: str,
    focus_root: str | None,
) -> str:
    """Build a full branch summary for the focused subtree, or empty string."""
    if not focus_root:
        return ""
    summary = compute_branch_summary(
        index, abreka_dir, focus_root, primary_metric, metric_direction,
    )
    body = render_branch_summary(summary, mode="default")
    return f"## Focused Branch Summary\n\n{body}"


def _build_focus_constraint(index: dict, focus_root: str) -> str:
    """Build prompt text constraining the agent to the focused subtree."""
    subtree_ids = get_subtree_ids(index, focus_root)
    valid_parents = sorted(
        eid for eid in subtree_ids
        if index[eid].get("status") == "completed"
    )
    ids_list = ", ".join(valid_parents)
    return (
        f"## FOCUS CONSTRAINT\n\n"
        f"You are focused on the subtree rooted at experiment **{focus_root}**.\n\n"
        f"**You MUST set `--parent` to one of these experiment IDs:** {ids_list}\n\n"
        f"Do NOT create an experiment without a parent. "
        f"Do NOT use a parent outside this list. "
        f"Your goal is to advance research within this specific branch of the experiment tree."
    )


def _build_backlog_context(abreka_dir, index: dict | None = None, focus_root: str | None = None) -> str:
    """Build a compact summary of open backlog items for prompt context.

    When *focus_root* is set, only include items whose source_exp is within
    the focused subtree.  Items with no source_exp are excluded under focus.
    """
    backlog_index = load_backlog_index(abreka_dir)
    if not backlog_index:
        return "No backlog ideas."

    # Filter to open items
    open_items = [e for e in backlog_index.values() if e.get("status") == "open"]

    # Filter by subtree when focused
    if focus_root and index:
        subtree_ids = get_subtree_ids(index, focus_root)
        open_items = [
            e for e in open_items
            if e.get("source_exp") and e["source_exp"] in subtree_ids
        ]

    if not open_items:
        return "No backlog ideas."

    # Sort by priority descending, then by ID
    open_items.sort(key=lambda e: (-e.get("priority", 0), e.get("id", "")))

    lines = []
    for e in open_items:
        bid = e["id"]
        pri = e.get("priority", 0)
        title = e.get("title", "")
        desc = e.get("description", "")
        source = e.get("source_exp", "")
        tags = e.get("tags", [])

        parts = [f"- **B{bid}** [priority={pri}] \"{title}\""]
        if desc:
            short_desc = desc[:120] + "..." if len(desc) > 120 else desc
            parts[0] += f" — {short_desc}"
        meta = []
        if source:
            meta.append(f"source: {source}")
        if tags:
            meta.append(f"tags: {', '.join(tags)}")
        if meta:
            parts[0] += f" ({'; '.join(meta)})"
        lines.append(parts[0])

    return "\n".join(lines)


def _format_metrics(metrics: dict) -> str:
    parts = []
    for split in sorted(metrics):
        for key in sorted(metrics[split]):
            parts.append(f"{split}:{key}={metrics[split][key]}")
    return ", ".join(parts) if parts else "no metrics"


def _get_primary_metric_value(
    entry: dict, primary_split: str, primary_key: str,
) -> float | None:
    return entry.get("metrics", {}).get(primary_split, {}).get(primary_key)


def _format_detail(entry: dict, abreka_dir, *, include_findings: bool = True) -> str:
    """Format a single experiment with full detail."""
    eid = entry["id"]
    status = entry.get("status", "?")
    hyp = entry.get("hypothesis", "")
    method = entry.get("method", "")
    tags = entry.get("tags", [])
    parents = ", ".join(entry.get("parents", [])) or "none"
    metrics_str = _format_metrics(entry.get("metrics", {}))

    tags_str = f" [{', '.join(tags)}]" if tags else ""
    lines = [
        f"- **{eid}** [{status}]{tags_str} {hyp}",
        f"  Method: {method}" if method else "",
        f"  Parents: {parents} | Metrics: {metrics_str}",
    ]

    if include_findings:
        findings_path = experiment_dir(abreka_dir, eid) / "findings.md"
        if findings_path.is_file():
            content = findings_path.read_text().strip()
            if content:
                truncated = content[:_FINDINGS_TRUNCATE]
                if len(content) > _FINDINGS_TRUNCATE:
                    truncated += "..."
                lines.append(f"  Findings: {truncated}")

    return "\n".join(line for line in lines if line)


def _format_compact(entry: dict, primary_split: str, primary_key: str) -> str:
    """Format a single experiment as a compact one-liner."""
    eid = entry["id"]
    tags = entry.get("tags", [])
    parents = ", ".join(entry.get("parents", [])) or "none"
    metric_val = _get_primary_metric_value(entry, primary_split, primary_key)
    metric_str = f"{metric_val}" if metric_val is not None else "-"
    tags_str = f" [{', '.join(tags)}]" if tags else ""
    hyp = entry.get("hypothesis", "")
    # Truncate hypothesis for compact view
    if len(hyp) > 80:
        hyp = hyp[:77] + "..."
    return f"- **{eid}**{tags_str} parents={parents} | {primary_split}:{primary_key}={metric_str} — {hyp}"


def _build_experiment_context(
    index: dict,
    abreka_dir,
    primary_metric: str = "val:accuracy",
    metric_direction: str = "maximize",
    focus_root: str | None = None,
) -> str:
    """Build a structured summary of all experiments for prompt context.

    Renders a frontier-focused view:
    1. Current state summary (best result, total count)
    2. Most recent experiment with full detail
    3. Top frontier experiments (no children) sorted by metric
    4. Remaining experiments in compact form

    When *focus_root* is set, the output is restructured to emphasize the
    subtree rooted at that experiment while keeping outside experiments
    visible in a collapsed section.
    """
    if not index:
        return "No experiments yet. This will be the first."

    primary_split, primary_key = parse_metric_ref(primary_metric)
    maximize = metric_direction == "maximize"

    # Classify experiments
    completed = {
        eid: e for eid, e in index.items() if e.get("status") == "completed"
    }

    def _sort_key(eid: str) -> float:
        val = _get_primary_metric_value(completed[eid], primary_split, primary_key)
        if val is None:
            return float("-inf") if maximize else float("inf")
        return val if maximize else -val

    if focus_root is not None:
        return _build_focused_context(
            index, abreka_dir, primary_split, primary_key,
            maximize, completed, _sort_key, focus_root,
        )

    return _build_unfocused_context(
        index, abreka_dir, primary_metric, primary_split, primary_key,
        maximize, completed, _sort_key,
    )


def _build_unfocused_context(
    index: dict,
    abreka_dir,
    primary_metric: str,
    primary_split: str,
    primary_key: str,
    maximize: bool,
    completed: dict,
    _sort_key,
) -> str:
    """Build experiment context without focus (original behaviour)."""
    # Build parent → children mapping for completed experiments
    children: dict[str, list[str]] = {}
    for eid, entry in completed.items():
        for pid in entry.get("parents", []):
            children.setdefault(pid, []).append(eid)

    # Frontier = completed experiments with no completed children
    frontier_ids = {eid for eid in completed if eid not in children}

    best_id = max(completed, key=_sort_key) if completed else None

    # Find most recent experiment
    most_recent_id = max(
        index.keys(),
        key=lambda eid: index[eid].get("created_at", ""),
    )

    # Sort frontier by metric
    frontier_sorted = sorted(frontier_ids, key=_sort_key, reverse=True)
    top_frontier = frontier_sorted[:_FRONTIER_LIMIT]
    remaining_frontier = frontier_sorted[_FRONTIER_LIMIT:]

    # Non-frontier completed experiments
    non_frontier_ids = sorted(
        eid for eid in completed if eid not in frontier_ids
    )

    # Non-completed experiments (proposed, running, etc.)
    non_completed = {
        eid: e for eid, e in index.items() if e.get("status") != "completed"
    }

    # === Build output ===
    sections = []

    # Section 1: Current state
    total = len(index)
    n_completed = len(completed)
    n_failed = sum(1 for e in index.values() if e.get("status") == "failed")
    n_other = total - n_completed - n_failed

    state_lines = [f"**Total experiments:** {total} ({n_completed} completed"]
    if n_failed:
        state_lines[0] += f", {n_failed} failed"
    if n_other:
        state_lines[0] += f", {n_other} in progress"
    state_lines[0] += ")"

    if best_id:
        best_val = _get_primary_metric_value(completed[best_id], primary_split, primary_key)
        if best_val is not None:
            state_lines.append(
                f"**Current best:** {best_id} ({primary_split}:{primary_key}={best_val})"
            )
    sections.append("\n".join(state_lines))

    # Section 2: Most recent experiment
    recent_lines = ["### Most Recent Experiment"]
    recent_lines.append(_format_detail(index[most_recent_id], abreka_dir))
    sections.append("\n".join(recent_lines))

    # Section 3: Top frontier
    if top_frontier:
        frontier_lines = [f"### Frontier — Top {len(top_frontier)} by {primary_metric}"]
        # Don't duplicate the most recent if it's already shown above
        for eid in top_frontier:
            if eid == most_recent_id:
                frontier_lines.append(f"- **{eid}** (see Most Recent above)")
            else:
                frontier_lines.append(_format_detail(completed[eid], abreka_dir))
        sections.append("\n".join(frontier_lines))

    # Section 4: Remaining experiments (compact)
    compact_ids = sorted(set(non_frontier_ids) | set(remaining_frontier))
    if compact_ids or non_completed:
        compact_lines = ["### All Other Experiments"]
        for eid in compact_ids:
            if eid == most_recent_id:
                continue  # already shown
            compact_lines.append(
                _format_compact(completed[eid], primary_split, primary_key)
            )
        for eid in sorted(non_completed):
            if eid == most_recent_id:
                continue  # already shown in Most Recent section
            entry = non_completed[eid]
            status = entry.get("status", "?")
            hyp = entry.get("hypothesis", "")[:80]
            compact_lines.append(f"- **{eid}** [{status}] {hyp}")
        if len(compact_lines) > 1:
            sections.append("\n".join(compact_lines))

    return "\n\n".join(sections)


def _build_focused_context(
    index: dict,
    abreka_dir,
    primary_split: str,
    primary_key: str,
    maximize: bool,
    completed: dict,
    _sort_key,
    focus_root: str,
) -> str:
    """Build experiment context focused on a subtree."""
    subtree_ids = get_subtree_ids(index, focus_root)
    primary_metric = f"{primary_split}:{primary_key}"

    # Partition
    subtree_completed = {eid: e for eid, e in completed.items() if eid in subtree_ids}
    outside_completed = {eid: e for eid, e in completed.items() if eid not in subtree_ids}

    # Frontier within subtree
    children: dict[str, list[str]] = {}
    for eid, entry in subtree_completed.items():
        for pid in entry.get("parents", []):
            if pid in subtree_completed:
                children.setdefault(pid, []).append(eid)
    subtree_frontier_ids = {eid for eid in subtree_completed if eid not in children}

    # Best overall and best within subtree
    best_overall_id = max(completed, key=_sort_key) if completed else None
    best_subtree_id = max(subtree_completed, key=_sort_key) if subtree_completed else None

    # Most recent within subtree
    subtree_all = {eid: e for eid, e in index.items() if eid in subtree_ids}
    most_recent_id = max(
        subtree_all.keys(),
        key=lambda eid: subtree_all[eid].get("created_at", ""),
    ) if subtree_all else None

    # Sort subtree frontier by metric
    frontier_sorted = sorted(subtree_frontier_ids, key=_sort_key, reverse=True)
    top_frontier = frontier_sorted[:_FRONTIER_LIMIT]
    remaining_frontier = frontier_sorted[_FRONTIER_LIMIT:]

    # Non-frontier subtree completed experiments
    subtree_non_frontier = sorted(
        eid for eid in subtree_completed if eid not in subtree_frontier_ids
    )

    # Non-completed subtree experiments
    subtree_non_completed = {
        eid: e for eid, e in index.items()
        if eid in subtree_ids and e.get("status") != "completed"
    }

    # === Build output ===
    sections = []

    # Section 1: State summary with focus label
    total = len(index)
    n_completed = len(completed)
    n_failed = sum(1 for e in index.values() if e.get("status") == "failed")
    n_other = total - n_completed - n_failed

    state_lines = [
        f"**FOCUS:** subtree rooted at **{focus_root}** ({len(subtree_ids)} experiments)",
        f"**Total experiments:** {total} ({n_completed} completed"
    ]
    if n_failed:
        state_lines[1] += f", {n_failed} failed"
    if n_other:
        state_lines[1] += f", {n_other} in progress"
    state_lines[1] += ")"

    if best_overall_id:
        best_val = _get_primary_metric_value(completed[best_overall_id], primary_split, primary_key)
        if best_val is not None:
            state_lines.append(
                f"**Best overall:** {best_overall_id} ({primary_metric}={best_val})"
            )
    if best_subtree_id and best_subtree_id != best_overall_id:
        best_sub_val = _get_primary_metric_value(subtree_completed[best_subtree_id], primary_split, primary_key)
        if best_sub_val is not None:
            state_lines.append(
                f"**Best in subtree:** {best_subtree_id} ({primary_metric}={best_sub_val})"
            )
    sections.append("\n".join(state_lines))

    # Section 2: Most recent within subtree
    if most_recent_id:
        recent_lines = ["### Most Recent Experiment (in subtree)"]
        recent_lines.append(_format_detail(index[most_recent_id], abreka_dir))
        sections.append("\n".join(recent_lines))

    # Section 3: Subtree frontier
    if top_frontier:
        frontier_lines = [f"### Subtree Frontier — Top {len(top_frontier)} by {primary_metric}"]
        for eid in top_frontier:
            if eid == most_recent_id:
                frontier_lines.append(f"- **{eid}** (see Most Recent above)")
            else:
                frontier_lines.append(_format_detail(subtree_completed[eid], abreka_dir))
        sections.append("\n".join(frontier_lines))

    # Section 4: Remaining subtree experiments (compact)
    compact_ids = sorted(set(subtree_non_frontier) | set(remaining_frontier))
    if compact_ids or subtree_non_completed:
        compact_lines = ["### Other Subtree Experiments"]
        for eid in compact_ids:
            if eid == most_recent_id:
                continue
            compact_lines.append(
                _format_compact(subtree_completed[eid], primary_split, primary_key)
            )
        for eid in sorted(subtree_non_completed):
            if eid == most_recent_id:
                continue
            entry = subtree_non_completed[eid]
            status = entry.get("status", "?")
            hyp = entry.get("hypothesis", "")[:80]
            compact_lines.append(f"- **{eid}** [{status}] {hyp}")
        if len(compact_lines) > 1:
            sections.append("\n".join(compact_lines))

    # Section 5: Outside experiments (condensed — top N + summary)
    outside_all = {eid: e for eid, e in index.items() if eid not in subtree_ids}
    if outside_all:
        outside_lines = [f"### Outside Subtree ({len(outside_all)} experiments)"]

        if outside_completed:
            top_outside = sorted(outside_completed, key=_sort_key, reverse=True)[:_OUTSIDE_TOP_N]
            outside_lines.append(f"**Top {len(top_outside)} by {primary_metric}:**")
            for eid in top_outside:
                outside_lines.append(
                    _format_compact(outside_completed[eid], primary_split, primary_key)
                )

            # Summary of the rest
            top_set = set(top_outside)
            remaining_completed = {
                eid: e for eid, e in outside_completed.items() if eid not in top_set
            }
            parts: list[str] = []
            if remaining_completed:
                remaining_vals = [
                    v for eid in remaining_completed
                    if (v := _get_primary_metric_value(
                        remaining_completed[eid], primary_split, primary_key,
                    )) is not None
                ]
                range_str = ""
                if remaining_vals:
                    lo, hi = min(remaining_vals), max(remaining_vals)
                    range_str = f" ({primary_metric} range {lo}\u2013{hi})"
                parts.append(f"{len(remaining_completed)} completed{range_str}")

            outside_non_completed = {
                eid: e for eid, e in outside_all.items()
                if e.get("status") != "completed"
            }
            status_counts: dict[str, int] = {}
            for e in outside_non_completed.values():
                s = e.get("status", "unknown")
                status_counts[s] = status_counts.get(s, 0) + 1
            for s, count in sorted(status_counts.items()):
                parts.append(f"{count} {s}")

            n_remaining = len(outside_all) - len(top_outside)
            if n_remaining > 0 and parts:
                outside_lines.append(f"Remaining {n_remaining}: {', '.join(parts)}")
        else:
            # No completed outside experiments — just status summary
            status_counts_nc: dict[str, int] = {}
            for e in outside_all.values():
                s = e.get("status", "unknown")
                status_counts_nc[s] = status_counts_nc.get(s, 0) + 1
            summary_parts = [f"{count} {s}" for s, count in sorted(status_counts_nc.items())]
            outside_lines.append(", ".join(summary_parts))

        if len(outside_lines) > 1:
            sections.append("\n".join(outside_lines))

    return "\n\n".join(sections)


def _read_guidance_interactive() -> str:
    """Read multi-line guidance from stdin until a line containing just 'EOF'."""
    lines: list[str] = []
    for line in sys.stdin:
        if line.rstrip("\n") == "EOF":
            break
        lines.append(line.rstrip("\n"))
    return "\n".join(lines)


def _resolve_guidance(guidance: str | None, guidance_file: str | None) -> str:
    """Resolve guidance from --guidance and --guidance-file options.

    Returns the raw guidance text (not yet wrapped in a section header).
    """
    if guidance is not None and guidance_file is not None:
        raise click.ClickException("--guidance and --guidance-file are mutually exclusive.")
    if guidance_file is not None:
        from pathlib import Path
        return Path(guidance_file).read_text()
    if guidance is not None:
        return guidance
    return ""


@click.command("propose")
@click.option("--focus", "focus_root", default=None, help="Focus on subtree rooted at this experiment ID.")
@click.option(
    "--guidance",
    default=None,
    is_flag=False,
    flag_value="__interactive__",
    help="Guidance for the proposal agent. Pass a string, or use bare --guidance for interactive input.",
)
@click.option(
    "--guidance-file",
    default=None,
    type=click.Path(exists=True),
    help="Path to a file containing guidance text.",
)
@click.option(
    "--resume",
    "resume_exp_id",
    default=None,
    help="Resume an interrupted propose session for the reserved experiment ID.",
)
@click.pass_context
def propose_cmd(
    ctx,
    focus_root: str | None,
    guidance: str | None,
    guidance_file: str | None,
    resume_exp_id: str | None,
):
    """Agent analyzes experiment history and proposes the next experiment."""
    debug = ctx.obj.get("debug", False) if ctx.obj else False
    abreka_dir = find_abreka_dir()
    config = load_config(abreka_dir)
    goal = load_goal(abreka_dir)
    index = load_index(abreka_dir)

    if resume_exp_id:
        exp_id = normalize_exp_id(resume_exp_id)
        reserved = load_experiment(abreka_dir, exp_id)
        if reserved.status != ExperimentStatus.PROPOSING:
            raise click.ClickException(
                f"Experiment {exp_id} is '{reserved.status.value}', expected 'proposing'."
            )
        index = {eid: entry for eid, entry in index.items() if eid != exp_id}
    else:
        reserved = reserve_experiment(abreka_dir)
        exp_id = reserved.id

    if focus_root and focus_root not in index:
        raise click.ClickException(f"Experiment {focus_root!r} not found.")

    # Resolve guidance
    if guidance == "__interactive__":
        click.echo("Enter guidance (end with EOF on its own line):")
        guidance = _read_guidance_interactive()
    guidance_text = format_guidance(_resolve_guidance(guidance, guidance_file))

    console = Console()
    focus_label = f" [bold yellow](focus={focus_root})[/bold yellow]" if focus_root else ""
    console.print(f"[bold]Proposing next experiment {exp_id}...[/bold]{focus_label}\n")

    template = load_template(abreka_dir, "propose_exploit")
    context = _build_experiment_context(
        index, abreka_dir,
        primary_metric=config.experiment.primary_metric,
        metric_direction=config.experiment.metric_direction,
        focus_root=focus_root,
    )
    focus_constraint = _build_focus_constraint(index, focus_root) if focus_root else ""
    backlog_context = _build_backlog_context(abreka_dir, index=index, focus_root=focus_root)
    branch_overview = _build_branch_overview(
        index, abreka_dir,
        primary_metric=config.experiment.primary_metric,
        metric_direction=config.experiment.metric_direction,
        focus_root=focus_root,
    )
    focused_branch_summary = _build_focused_branch_summary(
        index, abreka_dir,
        primary_metric=config.experiment.primary_metric,
        metric_direction=config.experiment.metric_direction,
        focus_root=focus_root,
    )
    wiki_text = format_wiki_index(load_wiki_index(abreka_dir))
    prompt_text = render_template(template, {
        "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,
        "branch_overview": branch_overview,
        "focused_branch_summary": focused_branch_summary,
        "guidance": guidance_text,
        "wiki_index": wiki_text,
        "reserved_exp_id": exp_id,
    })

    session_path = stage_session_path(abreka_dir, exp_id, "propose")
    if resume_exp_id:
        if not session_path.is_file():
            raise click.ClickException(
                f"No existing propose session found for experiment {exp_id}."
            )
        message = stage_resume_message("propose", exp_id)
    else:
        message = prompt_text

    model = config.agent.propose_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,
            debug=debug,
        ) as pi:
            pi.prompt(message, on_event=stream_event)

        updated = load_experiment(abreka_dir, exp_id)
        if updated.status == ExperimentStatus.PROPOSED:
            console.print(f"\n[bold green]Created experiment {exp_id}[/bold green]")
            return
        stage_status = "error"
        if updated.status == ExperimentStatus.FAILED:
            raise click.ClickException(
                updated.failure_reason or f"Proposal {exp_id} failed."
            )
        raise click.ClickException(
            f"Propose session for experiment {exp_id} did not finalize the reserved experiment. "
            "Resume it with 'abreka propose --resume <id>'."
        )
    except Exception:
        if stage_status == "ok":
            stage_status = "error"
        raise
    finally:
        record_stage_runtime(
            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),
        )
