"""SQLite index over the trial logs — a derived cache, never the source of truth.

ADR-0006: `trials.jsonl` is the record. This index is fully rebuildable from
those files, so deleting `.ablate/index.db` is never data loss; deleting a
run's `trials.jsonl` is.

It exists for one reason: **trial memoization**. Editing one section of a
config changes that variant's hash and nothing else, so a re-run should pay
only for the arms that actually changed. docs/06-cost.md puts a typical second
run at 30-60% of the first one's cost, and that is what makes iterative use
affordable rather than a thing you do once and never again.

The memo key is `(task_id, variant_hash, agent, model, replicate)`. Every part
earns its place:

- `variant_hash` over the config *text*, so renaming a file or moving it does
  not invalidate anything, and editing it invalidates exactly that arm.
- `model` and `agent`, because a reused trial from a different model would
  silently mix two experiments — the failure docs/05-statistics.md calls the
  variance floor. The model here is the one the run *requested*, not the
  versioned id the provider reports back; keying on the latter would mean the
  cache never hits and every re-run is billed in full, with nothing visibly
  wrong. The reported id is still recorded, and drives the mixed-model warning.
- `replicate`, so N=10 reuses the N=5 run's trials and pays only for the five
  new ones, rather than either re-running everything or collapsing ten
  requested trials into five distinct results.

Trials that errored are never memoized: a timeout is a property of that
attempt, not of the design point, and caching it would make a transient
failure permanent.
"""

from __future__ import annotations

import sqlite3
from collections.abc import Iterable
from contextlib import closing
from pathlib import Path

from ablate.record import TrialRecord
from ablate.runner.trial import TrialJob

SCHEMA = """
CREATE TABLE IF NOT EXISTS trials (
    task_id      TEXT NOT NULL,
    variant_hash TEXT NOT NULL,
    agent        TEXT NOT NULL,
    model        TEXT NOT NULL DEFAULT '',
    replicate    INTEGER NOT NULL,
    run_id       TEXT NOT NULL,
    payload      TEXT NOT NULL,
    PRIMARY KEY (task_id, variant_hash, agent, model, replicate)
);
CREATE INDEX IF NOT EXISTS trials_run ON trials(run_id);
"""

MemoKey = tuple[str, str, str, str, int]


def memo_key(record: TrialRecord) -> MemoKey:
    return (
        record.task_id,
        record.variant_hash,
        record.agent,
        record.requested_model or record.model or "",
        record.replicate,
    )


def job_memo_key(job: TrialJob, agent: str, model: str | None) -> MemoKey:
    """The same key, computed before a trial runs rather than after.

    Both sides must agree exactly or the cache silently never hits, so the
    variant hash comes from the one function that defines it rather than being
    recomputed here.
    """
    return (
        job.task.id,
        job.variant.content_hash,
        agent,
        model or "",
        job.replicate,
    )


class TrialIndex:
    """Content-addressed store of completed trials, keyed by design point."""

    def __init__(self, path: Path) -> None:
        self.path = path
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with closing(self._connect()) as conn:
            conn.executescript(SCHEMA)
            conn.commit()

    def _connect(self) -> sqlite3.Connection:
        return sqlite3.connect(self.path)

    def add(self, records: Iterable[TrialRecord]) -> int:
        """Index completed trials. Errored trials are skipped deliberately."""
        rows = [
            (*memo_key(r), r.run_id, r.model_dump_json())
            for r in records
            if r.error is None
        ]
        if not rows:
            return 0
        with closing(self._connect()) as conn:
            conn.executemany(
                "INSERT OR REPLACE INTO trials "
                "(task_id, variant_hash, agent, model, replicate, run_id, payload) "
                "VALUES (?, ?, ?, ?, ?, ?, ?)",
                rows,
            )
            conn.commit()
        return len(rows)

    def lookup(self, keys: Iterable[MemoKey]) -> dict[MemoKey, TrialRecord]:
        """Fetch whatever of `keys` has already been paid for."""
        keys = list(keys)
        if not keys:
            return {}

        found: dict[MemoKey, TrialRecord] = {}
        with closing(self._connect()) as conn:
            # Chunked to stay under SQLite's variable limit on large designs.
            for start in range(0, len(keys), 200):
                chunk = keys[start : start + 200]
                placeholders = ",".join(
                    ["(?, ?, ?, ?, ?)"] * len(chunk)
                )
                flat = [value for key in chunk for value in key]
                rows = conn.execute(
                    "SELECT task_id, variant_hash, agent, model, replicate, payload "
                    "FROM trials WHERE (task_id, variant_hash, agent, model, replicate) "
                    f"IN ({placeholders})",
                    flat,
                ).fetchall()
                for *key, payload in rows:
                    found[(key[0], key[1], key[2], key[3], key[4])] = (
                        TrialRecord.model_validate_json(payload)
                    )
        return found

    def count(self) -> int:
        with closing(self._connect()) as conn:
            return int(conn.execute("SELECT COUNT(*) FROM trials").fetchone()[0])

    def rebuild(self, run_dirs: Iterable[Path]) -> int:
        """Re-derive the whole index from the JSONL logs.

        The index is a cache. If it is ever lost, corrupted, or written by an
        older schema, this restores it from the actual source of truth.
        """
        from ablate.store.log import TrialLog

        with closing(self._connect()) as conn:
            conn.execute("DELETE FROM trials")
            conn.commit()

        total = 0
        for run_dir in run_dirs:
            total += self.add(TrialLog(run_dir).read_all())
        return total
