from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import pytest

from deskclaw.adapters.protocol import (
    AgentAdapter,
    AgentCapability,
    AgentMessage,
    AgentResponse,
)
from deskclaw.experiment.models import (
    AgentBlueprint,
    ExperimentConfig,
    ExperimentResult,
    RunResult,
    VariableMatrix,
)
from deskclaw.experiment.pipeline import ExperimentPipeline
from deskclaw.experiment.registry import FileRegistry, RegistryError
from deskclaw.experiment.reporter import ComparativeReporter
from deskclaw.experiment.scoring import MetadataScorer
from deskclaw.experiment.stats import (
    cohens_d,
    confidence_interval,
    pairwise_comparison,
    welch_t_test,
)
from deskclaw.experiment.template_adapter import TemplateAdapter
from deskclaw.experiment.workspace import WorkspaceBuilder


class _MockAdapter(AgentAdapter):
    def __init__(self, caps: set[AgentCapability] | None = None) -> None:
        self._caps = caps if caps is not None else set(AgentCapability)
        self._last_messages: list[AgentMessage] = []
        self._last_temperature: float = 0.0
        self._last_max_tokens: int = 4096

    def name(self) -> str:
        return "mock-experiment"

    def capabilities(self) -> set[AgentCapability]:
        return self._caps

    async def execute(
        self,
        messages: list[AgentMessage],
        *,
        max_tokens: int = 4096,
        temperature: float = 0.0,
        tools: list[dict] | None = None,
        timeout_s: int = 300,
    ) -> AgentResponse:
        self._last_messages = list(messages)
        self._last_temperature = temperature
        self._last_max_tokens = max_tokens
        msg = AgentMessage(role="assistant", content="mock experiment response")
        return AgentResponse(
            messages=list(messages) + [msg],
            final_answer="mock experiment response",
            trace={"mock": True},
            usage={"input_tokens": 100, "output_tokens": 50, "api_calls": 1},
            latency_ms=42,
        )

    async def stream(self, messages: list[AgentMessage], **kwargs: Any):
        yield AgentMessage(role="assistant", content="mock stream")


class TestAgentBlueprint:
    def test_creation(self):
        bp = AgentBlueprint(agent_name="test-agent", adapter_name="self")
        assert bp.agent_name == "test-agent"
        assert bp.adapter_name == "self"
        assert bp.temperature == 0.0
        assert bp.max_tokens == 4096

    def test_to_dict_from_dict_roundtrip(self):
        bp = AgentBlueprint(
            agent_name="agent-a",
            adapter_name="nanobot-v1",
            model_id="gpt-4",
            temperature=0.7,
            max_tokens=2048,
            system_prompt="Be helpful.",
            capabilities=frozenset(["text_generation", "tool_use"]),
            metadata=(("version", "1.0"), ("env", "test")),
        )
        d = bp.to_dict()
        bp2 = AgentBlueprint.from_dict(d)
        assert bp2.agent_name == bp.agent_name
        assert bp2.model_id == bp.model_id
        assert bp2.temperature == bp.temperature
        assert "text_generation" in bp2.capabilities

    def test_frozen(self):
        bp = AgentBlueprint(agent_name="test", adapter_name="self")
        with pytest.raises(AttributeError):
            bp.agent_name = "changed"  # type: ignore[misc]


class TestVariableMatrix:
    def test_cartesian_product(self):
        matrix = VariableMatrix.from_dict(
            {
                "temperature": [0.0, 0.5, 1.0],
                "max_tokens": [1024, 4096],
            }
        )
        expanded = matrix.expand()
        assert len(expanded) == 6
        temps = {e["temperature"] for e in expanded}
        assert temps == {0.0, 0.5, 1.0}

    def test_empty_matrix(self):
        matrix = VariableMatrix()
        expanded = matrix.expand()
        assert expanded == [{}]

    def test_single_variable(self):
        matrix = VariableMatrix.from_dict({"temperature": [0.0, 1.0]})
        expanded = matrix.expand()
        assert len(expanded) == 2

    def test_to_dict_roundtrip(self):
        data = {"temp": [0.0, 0.5], "tokens": [1024, 4096]}
        matrix = VariableMatrix.from_dict(data)
        d = matrix.to_dict()
        assert d["temp"] == [0.0, 0.5]
        assert d["tokens"] == [1024, 4096]


