"""Tests for the Experiment Memory (SQLite store)."""

import pytest

from bilgeai.memory.store import ExperimentStore, RunRecord


@pytest.fixture
def store(tmp_path):
    """Create a temporary ExperimentStore."""
    return ExperimentStore(project_dir=tmp_path)


@pytest.fixture
def sample_config():
    return {
        "model_name": "meta-llama/Llama-3-8B",
        "quantization": "4bit",
        "batch_size": 4,
        "learning_rate": 2e-4,
        "num_epochs": 3,
        "lora_rank": 16,
    }


@pytest.fixture
def sample_hardware():
    return {
        "cpu_cores": 8,
        "ram_total_gb": 32.0,
        "gpus": [{"name": "RTX 4060", "vram_gb": 8.0}],
    }


class TestExperimentStore:
    def test_create_store(self, store):
        assert store.db_path.exists()

    def test_create_run(self, store, sample_config, sample_hardware):
        record = store.create_run(
            config_dict=sample_config,
            hardware_dict=sample_hardware,
            model_name="meta-llama/Llama-3-8B",
            notes="test run",
            tags=["test"],
        )
        assert record.run_id.startswith("run_")
        assert record.status == "created"
        assert record.model_name == "meta-llama/Llama-3-8B"
        assert record.notes == "test run"
        assert record.tags == ["test"]

    def test_get_run(self, store, sample_config):
        created = store.create_run(config_dict=sample_config)
        fetched = store.get_run(created.run_id)
        assert fetched is not None
        assert fetched.run_id == created.run_id
        assert fetched.config_hash == created.config_hash

    def test_get_run_not_found(self, store):
        assert store.get_run("nonexistent") is None

    def test_update_run(self, store, sample_config):
        record = store.create_run(config_dict=sample_config)
        store.update_run(
            record.run_id,
            status="completed",
            final_loss=0.5,
            best_loss=0.4,
            total_steps=100,
            training_time_seconds=3600.0,
            output_dir="/tmp/output",
        )
        updated = store.get_run(record.run_id)
        assert updated.status == "completed"
        assert updated.final_loss == 0.5
        assert updated.best_loss == 0.4
        assert updated.total_steps == 100
        assert updated.training_time_seconds == 3600.0
        assert updated.output_dir == "/tmp/output"

    def test_update_run_partial(self, store, sample_config):
        record = store.create_run(config_dict=sample_config)
        store.update_run(record.run_id, status="running")
        updated = store.get_run(record.run_id)
        assert updated.status == "running"
        assert updated.final_loss is None

    def test_update_run_no_changes(self, store, sample_config):
        record = store.create_run(config_dict=sample_config)
        store.update_run(record.run_id)  # no-op
        updated = store.get_run(record.run_id)
        assert updated.status == "created"

    def test_list_runs(self, store, sample_config):
        store.create_run(config_dict=sample_config, model_name="model-a")
        store.create_run(config_dict=sample_config, model_name="model-b")
        store.create_run(config_dict=sample_config, model_name="model-c")

        runs = store.list_runs()
        assert len(runs) == 3
        # Most recent first
        assert runs[0].model_name == "model-c"

    def test_list_runs_with_limit(self, store, sample_config):
        for i in range(5):
            store.create_run(config_dict=sample_config, model_name=f"model-{i}")
        runs = store.list_runs(limit=2)
        assert len(runs) == 2

    def test_list_runs_filter_status(self, store, sample_config):
        r1 = store.create_run(config_dict=sample_config)
        store.create_run(config_dict=sample_config)
        store.update_run(r1.run_id, status="completed")

        completed = store.list_runs(status="completed")
        assert len(completed) == 1
        assert completed[0].run_id == r1.run_id

    def test_list_runs_filter_model(self, store, sample_config):
        store.create_run(config_dict=sample_config, model_name="llama-8B")
        store.create_run(config_dict=sample_config, model_name="mistral-7B")

        runs = store.list_runs(model_name="llama")
        assert len(runs) == 1
        assert runs[0].model_name == "llama-8B"

    def test_get_best_run(self, store, sample_config):
        r1 = store.create_run(config_dict=sample_config)
        r2 = store.create_run(config_dict=sample_config)
        store.update_run(r1.run_id, best_loss=0.5)
        store.update_run(r2.run_id, best_loss=0.3)

        best = store.get_best_run(metric="best_loss")
        assert best.run_id == r2.run_id

    def test_get_best_run_no_results(self, store):
        assert store.get_best_run() is None

    def test_get_best_run_invalid_metric(self, store):
        assert store.get_best_run(metric="invalid") is None

    def test_diff_runs_same_config(self, store, sample_config):
        r1 = store.create_run(config_dict=sample_config)
        r2 = store.create_run(config_dict=sample_config)

        diff = store.diff_runs(r1.run_id, r2.run_id)
        assert diff["same_config"] is True
        assert diff["config_diff"] == {}

    def test_diff_runs_different_config(self, store, sample_config):
        r1 = store.create_run(config_dict=sample_config)
        config2 = {**sample_config, "batch_size": 8, "learning_rate": 1e-4}
        r2 = store.create_run(config_dict=config2)

        diff = store.diff_runs(r1.run_id, r2.run_id)
        assert diff["same_config"] is False
        assert "batch_size" in diff["config_diff"]
        assert "learning_rate" in diff["config_diff"]

    def test_diff_runs_not_found(self, store, sample_config):
        r1 = store.create_run(config_dict=sample_config)
        diff = store.diff_runs(r1.run_id, "nonexistent")
        assert "error" in diff

    def test_diff_runs_results_comparison(self, store, sample_config):
        r1 = store.create_run(config_dict=sample_config)
        r2 = store.create_run(config_dict=sample_config)
        store.update_run(r1.run_id, final_loss=0.5, total_steps=100)
        store.update_run(r2.run_id, final_loss=0.3, total_steps=200)

        diff = store.diff_runs(r1.run_id, r2.run_id)
        assert "final_loss" in diff["results_diff"]
        assert "total_steps" in diff["results_diff"]


