"""Tests for abreka exp subcommands."""

import json
import os

from click.testing import CliRunner

from abreka.cli import main
from abreka.persistence import load_experiment, save_experiment, update_index_entry


class TestExpNew:
    def test_creates_experiment(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, [
            "exp", "new",
            "--hypothesis", "Test CNN baseline",
            "--parent", "0000",
            "--tag", "cnn",
        ])
        assert result.exit_code == 0, result.output
        assert "0001" in result.output

        # Verify experiment dir exists
        exp_dir = abreka_project / "abreka" / "experiments" / "0001"
        assert exp_dir.is_dir()
        assert (exp_dir / "experiment.toml").is_file()
        assert (exp_dir / "code").is_dir()
        assert (exp_dir / "artifacts" / "checkpoints").is_dir()
        assert (exp_dir / "artifacts" / "plots").is_dir()
        assert (exp_dir / "artifacts" / "eda").is_dir()

    def test_sequential_ids(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        r1 = runner.invoke(main, ["exp", "new", "--hypothesis", "First"])
        r2 = runner.invoke(main, ["exp", "new", "--hypothesis", "Second"])
        assert "0001" in r1.output
        assert "0002" in r2.output


class TestExpUpdate:
    def test_update_metrics(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Test"])
        result = runner.invoke(main, [
            "exp", "update", "1",
            "--metric", "val:accuracy=0.95",
            "--metric", "train:loss=0.01",
        ])
        assert result.exit_code == 0

        # Check index
        index = json.loads((abreka_project / "abreka" / "index.json").read_text())
        assert index["0001"]["metrics"]["val"]["accuracy"] == 0.95


class TestExpFinishFail:
    def test_finish(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new"])
        result = runner.invoke(main, ["exp", "finish", "1"])
        assert result.exit_code == 0

        index = json.loads((abreka_project / "abreka" / "index.json").read_text())
        assert index["0001"]["status"] == "completed"

    def test_fail(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new"])
        result = runner.invoke(main, ["exp", "fail", "1", "--reason", "OOM error"])
        assert result.exit_code == 0


class TestExpList:
    def test_empty(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["exp", "list"])
        assert "No experiments" in result.output

    def test_with_experiments(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "First"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Second"])
        result = runner.invoke(main, ["exp", "list"])
        assert result.exit_code == 0
        assert "0001" in result.output
        assert "0002" in result.output

    def test_filter_by_status(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "First"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Second"])
        runner.invoke(main, ["exp", "finish", "1"])
        result = runner.invoke(main, ["exp", "list", "--status", "completed"])
        assert "0001" in result.output
        assert "0002" not in result.output


class TestExpSearch:
    def test_search_word_boundary(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "SAM model baseline"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "SAME thing different"])
        result = runner.invoke(main, ["exp", "search", "SAM"])
        assert result.exit_code == 0
        assert "0001" in result.output
        assert "0002" not in result.output

    def test_search_regex(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Baseline CNN model"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Transformer model"])
        result = runner.invoke(main, ["exp", "search", "/CNN|Transformer/"])
        assert result.exit_code == 0
        assert "0001" in result.output
        assert "0002" in result.output

    def test_search_regex_case_insensitive(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Baseline CNN model"])
        result = runner.invoke(main, ["exp", "search", "/baseline/i"])
        assert result.exit_code == 0
        assert "0001" in result.output

    def test_search_regex_case_sensitive_no_match(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Baseline CNN model"])
        result = runner.invoke(main, ["exp", "search", "/baseline/"])
        assert result.exit_code == 0
        assert "No experiments" in result.output

    def test_search_filter_status(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "CNN model"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "CNN variant"])
        runner.invoke(main, ["exp", "finish", "1"])
        result = runner.invoke(main, ["exp", "search", "CNN", "--status", "completed"])
        assert "0001" in result.output
        assert "0002" not in result.output

    def test_search_filter_tag(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "CNN model", "--tag", "arch"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "CNN variant", "--tag", "hyperparam"])
        result = runner.invoke(main, ["exp", "search", "CNN", "--tag", "arch"])
        assert "0001" in result.output
        assert "0002" not in result.output

    def test_search_invalid_regex(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["exp", "search", "/[unclosed/"])
        assert result.exit_code != 0 or "Invalid regex" in result.output


class TestExpShow:
    def test_show_experiment(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Test show"])
        result = runner.invoke(main, ["exp", "show", "1"])
        assert result.exit_code == 0
        assert "Test show" in result.output

    def test_show_displays_runtime_costs(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Test show"])

        exp = load_experiment(abreka_project / "abreka", "0001")
        exp.runtime_costs = {
            "propose": {
                "attempts": 1,
                "seconds_total": 1.2,
                "seconds_last": 1.2,
                "last_status": "ok",
            },
            "run": {
                "attempts": 2,
                "seconds_total": 11.0,
                "seconds_last": 6.0,
                "last_status": "timed_out",
            },
        }
        save_experiment(exp, abreka_project / "abreka")
        update_index_entry(abreka_project / "abreka", exp)

        result = runner.invoke(main, ["exp", "show", "1"])
        assert result.exit_code == 0
        assert "Runtime Costs" in result.output
        assert "propose" in result.output
        assert "timed_out" in result.output


class TestExpTree:
    def test_tree(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Root"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Child", "--parent", "0001"])
        result = runner.invoke(main, ["exp", "tree"])
        assert result.exit_code == 0
        assert "0001" in result.output
        assert "0002" in result.output


class TestExpPrune:
    def _build_chain(self, runner, abreka_project):
        """root(0001) -> child(0002) -> grandchild(0003), plus sibling(0004)."""
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Root"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Child", "--parent", "0001"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Grandchild", "--parent", "0002"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Sibling"])

    def test_prune_inclusive_default(self, abreka_project):
        runner = CliRunner()
        self._build_chain(runner, abreka_project)
        result = runner.invoke(main, ["exp", "prune", "1", "--yes"])
        assert result.exit_code == 0, result.output

        exps = abreka_project / "abreka" / "experiments"
        assert not (exps / "0001").exists()
        assert not (exps / "0002").exists()
        assert not (exps / "0003").exists()
        assert (exps / "0004").exists()

        index = json.loads((abreka_project / "abreka" / "index.json").read_text())
        assert "0001" not in index
        assert "0002" not in index
        assert "0003" not in index
        assert "0004" in index

    def test_prune_exclusive_keeps_target(self, abreka_project):
        runner = CliRunner()
        self._build_chain(runner, abreka_project)
        result = runner.invoke(main, ["exp", "prune", "1", "--exclusive", "--yes"])
        assert result.exit_code == 0, result.output

        exps = abreka_project / "abreka" / "experiments"
        assert (exps / "0001").exists()
        assert not (exps / "0002").exists()
        assert not (exps / "0003").exists()
        assert (exps / "0004").exists()

        index = json.loads((abreka_project / "abreka" / "index.json").read_text())
        assert "0001" in index
        assert "0002" not in index
        assert "0003" not in index

    def test_prune_descends_dag_not_just_tree(self, abreka_project):
        """A grandchild with two parents inside the subtree is pruned cleanly."""
        runner = CliRunner()
        os.chdir(abreka_project)
        runner.invoke(main, ["exp", "new", "--hypothesis", "Root"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Left", "--parent", "0001"])
        runner.invoke(main, ["exp", "new", "--hypothesis", "Right", "--parent", "0001"])
        runner.invoke(main, [
            "exp", "new", "--hypothesis", "Convergence",
            "--parent", "0002", "--parent", "0003",
        ])
        result = runner.invoke(main, ["exp", "prune", "1", "--yes"])
        assert result.exit_code == 0, result.output

        exps = abreka_project / "abreka" / "experiments"
        for eid in ("0001", "0002", "0003", "0004"):
            assert not (exps / eid).exists()
        index = json.loads((abreka_project / "abreka" / "index.json").read_text())
        assert index == {}

    def test_prune_cascades_backlog(self, abreka_project):
        runner = CliRunner()
        self._build_chain(runner, abreka_project)
        runner.invoke(main, [
            "backlog", "add", "--title", "Idea from child", "--source-exp", "0002",
        ])
        runner.invoke(main, [
            "backlog", "add", "--title", "Idea from grandchild", "--source-exp", "0003",
        ])
        result = runner.invoke(main, ["exp", "prune", "1", "--yes"])
        assert result.exit_code == 0, result.output

        assert not (abreka_project / "abreka" / "backlog" / "0001.toml").exists()
        assert not (abreka_project / "abreka" / "backlog" / "0002.toml").exists()
        bl_index = json.loads(
            (abreka_project / "abreka" / "backlog_index.json").read_text()
        )
        assert "0001" not in bl_index
        assert "0002" not in bl_index

    def test_prune_leaves_unrelated_backlog(self, abreka_project):
        runner = CliRunner()
        self._build_chain(runner, abreka_project)
        # Sibling 0004 is outside the subtree rooted at 0001.
        runner.invoke(main, [
            "backlog", "add", "--title", "Unrelated idea", "--source-exp", "0004",
        ])
        runner.invoke(main, ["backlog", "add", "--title", "No source exp"])
        result = runner.invoke(main, ["exp", "prune", "1", "--yes"])
        assert result.exit_code == 0, result.output

        bl_index = json.loads(
            (abreka_project / "abreka" / "backlog_index.json").read_text()
        )
        assert "0001" in bl_index
        assert "0002" in bl_index
        assert (abreka_project / "abreka" / "backlog" / "0001.toml").exists()
        assert (abreka_project / "abreka" / "backlog" / "0002.toml").exists()

    def test_prune_requires_confirmation(self, abreka_project):
        runner = CliRunner()
        self._build_chain(runner, abreka_project)
        result = runner.invoke(main, ["exp", "prune", "1"], input="n\n")
        assert result.exit_code != 0

        exps = abreka_project / "abreka" / "experiments"
        assert (exps / "0001").exists()
        assert (exps / "0002").exists()
        assert (exps / "0003").exists()

    def test_prune_unknown_id(self, abreka_project):
        runner = CliRunner()
        os.chdir(abreka_project)
        result = runner.invoke(main, ["exp", "prune", "9999", "--yes"])
        assert result.exit_code != 0
        assert "not found" in result.output.lower()
