"""Integration tests for the strategy A 'all_at_once' executor (#000065).

Builds a tiny 2-shard test corpus, extracts audit events, plans, and
executes a reshard into a 4-shard target layout. Verifies:

  * every document survives
  * every document lands in the shard its document_root hashes to
  * the FTS5 indexes serve search queries on the target
  * the consolidated audit chain re-verifies on the target
  * the reshard's own audit event appears at the tail of that chain
  * every target shard has corpus_shard_count meta = M
"""
from __future__ import annotations

import hashlib
import json
import sqlite3
from pathlib import Path
from typing import Iterator

import pytest

from arborist.document import Document, shard_for_document
from arborist.ingest import ingest_source
from arborist.migrate import (
    HydrationPlanner,
    execute_plan,
    _table_exists,
)
from arborist.source import Source
from arborist.store import (
    connect,
    get_corpus_shard_count,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


class _FakeSource(Source):
    source_type = "html"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _make_doc(uri: str, seed: str) -> Document:
    # Long enough to chunk into 2+ chunks. ``seed`` varies content so
    # document_root is distinct per doc (the Merkle root keys on
    # CONTENT, not URI — identical content dedups in the cache).
    content = (
        f"This is the document at {uri} with content seed {seed}. " * 10
        + f"Body text {seed} with enough tokens for proper chunking. " * 30
    )
    return Document(
        uri=uri, content=content, source_type="html", title=uri.split("/")[-1]
    )


@pytest.fixture
def source_shards(tmp_path: Path) -> Path:
    """Two source shards, ~15 docs each."""
    src = tmp_path / "src"
    src.mkdir()
    docs_per_shard = 15
    for shard_idx in range(2):
        db = src / f"{shard_idx:03d}.db"
        conn = connect(db)
        try:
            docs = [
                _make_doc(
                    f"https://example.com/shard{shard_idx}/doc-{i}",
                    seed=f"s{shard_idx}d{i}",
                )
                for i in range(docs_per_shard)
            ]
            ingest_source(conn, _FakeSource(docs))
        finally:
            conn.close()
    return src


@pytest.fixture
def audit_ndjson(source_shards: Path, tmp_path: Path) -> Path:
    """Extract the source-shards' audit chain to ndjson (what
    bench/extract_audit_events.py would produce on the real host)."""
    out = tmp_path / "audit-events.ndjson"
    with out.open("w") as f:
        for shard in sorted(source_shards.glob("00[0-9].db")):
            idx = int(shard.stem)
            c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True)
            try:
                cur = c.execute(
                    "SELECT seq, event_hash, prev_event_hash, event_type, "
                    "subject_root, body, ts FROM audit_events ORDER BY seq ASC"
                )
                for seq, eh, peh, etype, subj, body, ts in cur:
                    f.write(json.dumps({
                        "src_shard": idx,
                        "src_seq": seq,
                        "src_event_hash": eh,
                        "prev_event_hash": peh,
                        "event_type": etype,
                        "subject_root": subj,
                        "body": body,
                        "ts": ts,
                    }, separators=(",", ":")))
                    f.write("\n")
            finally:
                c.close()
    return out


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestMigrationAllAtOnce:
    @pytest.fixture
    def migrated(self, source_shards: Path, audit_ndjson: Path, tmp_path: Path):
        target = tmp_path / "tgt"
        plan = HydrationPlanner().plan(
            source_dir=source_shards,
            target_dir=target,
            target_M=4,
            free_bytes=200 * 1024 ** 3,  # plenty
        )
        assert plan.strategy == "all_at_once"
        result = execute_plan(
            plan, audit_events_ndjson=audit_ndjson,
        )
        return target, plan, result

    def test_target_dir_has_M_shards(self, migrated):
        target, plan, _ = migrated
        assert sorted(p.name for p in target.glob("*.db")) == [
            "000.db", "001.db", "002.db", "003.db"
        ]

    def test_every_source_doc_lands_in_correct_shard(
        self, source_shards: Path, migrated
    ):
        target, plan, _ = migrated
        # Build the expected mapping from source shards.
        expected: dict[str, int] = {}
        for shard in sorted(source_shards.glob("*.db")):
            c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True)
            try:
                cur = c.execute("SELECT document_root FROM documents")
                for (root,) in cur:
                    expected[root] = shard_for_document(root, plan.target_M)
            finally:
                c.close()

        # Verify each target shard holds exactly the docs hashed to it.
        for tidx in range(plan.target_M):
            tdb = target / f"{tidx:03d}.db"
            c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
            try:
                actual = {
                    r[0] for r in c.execute("SELECT document_root FROM documents")
                }
            finally:
                c.close()
            for root, expect_idx in expected.items():
                if expect_idx == tidx:
                    assert root in actual, (
                        f"doc {root[:8]}... expected on shard {tidx} but missing"
                    )
                else:
                    assert root not in actual, (
                        f"doc {root[:8]}... on shard {tidx} but hashes to {expect_idx}"
                    )

    def test_every_chunk_followed_its_document(self, migrated):
        target, plan, _ = migrated
        # For each target shard, every chunk's document_root must route here.
        for tidx in range(plan.target_M):
            tdb = target / f"{tidx:03d}.db"
            c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
            try:
                cur = c.execute(
                    "SELECT DISTINCT document_root FROM chunks"
                )
                for (root,) in cur:
                    assert shard_for_document(root, plan.target_M) == tidx, (
                        f"chunk doc {root[:8]}... on shard {tidx} but routes "
                        f"to {shard_for_document(root, plan.target_M)}"
                    )
            finally:
                c.close()

    def test_audit_chain_valid_on_target(self, migrated):
        target, plan, result = migrated
        tdb = target / "000.db"
        c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
        try:
            cur = c.execute(
                "SELECT seq, event_hash, prev_event_hash, body "
                "FROM audit_events ORDER BY seq ASC"
            )
            prev = None
            for seq, event_hash, peh, body in cur:
                assert peh == prev, f"chain break at seq {seq}"
                h = hashlib.sha256()
                if prev is not None:
                    h.update(bytes.fromhex(prev))
                h.update(body.encode("utf-8", errors="surrogatepass"))
                assert h.hexdigest() == event_hash, (
                    f"hash mismatch at seq {seq}"
                )
                prev = event_hash
        finally:
            c.close()

    def test_reshard_event_at_tail(self, migrated):
        target, plan, result = migrated
        tdb = target / "000.db"
        c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
        try:
            row = c.execute(
                "SELECT event_type, body FROM audit_events "
                "ORDER BY seq DESC LIMIT 1"
            ).fetchone()
        finally:
            c.close()
        assert row[0] == "reshard"
        body = json.loads(row[1])
        assert body["kind"] == "reshard"
        assert body["plan"]["strategy"] == "all_at_once"
        assert body["plan"]["target_M"] == 4
        assert body["result"]["rows_moved_total"] > 0

    def test_corpus_shard_count_set_on_every_target(self, migrated):
        target, plan, _ = migrated
        for tidx in range(plan.target_M):
            tdb = target / f"{tidx:03d}.db"
            c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
            try:
                assert get_corpus_shard_count(c) == plan.target_M, (
                    f"shard {tidx} missing corpus_shard_count meta"
                )
            finally:
                c.close()

    def test_chunks_fts_searchable_on_target(self, migrated):
        target, plan, _ = migrated
        # Verify FTS5 index serves a hit on every target shard that has
        # documents — the test content uses the word "document" liberally.
        hit_count = 0
        for tidx in range(plan.target_M):
            tdb = target / f"{tidx:03d}.db"
            c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
            try:
                row = c.execute(
                    "SELECT COUNT(*) FROM chunks_fts WHERE chunks_fts MATCH ?",
                    ("document",),
                ).fetchone()
                hit_count += int(row[0])
            finally:
                c.close()
        assert hit_count > 0, "no chunks_fts hits across any target shard"

    def test_total_doc_count_preserved(
        self, source_shards: Path, migrated
    ):
        target, _, _ = migrated
        # 2 source × 15 docs = 30 unique docs.
        src_total = 0
        for shard in source_shards.glob("*.db"):
            c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True)
            try:
                src_total += c.execute(
                    "SELECT COUNT(*) FROM documents"
                ).fetchone()[0]
            finally:
                c.close()
        tgt_total = 0
        for shard in target.glob("*.db"):
            c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True)
            try:
                tgt_total += c.execute(
                    "SELECT COUNT(*) FROM documents"
                ).fetchone()[0]
            finally:
                c.close()
        assert tgt_total == src_total, (
            f"doc count drift: src={src_total} tgt={tgt_total}"
        )

    def test_dry_run_creates_no_targets(
        self, source_shards: Path, audit_ndjson: Path, tmp_path: Path
    ):
        target = tmp_path / "tgt-dry"
        plan = HydrationPlanner().plan(
            source_dir=source_shards,
            target_dir=target,
            target_M=4,
            free_bytes=200 * 1024 ** 3,
        )
        result = execute_plan(
            plan, audit_events_ndjson=audit_ndjson, dry_run=True
        )
        assert "DRY RUN" in " ".join(result.notes)
        # target dir exists but no DBs.
        assert list(target.glob("*.db")) == []


