"""Tests for abreka.branch — branch summary computation and rendering."""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from abreka.branch import (
    BranchSummary,
    compute_branch_summary,
    get_subtree_ids,
    render_branch_overview,
    render_branch_summary,
)


def _make_exp(
    eid: str,
    *,
    status: str = "completed",
    hypothesis: str = "",
    method: str = "",
    parents: list[str] | None = None,
    tags: list[str] | None = None,
    metrics: dict | None = None,
    created_at: str = "2026-04-08T10:00:00+00:00",
) -> dict:
    entry: dict = {
        "id": eid,
        "status": status,
        "hypothesis": hypothesis,
        "method": method,
        "parents": parents or [],
        "tags": tags or [],
        "created_at": created_at,
    }
    if metrics is not None:
        entry["metrics"] = metrics
    return entry


def _metric(v: float) -> dict:
    return {"val": {"accuracy": v}}


@pytest.fixture
def abreka_dir(tmp_path: Path) -> Path:
    d = tmp_path / "abreka"
    (d / "experiments").mkdir(parents=True)
    return d


class TestGetSubtreeIds:
    def test_single_node(self):
        index = {"0001": _make_exp("0001")}
        assert get_subtree_ids(index, "0001") == {"0001"}

    def test_linear_chain(self):
        index = {
            "0001": _make_exp("0001"),
            "0002": _make_exp("0002", parents=["0001"]),
            "0003": _make_exp("0003", parents=["0002"]),
        }
        assert get_subtree_ids(index, "0001") == {"0001", "0002", "0003"}
        assert get_subtree_ids(index, "0002") == {"0002", "0003"}

    def test_sibling_branch_excluded(self):
        index = {
            "0001": _make_exp("0001"),
            "0002": _make_exp("0002", parents=["0001"]),
            "0003": _make_exp("0003"),  # unrelated root
            "0004": _make_exp("0004", parents=["0003"]),
        }
        assert get_subtree_ids(index, "0001") == {"0001", "0002"}
        assert get_subtree_ids(index, "0003") == {"0003", "0004"}

    def test_unknown_root_raises(self):
        with pytest.raises(ValueError):
            get_subtree_ids({}, "0001")


class TestBranchSummaryBasics:
    def test_singleton(self, abreka_dir):
        index = {"0001": _make_exp("0001", metrics=_metric(0.9))}
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.is_singleton is True
        assert s.subtree_size == 1
        assert s.depth_max == 0
        assert s.n_completed == 1
        assert s.best_exp_id == "0001"
        assert s.best_metric_value == pytest.approx(0.9)
        assert s.health_status == "underexplored"

    def test_depth_computation(self, abreka_dir):
        index = {
            "0001": _make_exp("0001"),
            "0002": _make_exp("0002", parents=["0001"]),
            "0003": _make_exp("0003", parents=["0002"]),
            "0004": _make_exp("0004", parents=["0003"]),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.depth_max == 3

    def test_frontier_excludes_nodes_with_completed_children(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.9)),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.91)),
            "0003": _make_exp("0003", parents=["0001"], metrics=_metric(0.92)),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert set(s.frontier_ids) == {"0002", "0003"}
        assert "0001" not in s.frontier_ids

    def test_active_running_excluded_from_completed(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.9)),
            "0002": _make_exp("0002", status="running", parents=["0001"]),
            "0003": _make_exp("0003", status="failed", parents=["0001"]),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.n_completed == 1
        assert s.n_active == 1
        assert s.n_failed == 1
        assert s.active_ids == ["0002"]
        assert s.failed_ids == ["0003"]