class TestExperimentConfig:
    def test_creation(self):
        bp = AgentBlueprint(agent_name="a", adapter_name="self")
        config = ExperimentConfig(name="test-exp", blueprints=[bp])
        assert config.name == "test-exp"
        assert config.repetitions == 1

    def test_total_runs(self):
        bp1 = AgentBlueprint(agent_name="a", adapter_name="self")
        bp2 = AgentBlueprint(agent_name="b", adapter_name="self")
        matrix = VariableMatrix.from_dict({"temperature": [0.0, 0.5]})
        config = ExperimentConfig(
            name="test", blueprints=[bp1, bp2], matrix=matrix, repetitions=3
        )
        assert config.total_runs() == 2 * 2 * 3  # 12

    def test_to_dict_from_dict_roundtrip(self):
        bp = AgentBlueprint(agent_name="a", adapter_name="self")
        config = ExperimentConfig(
            name="roundtrip-test",
            blueprints=[bp],
            matrix=VariableMatrix.from_dict({"temp": [0.0]}),
            repetitions=2,
        )
        d = config.to_dict()
        config2 = ExperimentConfig.from_dict(d)
        assert config2.name == "roundtrip-test"
        assert config2.repetitions == 2
        assert len(config2.blueprints) == 1


class TestRunResult:
    def test_creation(self):
        r = RunResult(
            blueprint_name="test",
            variables=(("temp", 0.5),),
            repetition=0,
            status="passed",
            elapsed_ms=100,
            input_tokens=50,
            output_tokens=25,
        )
        assert r.status == "passed"
        assert r.elapsed_ms == 100

    def test_to_dict(self):
        r = RunResult(
            blueprint_name="test",
            variables=(("temp", 0.5),),
            repetition=0,
            status="error",
            error="boom",
        )
        d = r.to_dict()
        assert d["status"] == "error"
        assert d["error"] == "boom"
        assert d["variables"] == {"temp": 0.5}


class TestWorkspaceBuilder:
    def test_create_and_cleanup(self, tmp_path: Path):
        with WorkspaceBuilder(base_dir=tmp_path / "workspaces") as builder:
            ws = builder.create_workspace("run-1")
            assert ws.exists()
            assert ws.name == "run-1"
            assert len(builder.workspaces) == 1
        assert not (tmp_path / "workspaces" / "run-1").exists()

    def test_create_with_source(self, tmp_path: Path):
        src = tmp_path / "source"
        src.mkdir()
        (src / "config.yaml").write_text("key: value")

        with WorkspaceBuilder(base_dir=tmp_path / "ws", source_dir=src) as builder:
            ws = builder.create_workspace("cloned")
            assert (ws / "config.yaml").exists()
            assert (ws / "config.yaml").read_text() == "key: value"

    def test_duplicate_workspace_raises(self, tmp_path: Path):
        from deskclaw.experiment.workspace import WorkspaceError

        with WorkspaceBuilder(base_dir=tmp_path / "ws") as builder:
            builder.create_workspace("dup")
            with pytest.raises(WorkspaceError, match="already exists"):
                builder.create_workspace("dup")

    def test_cleanup_single(self, tmp_path: Path):
        builder = WorkspaceBuilder(base_dir=tmp_path / "ws")
        builder.__enter__()
        ws1 = builder.create_workspace("a")
        ws2 = builder.create_workspace("b")
        builder.cleanup(ws1)
        assert not ws1.exists()
        assert ws2.exists()
        assert len(builder.workspaces) == 1
        builder.__exit__(None, None, None)


