src/lib/system/migration.py

"""Move data between local and cloud backends, both directions.

Two components have backends today (more later — R2, DeepInfra):
    - inventory_db  : SQLite file  <->  Turso libsql
    - vector_db     : local Qdrant <->  Qdrant Cloud

Each migration is full-replace: target tables / collection are wiped first,
then every row / point copied over. Idempotent. No partial sync mode (kept
simple on purpose; we're a single-user side project).

Public API:
    migrate_inventory(direction)        # 'local_to_cloud' | 'cloud_to_local'
    migrate_vectors(direction, collection=None)
    InventoryReport, VectorReport       # what got moved

CLI:
    python -m src.lib.system.migrate inventory --to cloud
    python -m src.lib.system.migrate vectors   --to local --collection library__bge-m3
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field

from src.config import load_config
from src.stage1_inventory.db import open_local, open_turso

# --- Tables to migrate, in FK-safe insert order (parents before children) ---
INVENTORY_TABLES = (
    "books",
    "labels",
    "book_labels",
    "book_indexes",
    "processing_log",
    "llm_calls",
    "query_history",
    "printed_pages",
)

# Reference tables: deterministic primary keys, safe to merge with INSERT OR IGNORE.
# In merge mode we add new rows but leave existing target rows untouched.
INVENTORY_REFERENCE_TABLES = frozenset((
    "books",          # PK book_id
    "labels",         # PK label_id
    "book_labels",    # composite PK (book_id, label_id)
    "book_indexes",   # composite PK (book_id, collection_name)
    "printed_pages",  # composite PK (book_id, pdf_page)
))

# Log tables: AUTOINCREMENT integer PK; "merging" them would either overwrite
# unrelated rows that share an id, or duplicate them. Merge mode skips these
# tables on purpose — logs are local to where the work happened.
INVENTORY_LOG_TABLES = frozenset((
    "processing_log",
    "llm_calls",
    "query_history",
))

BATCH_INSERT = 200


# =====================================================================
# Models
# =====================================================================
@dataclass
class InventoryReport:
    direction: str
    mode: str = "replace"            # "merge" | "replace"
    rows_per_table: dict[str, int] = field(default_factory=dict)
    skipped_tables: list[str] = field(default_factory=list)
    elapsed_ms: int = 0

    @property
    def total_rows(self) -> int:
        return sum(self.rows_per_table.values())


@dataclass
class VectorReport:
    direction: str
    collection: str
    mode: str = "replace"            # "merge" | "replace"
    points: int = 0
    elapsed_ms: int = 0


@dataclass
class SideSummary:
    """Cheap snapshot of one side's data, used for the migration preview UI."""
    available: bool
    detail: str
    error: str = ""
    rows_per_table: dict[str, int] = field(default_factory=dict)
    points: int = 0  # vectors only

    @property
    def total_rows(self) -> int:
        return sum(self.rows_per_table.values())


# =====================================================================
# Inventory DB migration
# =====================================================================
def _columns_of(conn, table: str) -> list[str]:
    rows = conn.execute(f"PRAGMA table_info({table})").fetchall()
    # PRAGMA table_info returns rows shaped (cid, name, type, notnull, dflt_value, pk).
    # Both backends expose name at index 1. (Our DictRow wrapper supports indexing.)
    return [r[1] for r in rows]


