src/lib/storage/migrate.py

"""One-shot migration: copy PDFs from one storage backend to another.

Designed for the local -> r2 case but works in both directions. Reads
from the source backend, writes to the destination backend, and
optionally rewrites books.local_path so subsequent code paths know the
file is now hosted somewhere different.

Usage:
    python -m src.lib.storage.migrate            # local -> r2 (dry-run)
    python -m src.lib.storage.migrate --apply    # actually copy
    python -m src.lib.storage.migrate --apply --update-paths

Idempotent: skips any book that already exists at the destination
(checked via Storage.exists). Re-run after a partial failure to finish.
"""

from __future__ import annotations

import argparse
import sys

from src.stage1_inventory.db import connect

from .factory import _resolve_backend_name  # noqa: PLC2701 — internal but stable
from .local import LocalStorage
from .r2 import R2Storage


def _make(name: str):
    if name == "local":
        return LocalStorage()
    if name == "r2":
        return R2Storage()
    raise ValueError(f"unknown backend {name!r}")


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Copy PDFs between storage backends.")
    p.add_argument("--from", dest="src", default="local", choices=("local", "r2"))
    p.add_argument("--to",   dest="dst", default="r2",    choices=("local", "r2"))
    p.add_argument("--apply", action="store_true",
                   help="Actually copy. Without it, prints the plan only.")
    p.add_argument("--update-paths", action="store_true",
                   help="Rewrite books.local_path to the destination URI.")
    p.add_argument("--book-id", default=None,
                   help="Limit to one book; otherwise iterates every row in books.")
    args = p.parse_args(argv)

    if args.src == args.dst:
        print(f"src == dst ({args.src}); nothing to do.", file=sys.stderr)
        return 1

    src = _make(args.src)
    # Defer destination construction — dry-run shouldn't need R2 creds just
    # to print a plan. Built lazily on first use below.
    dst = None

    # Pull the inventory list (authoritative source of book_ids).
    sql = "SELECT book_id, local_path FROM books"
    sql_args: list = []
    if args.book_id:
        sql += " WHERE book_id = ?"
        sql_args.append(args.book_id)
    sql += " ORDER BY book_id"
    with connect() as conn:
        rows = conn.execute(sql, sql_args).fetchall()

    if not rows:
        print("No books in inventory.")
        return 0

    print(f"{args.src} -> {args.dst}: {len(rows)} book(s) in inventory")
    print(f"backends.storage active = {_resolve_backend_name()!r}")
    print(f"--apply        = {args.apply}")
    print(f"--update-paths = {args.update_paths}")
    print()

    copied = 0
    skipped = 0
    missing = 0
    text_synced = 0
    failed = 0

    for r in rows:
        bid = r["book_id"]
        has_src_pdf = src.exists(bid)

        # Dry run: never construct the destination, so no creds needed.
        if not args.apply:
            verb = "would copy PDF" if has_src_pdf else "no source PDF (text only)"
            print(f"  [plan] {bid}: {verb} {args.src} -> {args.dst}")
            continue

        if dst is None:
            dst = _make(args.dst)

        try:
            # --- PDF copy (only if the source actually has it) ---
            if not has_src_pdf:
                missing += 1
                print(f"  [skip-pdf] {bid}: no PDF in {args.src}")
            elif dst.exists(bid):
                skipped += 1
                print(f"  [done] {bid}: PDF already in {args.dst}")
                if args.update_paths:
                    _update_path_to(bid, dst, args.dst)
            else:
                local_path = src.ensure_local(bid)
                uri = dst.put_pdf(bid, local_path)
                copied += 1
                print(f"  [copy] {bid}: -> {uri}")
                if args.update_paths:
                    _update_path_to(bid, dst, args.dst, uri=uri)

            # --- Text bundle (independent of the PDF) ---
            # Always (re)push the OCR/clean text from local disk if present.
            # Idempotent, and decoupled from the PDF so a book whose source
            # PDF was already pruned locally still gets its text layer to the
            # cloud (the whole point of promoting a book is its text, too).
            for tu in dst.put_text(bid):
                text_synced += 1
                print(f"  [text] {bid}: -> {tu}")
        except Exception as e:  # noqa: BLE001
            failed += 1
            print(f"  [FAIL] {bid}: {type(e).__name__}: {e}", file=sys.stderr)

    print()
    print(f"copied={copied}  skipped(already)={skipped}  "
          f"missing(in src)={missing}  text_synced={text_synced}  failed={failed}")
    return 0 if failed == 0 else 2


def _update_path_to(book_id: str, dst, dst_name: str, *, uri: str | None = None) -> None:
    """Rewrite books.local_path so the inventory reflects the new home."""
    if uri is None:
        # Recompute the URI without re-uploading.
        if dst_name == "local":
            from src.config import load_config
            cfg = load_config()
            uri = f"file://{(cfg.paths.raw_dir / f'{book_id}.pdf').resolve()}"
        else:  # r2
            uri = f"r2://{dst._bucket}/{book_id}.pdf"  # noqa: SLF001
    with connect() as conn:
        conn.execute(
            "UPDATE books SET local_path = ?, updated_at = datetime('now') WHERE book_id = ?",
            (uri, book_id),
        )
        conn.commit()


if __name__ == "__main__":
    raise SystemExit(main())