class TestMetricAggregation:
    def test_maximize_best(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.80)),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.90)),
            "0003": _make_exp("0003", parents=["0001"], metrics=_metric(0.85)),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.best_exp_id == "0002"
        assert s.best_metric_value == pytest.approx(0.90)
        assert s.net_gain_from_root == pytest.approx(0.10)

    def test_minimize_best(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics={"val": {"loss": 0.5}}),
            "0002": _make_exp("0002", parents=["0001"], metrics={"val": {"loss": 0.3}}),
            "0003": _make_exp("0003", parents=["0001"], metrics={"val": {"loss": 0.4}}),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:loss", "minimize")
        assert s.best_exp_id == "0002"
        assert s.best_metric_value == pytest.approx(0.3)
        # Net gain from root: root 0.5 → best 0.3 is improvement of 0.2
        assert s.net_gain_from_root == pytest.approx(0.2)

    def test_gap_to_global_best(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.8)),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.85)),
            # Separate root with a higher value
            "0100": _make_exp("0100", metrics=_metric(0.95)),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.global_best_id == "0100"
        assert s.global_best_metric_value == pytest.approx(0.95)
        assert s.gap_to_global_best == pytest.approx(0.10)
        assert s.is_global_best_in_subtree is False

    def test_global_best_in_subtree(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.8)),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.95)),
            "0100": _make_exp("0100", metrics=_metric(0.85)),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.is_global_best_in_subtree is True
        assert s.gap_to_global_best == pytest.approx(0.0)


class TestTrendAndHealth:
    def _improving_branch(self) -> dict:
        return {
            "0001": _make_exp("0001", metrics=_metric(0.80), created_at="2026-04-01T00:00:00+00:00"),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.82), created_at="2026-04-02T00:00:00+00:00"),
            "0003": _make_exp("0003", parents=["0002"], metrics=_metric(0.84), created_at="2026-04-03T00:00:00+00:00"),
            "0004": _make_exp("0004", parents=["0003"], metrics=_metric(0.86), created_at="2026-04-04T00:00:00+00:00"),
            "0005": _make_exp("0005", parents=["0004"], metrics=_metric(0.88), created_at="2026-04-05T00:00:00+00:00"),
        }

    def test_improving_classification(self, abreka_dir):
        s = compute_branch_summary(self._improving_branch(), abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.health_status == "improving"
        assert s.avg_recent_delta is not None and s.avg_recent_delta > 0
        assert s.new_bests_in_last_n >= 1
        assert s.recommended_next_action in ("exploit", "discriminate")

    def test_plateauing_classification(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.90), created_at="2026-04-01T00:00:00+00:00"),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.9005), created_at="2026-04-02T00:00:00+00:00"),
            "0003": _make_exp("0003", parents=["0002"], metrics=_metric(0.9003), created_at="2026-04-03T00:00:00+00:00"),
            "0004": _make_exp("0004", parents=["0003"], metrics=_metric(0.9004), created_at="2026-04-04T00:00:00+00:00"),
            "0005": _make_exp("0005", parents=["0004"], metrics=_metric(0.9003), created_at="2026-04-05T00:00:00+00:00"),
            # Peak is at 0002 (0.9005); 0003-0005 all sit below it
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.health_status == "plateauing"
        assert s.recommended_next_action == "explore_elsewhere"

    def test_declining_classification(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.95), created_at="2026-04-01T00:00:00+00:00"),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.93), created_at="2026-04-02T00:00:00+00:00"),
            "0003": _make_exp("0003", parents=["0002"], metrics=_metric(0.91), created_at="2026-04-03T00:00:00+00:00"),
            "0004": _make_exp("0004", parents=["0003"], metrics=_metric(0.89), created_at="2026-04-04T00:00:00+00:00"),
            "0005": _make_exp("0005", parents=["0004"], metrics=_metric(0.87), created_at="2026-04-05T00:00:00+00:00"),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.health_status == "declining"
        assert s.recommended_next_action == "stop_or_pause"

    def test_underexplored_classification(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.9)),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.91)),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.health_status == "underexplored"
        assert s.recommended_next_action == "validate"

    def test_stalled_all_failed(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", status="failed"),
            "0002": _make_exp("0002", status="failed", parents=["0001"]),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.health_status == "stalled"
        assert s.recommended_next_action == "stop_or_pause"


class TestFindingsAndBacklog:
    def test_findings_snippets_are_truncated(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.9), created_at="2026-04-01T00:00:00+00:00"),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.91), created_at="2026-04-02T00:00:00+00:00"),
            "0003": _make_exp("0003", parents=["0002"], metrics=_metric(0.92), created_at="2026-04-03T00:00:00+00:00"),
        }
        long_text = "X" * 1000
        exp_dir = abreka_dir / "experiments" / "0003"
        exp_dir.mkdir(parents=True)
        (exp_dir / "findings.md").write_text(long_text)

        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.recent_findings_snippets
        assert "..." in s.recent_findings_snippets[0]
        assert len(s.recent_findings_snippets[0]) < 500

    def test_backlog_linkage_filters_by_source_exp(self, abreka_dir, tmp_path):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.9)),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.91)),
        }
        backlog = {
            "0001": {
                "id": "0001",
                "title": "In-branch idea",
                "description": "d",
                "tags": [],
                "priority": 5,
                "source_exp": "0002",
                "resolved_exp": "",
                "status": "open",
                "created_at": "2026-04-01T00:00:00+00:00",
                "updated_at": "2026-04-01T00:00:00+00:00",
            },
            "0002": {
                "id": "0002",
                "title": "Outside idea",
                "description": "d",
                "tags": [],
                "priority": 5,
                "source_exp": "9999",
                "resolved_exp": "",
                "status": "open",
                "created_at": "2026-04-01T00:00:00+00:00",
                "updated_at": "2026-04-01T00:00:00+00:00",
            },
        }
        (abreka_dir / "backlog_index.json").write_text(json.dumps(backlog))
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        assert s.open_backlog_count == 1
        assert any("In-branch idea" in item for item in s.top_backlog_items)


