"""Tests for the SQLite index and trial memoization.

Memoization spends the user's money when it misses and corrupts their
comparison when it hits wrongly, so both directions are tested: what must be
reused, and what must never be.
"""

from __future__ import annotations

from pathlib import Path

import pytest

from ablate.config.variant import ConfigVariant
from ablate.record import AssertionRecord, TelemetryRecord, TrialRecord
from ablate.runner.trial import TrialJob
from ablate.store.index import TrialIndex, job_memo_key, memo_key
from ablate.store.log import TrialLog
from ablate.task import Task


def _record(
    *,
    task: str = "fix-it",
    variant_hash: str = "sha256:aaa",
    agent: str = "claude-code",
    model: str | None = "claude-sonnet-5",
    replicate: int = 1,
    passed: bool = True,
    error: str | None = None,
) -> TrialRecord:
    return TrialRecord(
        trial_id=f"{task}-{variant_hash[-3:]}-r{replicate}",
        run_id="run1",
        task_id=task,
        variant_id="A",
        variant_hash=variant_hash,
        replicate=replicate,
        agent=agent,
        model=model,
        started_at="2026-07-27T00:00:00Z",
        duration_s=1.0,
        telemetry=TelemetryRecord(input_tokens=10, turns=2, cost_usd=0.3),
        assertions=[AssertionRecord(id="a", type="exit_code", passed=passed)],
        passed=passed,
        error=error,
    )


def _task(task_id: str = "fix-it") -> Task:
    return Task(
        id=task_id,
        prompt="fix the bug",
        base_commit="HEAD",
        assertions=[{"id": "tests-pass", "type": "exit_code", "cmd": "true"}],
    )


@pytest.fixture
def index(tmp_path: Path) -> TrialIndex:
    return TrialIndex(tmp_path / ".ablate" / "index.db")


# --- round trip -------------------------------------------------------------


def test_indexed_trial_is_found_again(index: TrialIndex):
    record = _record()
    assert index.add([record]) == 1

    found = index.lookup([memo_key(record)])
    assert list(found) == [memo_key(record)]
    assert found[memo_key(record)].trial_id == record.trial_id


def test_lookup_of_an_unseen_design_point_is_empty(index: TrialIndex):
    index.add([_record()])
    assert index.lookup([("other-task", "sha256:aaa", "claude-code", "claude-sonnet-5", 1)]) == {}


def test_reindexing_the_same_point_does_not_duplicate(index: TrialIndex):
    index.add([_record()])
    index.add([_record()])
    assert index.count() == 1


def test_empty_lookup_short_circuits(index: TrialIndex):
    assert index.lookup([]) == {}


def test_lookup_handles_more_keys_than_sqlite_variable_limit(index: TrialIndex):
    """Chunking is load-bearing on a large ablation design."""
    records = [_record(task=f"task-{i:04d}") for i in range(500)]
    index.add(records)
    found = index.lookup(memo_key(r) for r in records)
    assert len(found) == 500


# --- what must never be reused ---------------------------------------------


def test_editing_the_config_invalidates_only_that_arm(index: TrialIndex):
    index.add([_record(variant_hash="sha256:aaa")])
    assert index.lookup([("fix-it", "sha256:bbb", "claude-code", "claude-sonnet-5", 1)]) == {}


def test_a_different_model_is_a_different_design_point(index: TrialIndex):
    """Reusing a Haiku trial in a Sonnet run would silently mix two
    experiments — the variance floor in docs/05-statistics.md."""
    index.add([_record(model="claude-haiku-4-5-20251001")])
    assert index.lookup([("fix-it", "sha256:aaa", "claude-code", "claude-sonnet-5", 1)]) == {}


def test_a_different_agent_is_a_different_design_point(index: TrialIndex):
    index.add([_record(agent="claude-code")])
    assert index.lookup([("fix-it", "sha256:aaa", "codex", "claude-sonnet-5", 1)]) == {}


def test_each_replicate_is_its_own_entry(index: TrialIndex):
    """Raising --trials 5 to 10 reuses five and buys five."""
    index.add([_record(replicate=r) for r in range(1, 6)])
    keys = [("fix-it", "sha256:aaa", "claude-code", "claude-sonnet-5", r) for r in range(1, 11)]
    assert len(index.lookup(keys)) == 5


def test_errored_trials_are_never_memoized(index: TrialIndex):
    """A timeout is a property of that attempt, not of the design point.
    Caching it would make a transient failure permanent."""
    assert index.add([_record(error="agent timed out")]) == 0
    assert index.count() == 0


# --- key agreement ----------------------------------------------------------