class TestInPlacePromotion:
    """Verify the .db.new → .db atomic-promote dance (α-shape)."""

    def test_target_dir_only_has_db_files_after_completion(
        self, source_shards: Path, audit_ndjson: Path, tmp_path: Path
    ):
        """No .new sidecars left after a successful run."""
        target = tmp_path / "tgt"
        plan = HydrationPlanner().plan(
            source_dir=source_shards,
            target_dir=target,
            target_M=4,
            free_bytes=200 * 1024 ** 3,
        )
        execute_plan(plan, audit_events_ndjson=audit_ndjson)
        files = sorted(p.name for p in target.glob("*"))
        # 4 final .db files; -wal/-shm sidecars after close are
        # permissible but no .db.new files.
        assert all(not f.endswith(".db.new") for f in files), files
        assert sum(1 for f in files if f.endswith(".db")) == 4

    def test_inplace_reshard_overwrites_originals(
        self, source_shards: Path, audit_ndjson: Path
    ):
        """Plan with target_dir == source_dir overwrites originals."""
        # Capture source roots so we can later detect topology change.
        source_roots: set[str] = set()
        for shard in source_shards.glob("*.db"):
            c = sqlite3.connect(f"file:{shard}?mode=ro", uri=True)
            try:
                for (root,) in c.execute("SELECT document_root FROM documents"):
                    source_roots.add(root)
            finally:
                c.close()

        plan = HydrationPlanner().plan(
            source_dir=source_shards,
            target_dir=source_shards,   # in-place
            target_M=4,
            free_bytes=200 * 1024 ** 3,
        )
        execute_plan(plan, audit_events_ndjson=audit_ndjson)

        # After in-place: target_dir has exactly M files (no .new, no
        # parallel dir), each carrying corpus_shard_count = 4, and
        # each doc lands on its hashed shard.
        db_files = sorted(p.name for p in source_shards.glob("*.db"))
        assert db_files == ["000.db", "001.db", "002.db", "003.db"]
        assert not list(source_shards.glob("*.db.new"))

        for tidx in range(4):
            tdb = source_shards / f"{tidx:03d}.db"
            c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
            try:
                # corpus_shard_count meta set.
                row = c.execute(
                    "SELECT value FROM meta WHERE key='corpus_shard_count'"
                ).fetchone()
                assert row and int(row[0]) == 4
                # All docs route to this shard.
                for (root,) in c.execute("SELECT document_root FROM documents"):
                    assert shard_for_document(root, 4) == tidx
            finally:
                c.close()

    def test_validation_failure_leaves_new_files(
        self, source_shards: Path, audit_ndjson: Path, tmp_path: Path
    ):
        """When row-count validation fails, .db.new files survive for
        inspection and no atomic rename runs."""
        from arborist.migrate import execute_plan as run
        target = tmp_path / "tgt"
        plan = HydrationPlanner().plan(
            source_dir=source_shards,
            target_dir=target,
            target_M=4,
            free_bytes=200 * 1024 ** 3,
        )
        # Pass a deliberately wrong expected row count to trip the
        # >1% tolerance guard.
        with pytest.raises(RuntimeError, match="validation failed"):
            run(
                plan,
                audit_events_ndjson=audit_ndjson,
                expected_row_counts={"documents": 9_999_999},
            )
        # Inspection artefacts must remain.
        assert sorted(p.name for p in target.glob("*.db.new")) == [
            "000.db.new", "001.db.new", "002.db.new", "003.db.new",
        ]
        # No promoted .db files yet.
        assert not list(target.glob("00[0-9].db"))