class TestSerialization:
    def test_to_dict_is_json_serializable(self, abreka_dir):
        index = {"0001": _make_exp("0001", metrics=_metric(0.9))}
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        blob = json.dumps(s.to_dict())
        reloaded = json.loads(blob)
        assert reloaded["root_id"] == "0001"
        assert reloaded["health_status"] == "underexplored"


class TestRendering:
    def test_compact_shows_key_fields(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.9)),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.92)),
            "0003": _make_exp("0003", parents=["0002"], metrics=_metric(0.94)),
            "0004": _make_exp("0004", parents=["0003"], metrics=_metric(0.96)),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        out = render_branch_summary(s, mode="compact")
        assert "Branch 0001" in out
        assert "Health:" in out
        assert "Suggested next action" in out

    def test_default_render_has_performance_and_health(self, abreka_dir):
        index = {
            "0001": _make_exp("0001", metrics=_metric(0.9)),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.92)),
            "0003": _make_exp("0003", parents=["0002"], metrics=_metric(0.94)),
            "0004": _make_exp("0004", parents=["0003"], metrics=_metric(0.96)),
        }
        s = compute_branch_summary(index, abreka_dir, "0001", "val:accuracy", "maximize")
        out = render_branch_summary(s)
        assert "Performance" in out
        assert "Health" in out
        assert "Frontier" in out
        assert "0004" in out  # frontier contains the newest leaf

    def test_overview_empty(self):
        out = render_branch_overview([])
        assert "No branches" in out

    def test_overview_sorts_by_attention_by_default(self, abreka_dir):
        # Improving high-attention branch
        idx_a = {
            "0001": _make_exp("0001", metrics=_metric(0.80), created_at="2026-04-01T00:00:00+00:00"),
            "0002": _make_exp("0002", parents=["0001"], metrics=_metric(0.82), created_at="2026-04-02T00:00:00+00:00"),
            "0003": _make_exp("0003", parents=["0002"], metrics=_metric(0.88), created_at="2026-04-03T00:00:00+00:00"),
            "0004": _make_exp("0004", parents=["0003"], metrics=_metric(0.90), created_at="2026-04-04T00:00:00+00:00"),
        }
        s_improving = compute_branch_summary(idx_a, abreka_dir, "0001", "val:accuracy", "maximize")

        # Underexplored medium-attention branch
        idx_b = {"0100": _make_exp("0100", metrics=_metric(0.75))}
        s_under = compute_branch_summary(idx_b, abreka_dir, "0100", "val:accuracy", "maximize")

        out = render_branch_overview([s_under, s_improving])
        a = out.index("Branch 0001")
        b = out.index("Branch 0100")
        assert a < b  # improving/high attention comes first