class TestTemplateAdapter:
    async def test_name(self):
        bp = AgentBlueprint(agent_name="agent-a", adapter_name="self", model_id="gpt-4")
        inner = _MockAdapter()
        adapter = TemplateAdapter(bp, inner)
        assert adapter.name() == "agent-a:gpt-4"

    async def test_name_without_model(self):
        bp = AgentBlueprint(agent_name="agent-a", adapter_name="self")
        inner = _MockAdapter()
        adapter = TemplateAdapter(bp, inner)
        assert adapter.name() == "agent-a"

    async def test_controlled_temperature(self):
        bp = AgentBlueprint(
            agent_name="hot", adapter_name="self", temperature=0.9, max_tokens=2048
        )
        inner = _MockAdapter()
        adapter = TemplateAdapter(bp, inner)

        messages = [AgentMessage(role="user", content="hello")]
        await adapter.execute(messages)

        assert inner._last_temperature == 0.9
        assert inner._last_max_tokens == 2048

    async def test_system_prompt_injection(self):
        bp = AgentBlueprint(
            agent_name="prompted",
            adapter_name="self",
            system_prompt="Always be concise.",
        )
        inner = _MockAdapter()
        adapter = TemplateAdapter(bp, inner)

        messages = [AgentMessage(role="user", content="hello")]
        await adapter.execute(messages)

        assert inner._last_messages[0].role == "system"
        assert "Always be concise" in inner._last_messages[0].content

    async def test_delegates_capabilities(self):
        bp = AgentBlueprint(agent_name="a", adapter_name="self")
        inner = _MockAdapter(caps={AgentCapability.TEXT_GENERATION})
        adapter = TemplateAdapter(bp, inner)
        assert AgentCapability.TEXT_GENERATION in adapter.capabilities()


class TestExperimentPipeline:
    def test_run_with_mock(self, tmp_path: Path):
        import deskclaw.adapters.nanoclaw  # noqa: F401

        bp = AgentBlueprint(agent_name="self-test", adapter_name="self")
        config = ExperimentConfig(
            name="pipeline-test",
            blueprints=[bp],
            repetitions=2,
            output_dir=str(tmp_path / "results"),
            timeout_s=30,
        )
        pipeline = ExperimentPipeline(config)
        result = pipeline.run()

        assert len(result.runs) == 2
        assert result.experiment_id
        assert result.started_at
        assert result.finished_at

    def test_matrix_expansion(self, tmp_path: Path):
        import deskclaw.adapters.nanoclaw  # noqa: F401

        bp = AgentBlueprint(agent_name="self-test", adapter_name="self")
        matrix = VariableMatrix.from_dict({"temperature": [0.0, 0.5]})
        config = ExperimentConfig(
            name="matrix-test",
            blueprints=[bp],
            matrix=matrix,
            repetitions=1,
            output_dir=str(tmp_path / "results"),
            timeout_s=30,
        )
        pipeline = ExperimentPipeline(config)
        result = pipeline.run()

        assert len(result.runs) == 2

    def test_saves_results_json(self, tmp_path: Path):
        import deskclaw.adapters.nanoclaw  # noqa: F401

        bp = AgentBlueprint(agent_name="self-test", adapter_name="self")
        config = ExperimentConfig(
            name="save-test",
            blueprints=[bp],
            output_dir=str(tmp_path / "results"),
        )
        pipeline = ExperimentPipeline(config)
        pipeline.run()

        results_dir = tmp_path / "results"
        json_files = list(results_dir.glob("*.json"))
        assert len(json_files) >= 1

    def test_invalid_adapter_error(self, tmp_path: Path):
        bp = AgentBlueprint(agent_name="broken", adapter_name="nonexistent-xyz")
        config = ExperimentConfig(
            name="error-test",
            blueprints=[bp],
            output_dir=str(tmp_path / "results"),
        )
        pipeline = ExperimentPipeline(config)
        result = pipeline.run()

        assert len(result.runs) == 1
        assert result.runs[0].status == "error"
        assert result.runs[0].error is not None