class TestConfigHash:
    def test_deterministic(self, sample_config):
        h1 = ExperimentStore.compute_config_hash(sample_config)
        h2 = ExperimentStore.compute_config_hash(sample_config)
        assert h1 == h2

    def test_different_configs(self, sample_config):
        h1 = ExperimentStore.compute_config_hash(sample_config)
        h2 = ExperimentStore.compute_config_hash({**sample_config, "batch_size": 8})
        assert h1 != h2

    def test_ignores_explanations(self, sample_config):
        h1 = ExperimentStore.compute_config_hash(sample_config)
        h2 = ExperimentStore.compute_config_hash({**sample_config, "explanations": ["x"]})
        assert h1 == h2

    def test_ignores_schema_version(self, sample_config):
        h1 = ExperimentStore.compute_config_hash(sample_config)
        h2 = ExperimentStore.compute_config_hash({**sample_config, "schema_version": "2.0"})
        assert h1 == h2


class TestRunRecord:
    def test_defaults(self):
        record = RunRecord(
            run_id="test_001",
            timestamp="2024-01-01T00:00:00",
            config_hash="abc123",
            model_name="test-model",
        )
        assert record.status == "created"
        assert record.final_loss is None
        assert record.tags == []

    def test_tags_serialization(self, store, sample_config):
        record = store.create_run(
            config_dict=sample_config,
            tags=["sft", "lora", "8bit"],
        )
        fetched = store.get_run(record.run_id)
        assert fetched.tags == ["sft", "lora", "8bit"]


class TestEnvironmentCapture:
    def test_captures_python_version(self):
        env = ExperimentStore.capture_environment()
        assert "python_version" in env
        assert "platform" in env

    def test_captures_pip_freeze(self):
        env = ExperimentStore.capture_environment()
        # pip_freeze may or may not be present depending on environment
        # but the function should not raise
        assert isinstance(env, dict)