class TestWalCheckpointing:
    """Regression test: WAL files must be reclaimed between executor
    phases so they don't pin disk through subsequent passes. The
    2026-05-26 production migration came within 7 GB of ENOSPC because
    SQLite auto-checkpoint can't run while a reader cursor is open and
    the FTS rebuild keeps a cursor open through every chunk. Manual
    PRAGMA wal_checkpoint(TRUNCATE) freed 27 GB."""

    def test_no_large_wal_after_migration(
        self, source_shards: Path, audit_ndjson: Path, tmp_path: Path
    ):
        target = tmp_path / "tgt"
        plan = HydrationPlanner().plan(
            source_dir=source_shards,
            target_dir=target,
            target_M=4,
            free_bytes=200 * 1024 ** 3,
        )
        execute_plan(plan, audit_events_ndjson=audit_ndjson)

        # After completion, no .db-wal file should retain more than a
        # few MB. The checkpoint(TRUNCATE) calls inside the executor
        # collapse the WAL between phases; final VACUUM produces a
        # zero-WAL state.
        for i in range(plan.target_M):
            for suffix in (".db-wal", ".db.new-wal"):
                wal = target / f"{i:03d}{suffix}"
                if wal.exists():
                    assert wal.stat().st_size < 4 * 1024 * 1024, (
                        f"WAL retained {wal.stat().st_size / 1e6:.1f} MB "
                        f"at {wal} — checkpoint between phases didn't fire"
                    )