class TestMetadataScorer:
    def test_score_run(self):
        run = RunResult(
            blueprint_name="test",
            variables=(),
            repetition=0,
            status="passed",
            elapsed_ms=150,
            input_tokens=100,
            output_tokens=50,
            api_calls=1,
            cost_usd=0.001,
        )
        scorer = MetadataScorer()
        metrics = scorer.score_run(run)
        assert metrics.total_tokens == 150
        assert metrics.latency_ms == 150

    def test_score_experiment(self):
        runs = [
            RunResult(
                blueprint_name="agent-a",
                variables=(),
                repetition=i,
                status="passed",
                elapsed_ms=100 + i * 10,
                input_tokens=50,
                output_tokens=25,
                cost_usd=0.001,
            )
            for i in range(5)
        ]
        config = ExperimentConfig(
            name="scoring-test",
            blueprints=[AgentBlueprint(agent_name="agent-a", adapter_name="self")],
        )
        result = ExperimentResult(config=config, runs=runs)
        scorer = MetadataScorer()
        metrics = scorer.score_experiment(result)

        assert "agent-a" in metrics.blueprint_metrics
        bp_m = metrics.blueprint_metrics["agent-a"]
        assert bp_m.run_count == 5
        assert bp_m.pass_rate == 1.0
        assert bp_m.total_tokens == 375  # 5 * (50 + 25)

    def test_handles_empty_runs(self):
        config = ExperimentConfig(name="empty", blueprints=[])
        result = ExperimentResult(config=config, runs=[])
        scorer = MetadataScorer()
        metrics = scorer.score_experiment(result)
        assert metrics.total_runs == 0


class TestFileRegistry:
    def test_register_and_get(self, tmp_path: Path):
        test_file = tmp_path / "data.json"
        test_file.write_text('{"key": "value"}')

        registry = FileRegistry()
        record = registry.register(test_file, "result")
        assert record.artifact_type == "result"
        assert record.size_bytes > 0
        assert len(record.checksum) == 64

        got = registry.get(test_file)
        assert got is not None
        assert got.path == str(test_file)

    def test_list_filter(self, tmp_path: Path):
        f1 = tmp_path / "a.json"
        f1.write_text("{}")
        f2 = tmp_path / "b.json"
        f2.write_text("{}")

        registry = FileRegistry()
        registry.register(f1, "result")
        registry.register(f2, "report")

        all_records = registry.list_artifacts()
        assert len(all_records) == 2

        results_only = registry.list_artifacts(type_filter="result")
        assert len(results_only) == 1

    def test_save_load_roundtrip(self, tmp_path: Path):
        f1 = tmp_path / "a.json"
        f1.write_text("{}")

        registry = FileRegistry()
        registry.register(f1, "result", metadata={"key": "val"})
        save_path = tmp_path / "registry.json"
        registry.save(save_path)
        assert save_path.exists()

        registry2 = FileRegistry()
        loaded = registry2.load(save_path)
        assert loaded == 1
        assert registry2.count == 1

        record = registry2.get(f1)
        assert record is not None
        assert dict(record.metadata).get("key") == "val"

    def test_load_missing_raises(self, tmp_path: Path):
        registry = FileRegistry()
        with pytest.raises(RegistryError, match="not found"):
            registry.load(tmp_path / "nonexistent.json")

    def test_register_directory(self, tmp_path: Path):
        d = tmp_path / "workspace"
        d.mkdir()
        registry = FileRegistry()
        record = registry.register(d, "workspace")
        assert record.checksum == ""


class TestStatisticalAnalysis:
    def test_welch_t_test_significant(self):
        a = [1.0, 2.0, 3.0, 4.0, 5.0]
        b = [6.0, 7.0, 8.0, 9.0, 10.0]
        result = welch_t_test(a, b)
        assert result.significant
        assert result.p_value < 0.05
        assert result.t_statistic < 0

    def test_welch_t_test_not_significant(self):
        a = [1.0, 2.0, 3.0, 4.0, 5.0]
        b = [1.1, 2.1, 3.1, 4.1, 5.1]
        result = welch_t_test(a, b)
        assert not result.significant

    def test_cohens_d(self):
        a = [1.0, 2.0, 3.0, 4.0, 5.0]
        b = [6.0, 7.0, 8.0, 9.0, 10.0]
        d = cohens_d(a, b)
        assert d < -2.0  # large effect size

    def test_cohens_d_small_sample(self):
        assert cohens_d([1.0], [2.0]) == 0.0

    def test_confidence_interval(self):
        data = [10.0, 11.0, 12.0, 13.0, 14.0]
        lo, hi = confidence_interval(data)
        assert lo < 12.0
        assert hi > 12.0

    def test_confidence_interval_single_value(self):
        lo, hi = confidence_interval([5.0])
        assert lo == 5.0
        assert hi == 5.0

    def test_pairwise_comparison(self):
        groups = {
            "agent-a": [1.0, 2.0, 3.0],
            "agent-b": [4.0, 5.0, 6.0],
            "agent-c": [1.5, 2.5, 3.5],
        }
        results = pairwise_comparison(groups)
        assert len(results) == 3
        assert ("agent-a", "agent-b") in results

    def test_t_test_to_dict(self):
        result = welch_t_test([1.0, 2.0, 3.0], [4.0, 5.0, 6.0])
        d = result.to_dict()
        assert "t_statistic" in d
        assert "p_value" in d
        assert "significant" in d


