"""abreka exp tree — show the experiment DAG."""

from __future__ import annotations

import click
from rich.console import Console
from rich.tree import Tree

from abreka.config import find_abreka_dir, load_config, load_index
from abreka.metrics import parse_metric_ref


@click.command("tree")
def tree_cmd():
    """Show the experiment DAG as a tree."""
    abreka_dir = find_abreka_dir()
    config = load_config(abreka_dir)
    index = load_index(abreka_dir)

    if not index:
        click.echo("No experiments yet.")
        return

    primary_split, primary_key = parse_metric_ref(config.experiment.primary_metric)

    # Build parent -> children mapping
    children: dict[str, list[str]] = {}
    roots: list[str] = []
    for eid, entry in sorted(index.items()):
        parents = entry.get("parents", [])
        if not parents:
            roots.append(eid)
        for pid in parents:
            if pid not in children:
                children[pid] = []
            children[pid].append(eid)

    def _label(eid: str) -> str:
        entry = index.get(eid, {})
        status = entry.get("status", "?")
        hyp = entry.get("hypothesis", "")[:40]
        metric_val = entry.get("metrics", {}).get(primary_split, {}).get(primary_key)
        metric_str = f" ({primary_split}:{primary_key}={metric_val:.4f})" if metric_val is not None else ""

        status_styles = {
            "completed": "green",
            "running": "yellow",
            "proposed": "cyan",
            "failed": "red",
        }
        style = status_styles.get(status, "")
        return f"[bold]{eid}[/bold] [{style}]{status}[/{style}] {hyp}{metric_str}"

    def _add_children(tree_node: Tree, parent_id: str, visited: set[str]):
        for child_id in children.get(parent_id, []):
            if child_id in visited:
                tree_node.add(f"[dim]{child_id} (see above)[/dim]")
                continue
            visited.add(child_id)
            child_node = tree_node.add(_label(child_id))
            _add_children(child_node, child_id, visited)

    console = Console()
    tree = Tree("[bold]Experiments[/bold]")
    visited: set[str] = set()

    for root_id in roots:
        visited.add(root_id)
        root_node = tree.add(_label(root_id))
        _add_children(root_node, root_id, visited)

    # Any experiments not reachable from roots (orphans with missing parents)
    for eid in sorted(index.keys()):
        if eid not in visited:
            visited.add(eid)
            node = tree.add(_label(eid))
            _add_children(node, eid, visited)

    console.print(tree)