class TestCrossShardForeignKeys:
    """Regression test: derivations.src_root can legitimately reference
    a surface doc that hash-routes to a different target shard. FK
    enforcement on the migration writer connection must be OFF or the
    INSERT fails (the 2026-05-26 cutover-crash reproducer)."""

    def test_cross_shard_derivation_succeeds(
        self, tmp_path: Path
    ):
        """Synthesize a derivation row whose core_root and src_root
        hash to different target shards, run migration, verify it
        lands without FK violation."""
        # Build a source shard with two docs A and B that we KNOW hash
        # to different target shards under M=4. Then INSERT a
        # derivation linking them, run migration, expect success.
        src_dir = tmp_path / "src"
        src_dir.mkdir()
        src_db = src_dir / "000.db"
        conn = connect(src_db)
        try:
            # Two docs with content chosen so their roots split across
            # M=4 targets. Find them by trial.
            roots: dict[int, str] = {}
            for i in range(200):
                content = f"crosshard fixture seed {i} body text " * 30
                ds = _FakeSource([
                    Document(
                        uri=f"https://test/cs-{i}",
                        content=content,
                        source_type="html",
                        title=f"cs-{i}",
                    )
                ])
                ingest_source(conn, ds)
                row = conn.execute(
                    "SELECT document_root FROM documents WHERE document_uri = ?",
                    (f"https://test/cs-{i}",),
                ).fetchone()
                root = row[0]
                idx = shard_for_document(root, 4)
                roots.setdefault(idx, root)
                if len(roots) >= 2:
                    break
            assert len(roots) >= 2, "couldn't find two docs hashing to different shards"
            two_idxs = sorted(roots.keys())[:2]
            core_root = roots[two_idxs[0]]
            src_root = roots[two_idxs[1]]
            # Synthesize a derivation linking core_root (target A) to
            # src_root (target B != A).
            conn.execute(
                "INSERT INTO derivations "
                "(core_root, src_root, proof_blob, process_id, distilled_at) "
                "VALUES (?, ?, ?, ?, ?)",
                (core_root, src_root, "{}", "test_cross_shard", 0),
            )
            conn.commit()
        finally:
            conn.close()

        # Extract audit events.
        ndjson = tmp_path / "a.ndjson"
        with ndjson.open("w") as f:
            c = sqlite3.connect(f"file:{src_db}?mode=ro", uri=True)
            try:
                cur = c.execute(
                    "SELECT seq, event_hash, prev_event_hash, event_type, "
                    "subject_root, body, ts FROM audit_events ORDER BY seq"
                )
                for seq, eh, peh, et, subj, body, ts in cur:
                    f.write(json.dumps({
                        "src_shard": 0, "src_seq": seq, "src_event_hash": eh,
                        "prev_event_hash": peh, "event_type": et,
                        "subject_root": subj, "body": body, "ts": ts,
                    }, separators=(",", ":")) + "\n")
            finally:
                c.close()

        target = tmp_path / "tgt"
        plan = HydrationPlanner().plan(
            source_dir=src_dir, target_dir=target, target_M=4,
            free_bytes=200 * 1024 ** 3,
        )
        # The reproducer: pre-fix this raised sqlite3.IntegrityError
        # ("FOREIGN KEY constraint failed") inside _route_per_doc_table.
        from arborist.migrate import execute_plan as run
        run(plan, audit_events_ndjson=ndjson)

        # Verify: derivation row lands on the core_root's target shard
        # (where the row is routed by `core_root`); src_root reference
        # points cross-shard.
        tdb = target / f"{shard_for_document(core_root, 4):03d}.db"
        c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
        try:
            row = c.execute(
                "SELECT core_root, src_root FROM derivations "
                "WHERE process_id = 'test_cross_shard'"
            ).fetchone()
        finally:
            c.close()
        assert row is not None
        assert row[0] == core_root
        assert row[1] == src_root
        # src_root is NOT in this same target shard — that's the
        # cross-shard semantic this test exists to protect.
        c = sqlite3.connect(f"file:{tdb}?mode=ro", uri=True)
        try:
            row = c.execute(
                "SELECT document_root FROM documents WHERE document_root = ?",
                (src_root,),
            ).fetchone()
        finally:
            c.close()
        assert row is None, "src_root unexpectedly local; fixture didn't cross-shard"


class TestExecutorErrorCases:
    def test_missing_ndjson_raises(
        self, source_shards: Path, tmp_path: Path
    ):
        target = tmp_path / "tgt"
        plan = HydrationPlanner().plan(
            source_dir=source_shards,
            target_dir=target,
            target_M=4,
            free_bytes=200 * 1024 ** 3,
        )
        with pytest.raises(FileNotFoundError, match="audit-event ndjson"):
            execute_plan(
                plan, audit_events_ndjson=tmp_path / "does-not-exist.ndjson"
            )

    def test_per_source_shard_not_implemented_yet(
        self, source_shards: Path, audit_ndjson: Path, tmp_path: Path
    ):
        target = tmp_path / "tgt"
        plan = HydrationPlanner().plan(
            source_dir=source_shards,
            target_dir=target,
            target_M=4,
            free_bytes=200 * 1024 ** 3,
            force_strategy="per_source_shard",
        )
        with pytest.raises(NotImplementedError, match="per_source_shard"):
            execute_plan(plan, audit_events_ndjson=audit_ndjson)