def _copy_table(src_conn, dst_conn, table: str, mode: str) -> int:
    """Copy `table` from src to dst.

    - mode='replace': DELETE all target rows first, then bulk INSERT from source.
    - mode='merge':   INSERT OR IGNORE — adds new rows, keeps existing target
                      rows untouched. Safe to re-run.
    """
    cols = _columns_of(src_conn, table)
    if not cols:
        return 0
    col_list = ", ".join(cols)
    placeholders = ", ".join("?" * len(cols))

    if mode == "replace":
        dst_conn.execute(f"DELETE FROM {table}")
        dst_conn.commit()
        insert_verb = "INSERT"
    else:  # merge
        insert_verb = "INSERT OR IGNORE"

    src_rows = src_conn.execute(f"SELECT {col_list} FROM {table}").fetchall()
    if not src_rows:
        return 0

    sql = f"{insert_verb} INTO {table} ({col_list}) VALUES ({placeholders})"
    payload = [tuple(r[i] for i in range(len(cols))) for r in src_rows]
    for i in range(0, len(payload), BATCH_INSERT):
        dst_conn.executemany(sql, payload[i : i + BATCH_INSERT])
    dst_conn.commit()
    return len(payload)


def migrate_inventory(direction: str, mode: str = "merge") -> InventoryReport:
    """direction = 'local_to_cloud' | 'cloud_to_local'.

    mode = 'merge'   (safer default): INSERT OR IGNORE rows on every reference
                     table; skip the AUTOINCREMENT log tables (processing_log,
                     llm_calls, query_history) because re-using their integer
                     ids across machines is meaningless. Idempotent.

    mode = 'replace': DELETE every table on the target side, then INSERT a
                     full copy from the source. Destructive on target.

    Both connections are opened directly (bypassing the config dispatch) so
    the active backend setting doesn't matter. Schema is auto-created on the
    target side via init_schema().
    """
    if direction not in ("local_to_cloud", "cloud_to_local"):
        raise ValueError(f"unknown direction {direction!r}")
    if mode not in ("merge", "replace"):
        raise ValueError(f"unknown mode {mode!r}")

    started = time.monotonic()
    if direction == "local_to_cloud":
        src = open_local()
        dst = open_turso()
    else:
        src = open_turso()
        dst = open_local()

    report = InventoryReport(direction=direction, mode=mode)
    try:
        # Foreign keys are ON on local; defer them while we wipe + bulk insert.
        try:
            dst.execute("PRAGMA foreign_keys = OFF")
        except Exception:  # noqa: BLE001
            pass  # turso may not support PRAGMA the same way; best effort

        for table in INVENTORY_TABLES:
            # Merge skips log tables on purpose
            if mode == "merge" and table in INVENTORY_LOG_TABLES:
                report.skipped_tables.append(table)
                continue
            try:
                report.rows_per_table[table] = _copy_table(src, dst, table, mode)
            except Exception as e:  # noqa: BLE001
                report.rows_per_table[table] = -1
                print(f"[migrate_inventory] {table}: FAILED {type(e).__name__}: {e}")

        try:
            dst.execute("PRAGMA foreign_keys = ON")
        except Exception:  # noqa: BLE001
            pass
    finally:
        try: src.close()
        except Exception: pass  # noqa: BLE001
        try: dst.close()
        except Exception: pass  # noqa: BLE001

    report.elapsed_ms = int((time.monotonic() - started) * 1000)
    return report


# =====================================================================
# Vector DB migration
# =====================================================================
def _local_qdrant_client():
    from qdrant_client import QdrantClient

    cfg = load_config()
    vc = cfg.section("vector_db")
    return QdrantClient(host=vc.get("host", "localhost"), port=int(vc.get("port", 6333)))


def _copy_payload_indexes(src, dst, coll: str, src_info=None) -> None:
    """Replicate the source collection's payload indexes onto the target.

    Without this, filter queries on cloud (e.g. WHERE language='ar') fail with
    'Index required but not found'. We read `payload_schema` from source and
    re-create each entry on target. Idempotent — if a target index already
    exists, Qdrant treats it as a no-op.
    """
    if src_info is None:
        src_info = src.get_collection(coll)
    schema = src_info.payload_schema or {}
    for field_name, field_info in schema.items():
        try:
            dst.create_payload_index(
                collection_name=coll,
                field_name=field_name,
                field_schema=field_info.data_type,
            )
        except Exception as e:  # noqa: BLE001
            print(f"[migrate_vectors] payload index for {field_name!r} failed: {type(e).__name__}: {e}")


