"""Tests for `abreka exp chat` — context assembly, argv building, prompt rendering."""

from __future__ import annotations

from pathlib import Path

from abreka.cli.exp.chat import _build_argv, _build_context, _render_system_prompt
from abreka.models import Experiment, ExperimentStatus
from abreka.persistence import create_experiment_dir, save_experiment, update_index_entry


def _make_experiment(
    abreka_dir: Path,
    exp_id: str,
    *,
    status: ExperimentStatus = ExperimentStatus.COMPLETED,
    hypothesis: str = "Bigger batch size improves throughput.",
    method: str = "Train CNN with batch size 256.",
    parents: list[str] | None = None,
    metrics: dict | None = None,
    findings: str = "",
    failure_reason: str = "",
) -> Experiment:
    exp = Experiment(
        id=exp_id,
        status=status,
        hypothesis=hypothesis,
        method=method,
        parents=parents or [],
        metrics=metrics or {},
        failure_reason=failure_reason,
    )
    create_experiment_dir(abreka_dir, exp_id)
    save_experiment(exp, abreka_dir)
    update_index_entry(abreka_dir, exp)
    if findings:
        (abreka_dir / "experiments" / exp_id / "findings.md").write_text(findings)
    return exp


class TestBuildContext:
    def test_collects_experiment_metadata(self, abreka_project):
        abreka_dir = abreka_project / "abreka"
        exp = _make_experiment(
            abreka_dir,
            "0001",
            metrics={"val": {"accuracy": 0.931}, "train": {"loss": 0.12}},
            findings="## What happened\nIt worked.\n",
        )

        ctx = _build_context(abreka_dir, exp)

        assert ctx["exp_id"] == "0001"
        assert ctx["status"] == "completed"
        assert ctx["hypothesis"] == "Bigger batch size improves throughput."
        assert ctx["method"] == "Train CNN with batch size 256."
        assert ctx["project_name"] == "test-project"
        assert ctx["lib_name"] == "test_project"
        assert ctx["primary_metric"] == "val:accuracy"
        assert "Achieve high accuracy" in ctx["goal_text"]
        assert "It worked" in ctx["findings_text"]
        assert "val" in ctx["metrics_table"]
        assert "accuracy" in ctx["metrics_table"]
        assert "0.93" in ctx["metrics_table"]

    def test_handles_missing_findings_and_metrics(self, abreka_project):
        abreka_dir = abreka_project / "abreka"
        exp = _make_experiment(abreka_dir, "0002", metrics={}, findings="")

        ctx = _build_context(abreka_dir, exp)

        assert ctx["findings_text"].strip() != ""  # placeholder, not empty string
        assert ctx["metrics_table"].strip() != ""

    def test_parent_summaries(self, abreka_project):
        abreka_dir = abreka_project / "abreka"
        _make_experiment(
            abreka_dir,
            "0001",
            hypothesis="Baseline CNN.",
            metrics={"val": {"accuracy": 0.85}},
        )
        exp = _make_experiment(
            abreka_dir,
            "0002",
            parents=["0001"],
            hypothesis="Add dropout.",
        )

        ctx = _build_context(abreka_dir, exp)

        assert "0001" in ctx["parent_summaries"]
        assert "Baseline CNN" in ctx["parent_summaries"]

    def test_failure_reason_included(self, abreka_project):
        abreka_dir = abreka_project / "abreka"
        exp = _make_experiment(
            abreka_dir,
            "0003",
            status=ExperimentStatus.FAILED,
            failure_reason="OOM on GPU.",
        )

        ctx = _build_context(abreka_dir, exp)

        assert "OOM on GPU" in ctx["failure_reason"]


class TestRenderSystemPrompt:
    def test_template_substitutes_experiment_id_and_hypothesis(self, abreka_project):
        abreka_dir = abreka_project / "abreka"
        exp = _make_experiment(
            abreka_dir,
            "0007",
            hypothesis="Swapping optimizer helps convergence.",
        )

        ctx = _build_context(abreka_dir, exp)
        rendered = _render_system_prompt(ctx)

        assert "0007" in rendered
        assert "Swapping optimizer helps convergence." in rendered
        assert "test-project" in rendered
        assert "val:accuracy" in rendered


class TestBuildArgv:
    def test_minimal_argv_has_session_dir_and_system_prompt(self, tmp_path):
        chats_dir = tmp_path / "chats"
        argv = _build_argv(
            system_prompt="SYS",
            chats_dir=chats_dir,
            continue_=False,
            resume=False,
            messages=(),
            exp_id="0001",
        )
        assert argv[0] == "pi"
        assert "--append-system-prompt" in argv
        assert "SYS" in argv
        assert "--session-dir" in argv
        assert str(chats_dir) in argv
        # Default initial message when none given
        assert any("0001" in a for a in argv[-1:])

    def test_continue_and_resume_pass_through(self, tmp_path):
        chats_dir = tmp_path / "chats"
        argv = _build_argv(
            system_prompt="SYS",
            chats_dir=chats_dir,
            continue_=True,
            resume=False,
            messages=(),
            exp_id="0001",
        )
        assert "--continue" in argv

        argv = _build_argv(
            system_prompt="SYS",
            chats_dir=chats_dir,
            continue_=False,
            resume=True,
            messages=(),
            exp_id="0001",
        )
        assert "--resume" in argv

    def test_explicit_messages_replace_default(self, tmp_path):
        chats_dir = tmp_path / "chats"
        argv = _build_argv(
            system_prompt="SYS",
            chats_dir=chats_dir,
            continue_=False,
            resume=False,
            messages=("hello", "world"),
            exp_id="0001",
        )
        assert argv[-2:] == ["hello", "world"]