def test_job_key_matches_the_key_of_the_trial_it_produces():
    """If these two ever disagree the cache silently never hits and every
    re-run is billed in full, with nothing visibly broken."""
    variant = ConfigVariant(id="A", text="# Conventions\n", source="CLAUDE.md")
    job = TrialJob(task=_task(), variant=variant, replicate=3)

    produced = _record(
        task="fix-it", variant_hash=variant.content_hash, replicate=3, model="claude-sonnet-5"
    )
    assert job_memo_key(job, "claude-code", "claude-sonnet-5") == memo_key(produced)


def test_variant_hash_follows_content_not_path():
    same_text = "# Conventions\n"
    assert (
        ConfigVariant("A", same_text, "CLAUDE.md").content_hash
        == ConfigVariant("A", same_text, "docs/AGENTS.md").content_hash
    )


def test_variant_hash_changes_when_the_text_changes():
    assert (
        ConfigVariant("A", "# One\n", "CLAUDE.md").content_hash
        != ConfigVariant("A", "# Two\n", "CLAUDE.md").content_hash
    )


def test_no_config_hashes_consistently():
    assert ConfigVariant("B", None, "/dev/null").content_hash == (
        ConfigVariant("B", "", "/dev/null").content_hash
    )


# --- rebuild ----------------------------------------------------------------


def test_index_rebuilds_from_the_jsonl_logs(tmp_path: Path):
    """The index is a derived cache (ADR-0006). Losing it is never data loss."""
    run_dir = tmp_path / "runs" / "run1"
    log = TrialLog(run_dir)
    for r in range(1, 4):
        log.append(_record(replicate=r))

    index = TrialIndex(tmp_path / "index.db")
    assert index.rebuild([run_dir]) == 3
    assert index.count() == 3


def test_rebuild_discards_stale_entries(tmp_path: Path):
    index = TrialIndex(tmp_path / "index.db")
    index.add([_record(task="gone")])

    run_dir = tmp_path / "runs" / "run1"
    TrialLog(run_dir).append(_record(task="kept"))

    index.rebuild([run_dir])
    assert index.count() == 1
    assert index.lookup([("gone", "sha256:aaa", "claude-code", "claude-sonnet-5", 1)]) == {}


# --- end to end -------------------------------------------------------------


def test_memoization_round_trips_through_a_real_run(git_repo: Path, tmp_path: Path):
    """The key must agree across an actual `run_jobs` -> index -> lookup cycle.

    The unit tests above check the key function against itself; this checks it
    against a trial that a runner really produced, which is where a mismatch
    would actually cost money.
    """
    from ablate.runner.schedule import build_jobs, run_jobs
    from tests.test_trial_e2e import FakeAdapter
    from tests.test_trial_e2e import _task as _e2e_task

    variants = [ConfigVariant(id="A", text="# conventions\n", source="CLAUDE.md")]
    jobs = build_jobs([_e2e_task()], variants, trials=2, seed=1)

    records, _ = run_jobs(
        jobs=jobs,
        repo=git_repo,
        worktrees_root=tmp_path / "worktrees",
        adapter=FakeAdapter(),
        run_id="run1",
        concurrency=1,
        model="claude-sonnet-5",
    )

    index = TrialIndex(tmp_path / "index.db")
    assert index.add(records) == len(records)

    hits = index.lookup(job_memo_key(job, "fake", "claude-sonnet-5") for job in jobs)
    assert len(hits) == len(jobs), "every job should hit the cache it just populated"


def test_reported_model_does_not_break_the_cache_key(git_repo: Path, tmp_path: Path):
    """Providers report versioned ids the caller never typed. Keying on those
    would make the cache silently useless."""
    from ablate.agents.base import TrialTelemetry
    from ablate.runner.schedule import build_jobs, run_jobs
    from tests.test_trial_e2e import FakeAdapter
    from tests.test_trial_e2e import _task as _e2e_task

    class VersionReportingAdapter(FakeAdapter):
        def invoke(self, spec):
            telemetry = super().invoke(spec)
            return TrialTelemetry(
                **{**telemetry.__dict__, "model": "claude-sonnet-5-20260101"}
            )

    variants = [ConfigVariant(id="A", text="# conventions\n", source="CLAUDE.md")]
    jobs = build_jobs([_e2e_task()], variants, trials=1, seed=1)

    records, _ = run_jobs(
        jobs=jobs, repo=git_repo, worktrees_root=tmp_path / "worktrees",
        adapter=VersionReportingAdapter(), run_id="run1", concurrency=1,
        model="claude-sonnet-5",
    )
    assert records[0].model == "claude-sonnet-5-20260101"
    assert records[0].requested_model == "claude-sonnet-5"

    index = TrialIndex(tmp_path / "index.db")
    index.add(records)
    hits = index.lookup(job_memo_key(job, "fake", "claude-sonnet-5") for job in jobs)
    assert len(hits) == 1