def _cloud_qdrant_client():
    import os

    from qdrant_client import QdrantClient

    from src.stage6_indexing.qdrant_io import CLOUD_TIMEOUT_S

    url = os.environ.get("QDRANT_URL")
    if not url:
        raise RuntimeError("QDRANT_URL not set in .env (cannot migrate vectors)")
    return QdrantClient(url=url, api_key=os.environ.get("QDRANT_API_KEY"), timeout=CLOUD_TIMEOUT_S)


VECTOR_BATCH_LOCAL = 256        # local-target uploads can be big
VECTOR_BATCH_CLOUD = 64         # cloud-target uploads need to fit in one HTTP timeout


def migrate_vectors(
    direction: str,
    collection: str | None = None,
    batch: int | None = None,
    mode: str = "merge",
) -> VectorReport:
    """Copy every point in `collection` from source backend to target backend.

    mode = 'merge'   (default): if the target collection doesn't exist, create
                     it with the source's vectors_config; then upsert every
                     source point. Existing target points with the same id get
                     updated (Qdrant point ids are deterministic uuid5 of
                     chunk_id, so re-running is idempotent). Existing points
                     not present in source are kept.

    mode = 'replace': delete target collection if present, recreate from
                     source's config, upsert every source point. Destructive.
    """
    if direction not in ("local_to_cloud", "cloud_to_local"):
        raise ValueError(f"unknown direction {direction!r}")
    if mode not in ("merge", "replace"):
        raise ValueError(f"unknown mode {mode!r}")

    from src.stage6_indexing.qdrant_io import default_collection_name

    coll = collection or default_collection_name()
    started = time.monotonic()

    if direction == "local_to_cloud":
        src = _local_qdrant_client()
        dst = _cloud_qdrant_client()
        target_is_cloud = True
    else:
        src = _cloud_qdrant_client()
        dst = _local_qdrant_client()
        target_is_cloud = False

    # Default batch depends on which side is the upload target
    if batch is None:
        batch = VECTOR_BATCH_CLOUD if target_is_cloud else VECTOR_BATCH_LOCAL

    # Read source collection config (always needed if target needs creation)
    info = src.get_collection(coll)
    vectors_config = info.config.params.vectors
    sparse_config = info.config.params.sparse_vectors

    if mode == "replace":
        if dst.collection_exists(coll):
            dst.delete_collection(coll)
        dst.create_collection(
            collection_name=coll,
            vectors_config=vectors_config,
            sparse_vectors_config=sparse_config,
        )
        _copy_payload_indexes(src, dst, coll, info)
    else:  # merge
        if not dst.collection_exists(coll):
            dst.create_collection(
                collection_name=coll,
                vectors_config=vectors_config,
                sparse_vectors_config=sparse_config,
            )
            _copy_payload_indexes(src, dst, coll, info)
        else:
            # Existing target — make sure every source index also exists on
            # target (idempotent; create_payload_index is a no-op if present).
            _copy_payload_indexes(src, dst, coll, info)

    # Scroll source points and upsert (upsert is idempotent on point id)
    moved = 0
    next_offset = None
    while True:
        points, next_offset = src.scroll(
            collection_name=coll,
            limit=batch,
            offset=next_offset,
            with_payload=True,
            with_vectors=True,
        )
        if not points:
            break

        from qdrant_client.http import models as qm
        upsert_payload = [
            qm.PointStruct(id=p.id, payload=p.payload, vector=p.vector)
            for p in points
        ]
        # wait=False on cloud so the server queues the batch and returns
        # immediately — the HTTP round-trip stays well under the timeout.
        dst.upsert(collection_name=coll, points=upsert_payload, wait=not target_is_cloud)
        moved += len(points)
        if next_offset is None:
            break

    # On cloud target with wait=False, give the server a moment to finish
    # indexing before reporting. Cheap: one count() round-trip.
    if target_is_cloud:
        try:
            final = dst.count(collection_name=coll, exact=True).count
            print(f"[migrate_vectors] cloud collection now has {final} points")
        except Exception:  # noqa: BLE001
            pass

    elapsed_ms = int((time.monotonic() - started) * 1000)
    return VectorReport(direction=direction, collection=coll, mode=mode, points=moved, elapsed_ms=elapsed_ms)