class TestComparativeReporter:
    def test_generate_markdown(self):
        runs = [
            RunResult(
                blueprint_name="agent-a",
                variables=(),
                repetition=i,
                status="passed",
                elapsed_ms=100,
                input_tokens=50,
                output_tokens=25,
            )
            for i in range(3)
        ]
        config = ExperimentConfig(
            name="report-test",
            blueprints=[AgentBlueprint(agent_name="agent-a", adapter_name="self")],
        )
        result = ExperimentResult(config=config, runs=runs, experiment_id="test-123")

        scorer = MetadataScorer()
        metrics = scorer.score_experiment(result)
        reporter = ComparativeReporter(result, metrics)

        md = reporter.generate_markdown()
        assert "# Experiment Report: report-test" in md
        assert "agent-a" in md
        assert "Blueprint Comparison" in md

    def test_write_report(self, tmp_path: Path):
        runs = [
            RunResult(
                blueprint_name="a",
                variables=(),
                repetition=0,
                status="passed",
                elapsed_ms=50,
                input_tokens=10,
                output_tokens=5,
            )
        ]
        config = ExperimentConfig(
            name="write-test",
            blueprints=[AgentBlueprint(agent_name="a", adapter_name="self")],
        )
        result = ExperimentResult(config=config, runs=runs, experiment_id="wr-test")

        scorer = MetadataScorer()
        metrics = scorer.score_experiment(result)
        reporter = ComparativeReporter(result, metrics)

        report_dir = tmp_path / "reports"
        written = reporter.write_report(report_dir)
        assert len(written) >= 2
        assert any(p.suffix == ".md" for p in written)
        assert any(p.suffix == ".json" for p in written)


class TestExperimentCLI:
    def test_experiment_help(self):
        from click.testing import CliRunner

        from deskclaw.cli.main import cli

        runner = CliRunner()
        result = runner.invoke(cli, ["experiment", "--help"])
        assert result.exit_code == 0
        assert "run" in result.output
        assert "list" in result.output

    def test_experiment_run_help(self):
        from click.testing import CliRunner

        from deskclaw.cli.main import cli

        runner = CliRunner()
        result = runner.invoke(cli, ["experiment", "run", "--help"])
        assert result.exit_code == 0
        assert "--config" in result.output
        assert "--output-dir" in result.output
        assert "--repetitions" in result.output

    def test_experiment_run_with_config(self, tmp_path: Path):
        from click.testing import CliRunner

        import deskclaw.adapters.nanoclaw  # noqa: F401
        from deskclaw.cli.main import cli

        config = {
            "name": "cli-test",
            "blueprints": [{"agent_name": "self-test", "adapter_name": "self"}],
            "matrix": {},
            "repetitions": 1,
            "output_dir": str(tmp_path / "results"),
            "timeout_s": 30,
        }
        config_file = tmp_path / "experiment.json"
        config_file.write_text(json.dumps(config))

        runner = CliRunner()
        result = runner.invoke(
            cli,
            ["experiment", "run", "--config", str(config_file)],
        )
        assert result.exit_code == 0
        assert "Experiment complete" in result.output

    def test_experiment_run_missing_config(self):
        from click.testing import CliRunner

        from deskclaw.cli.main import cli

        runner = CliRunner()
        result = runner.invoke(
            cli,
            ["experiment", "run", "--config", "/nonexistent/config.json"],
        )
        assert result.exit_code != 0

    def test_experiment_list_empty(self, tmp_path: Path):
        from click.testing import CliRunner

        from deskclaw.cli.main import cli

        runner = CliRunner()
        result = runner.invoke(
            cli,
            ["experiment", "list", "--output-dir", str(tmp_path / "empty")],
        )
        assert result.exit_code == 0
        assert "No experiments found" in result.output