# =====================================================================
# Side summaries (no data is moved — just counts for the preview UI)
# =====================================================================
def inventory_summary() -> dict[str, SideSummary]:
    """Return {'local': SideSummary, 'cloud': SideSummary} for the inventory DB.

    Each summary lists row counts per table (so the user can see what's at
    stake before migrating). Errors are caught and surfaced in the `error`
    field so the UI never crashes if e.g. Turso is unreachable.
    """
    out: dict[str, SideSummary] = {}

    # --- Local SQLite ---
    try:
        src = open_local()
        rows: dict[str, int] = {}
        for t in INVENTORY_TABLES:
            try:
                row = src.execute(f"SELECT COUNT(*) FROM {t}").fetchone()
                rows[t] = int(row[0])
            except Exception:  # noqa: BLE001
                rows[t] = 0
        src.close()
        out["local"] = SideSummary(
            available=True,
            detail=f"{rows.get('books', 0)} books · {sum(rows.values())} total rows",
            rows_per_table=rows,
        )
    except Exception as e:  # noqa: BLE001
        out["local"] = SideSummary(available=False, detail="local SQLite unreachable", error=f"{type(e).__name__}: {e}")

    # --- Cloud Turso ---
    try:
        dst = open_turso()
        rows = {}
        for t in INVENTORY_TABLES:
            try:
                row = dst.execute(f"SELECT COUNT(*) FROM {t}").fetchone()
                rows[t] = int(row[0])
            except Exception:  # noqa: BLE001
                rows[t] = 0
        dst.close()
        out["cloud"] = SideSummary(
            available=True,
            detail=f"{rows.get('books', 0)} books · {sum(rows.values())} total rows",
            rows_per_table=rows,
        )
    except Exception as e:  # noqa: BLE001
        out["cloud"] = SideSummary(available=False, detail="Turso unreachable / not configured", error=f"{type(e).__name__}: {e}")

    return out


def vector_summary(collection: str | None = None) -> dict[str, SideSummary]:
    """Return {'local': SideSummary, 'cloud': SideSummary} for a Qdrant collection."""
    from src.stage6_indexing.qdrant_io import default_collection_name

    coll = collection or default_collection_name()
    out: dict[str, SideSummary] = {}

    # --- Local Docker Qdrant ---
    try:
        client = _local_qdrant_client()
        if client.collection_exists(coll):
            n = client.get_collection(coll).points_count
            out["local"] = SideSummary(available=True, detail=f"{n} points · {coll}", points=int(n or 0))
        else:
            out["local"] = SideSummary(available=True, detail=f"collection {coll!r} not present", points=0)
    except Exception as e:  # noqa: BLE001
        out["local"] = SideSummary(available=False, detail="local Qdrant unreachable", error=f"{type(e).__name__}: {e}")

    # --- Cloud Qdrant Cloud ---
    try:
        client = _cloud_qdrant_client()
        if client.collection_exists(coll):
            n = client.get_collection(coll).points_count
            out["cloud"] = SideSummary(available=True, detail=f"{n} points · {coll}", points=int(n or 0))
        else:
            out["cloud"] = SideSummary(available=True, detail=f"collection {coll!r} not present", points=0)
    except Exception as e:  # noqa: BLE001
        out["cloud"] = SideSummary(available=False, detail="Qdrant Cloud unreachable / not configured", error=f"{type(e).__name__}: {e}")

    return out