#!/usr/bin/env python3
"""Convert the inline `fund_statements.coverage` JSON of existing rows to `fund_coverage_blob` references.

    cd ~/hfmarketdata && HFMD_DATA_ROOT=~/firstratedata venv/bin/python scripts/migrate_coverage.py            # all rows
    venv/bin/python scripts/migrate_coverage.py --batch 5000 --max-batches 50                                    # a slice
    venv/bin/python scripts/migrate_coverage.py --vacuum                                                         # reclaim space (offline!)

Runs while the API serves (WAL, short transactions of `--batch` rows; each batch commits, so it can be stopped
and resumed at any time: the work set is simply `WHERE coverage IS NOT NULL`). New rows are written through the
dictionary already (`ingest._replace_statements`), the read side understands both shapes, so this only shrinks the
past: 1.22 M rows × ~1.1 KB ≈ 1.3 GB of JSON → a few tens of thousands of blobs.

Space is NOT returned to the OS by the conversion itself (free pages are reused by later inserts). `--vacuum`
runs `PRAGMA auto_vacuum=INCREMENTAL` + `VACUUM` afterwards: needs ~2× the file size of free disk, exclusive access
(stop the API and the jobs) and a few minutes for 3 GB — run it once, off-hours.

Author: Simon-Pierre Boucher <contact@spboucher.ai>
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "hfmarketdata" / "api"))

log = logging.getLogger("migrate_coverage")


def convert(batch: int, max_batches: int | None, sleep_s: float) -> tuple[int, int]:
    """Convert rows in id order; returns (rows converted, batches)."""
    from sqlalchemy import text

    from core.db import engine
    from fundamentals import coverage_store as CS
    from fundamentals.models import init_db

    init_db()   # makes sure coverage_id / fund_coverage_blob exist
    done = batches = 0
    last_id = 0
    t0 = time.time()
    while True:
        with engine.begin() as con:
            rows = con.execute(text("SELECT id, coverage FROM fund_statements WHERE id > :last AND coverage IS NOT NULL "
                                    "ORDER BY id LIMIT :n"), {"last": last_id, "n": batch}).all()
            if not rows:
                break
            docs = []
            for _, cov in rows:
                if isinstance(cov, (bytes, bytearray)):
                    cov = cov.decode()
                docs.append(json.loads(cov) if isinstance(cov, str) else (cov or {}))
            ids = CS.intern(docs)
            con.execute(text("UPDATE fund_statements SET coverage = NULL, coverage_id = :cid WHERE id = :id"),
                        [{"cid": cid, "id": rid} for (rid, _), cid in zip(rows, ids)])
            last_id = int(rows[-1][0])
        done += len(rows)
        batches += 1
        if batches % 20 == 0:
            log.info("converted %d rows (last id %d) in %.0fs", done, last_id, time.time() - t0)
        if max_batches is not None and batches >= max_batches:
            break
        if sleep_s:
            time.sleep(sleep_s)
    return done, batches


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--batch", type=int, default=2000, help="rows per transaction (default 2000)")
    ap.add_argument("--max-batches", type=int, help="stop after N batches (resume later)")
    ap.add_argument("--sleep", type=float, default=0.0, help="pause between batches (seconds) to stay gentle with the API")
    ap.add_argument("--vacuum", action="store_true", help="after conversion: auto_vacuum=INCREMENTAL + VACUUM (exclusive, slow)")
    ap.add_argument("-v", "--verbose", action="store_true")
    args = ap.parse_args()
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")

    from sqlalchemy import text

    from core.config import settings
    from core.db import engine

    log.info("database: %s", settings.state_db)
    t0 = time.time()
    done, batches = convert(args.batch, args.max_batches, args.sleep)
    with engine.connect() as con:
        left = con.execute(text("SELECT count(*) FROM fund_statements WHERE coverage IS NOT NULL")).scalar()
        blobs = con.execute(text("SELECT count(*) FROM fund_coverage_blob")).scalar()
    log.info("converted %d rows in %d batches (%.0fs); rows still inline: %d; distinct blobs: %d", done, batches,
             time.time() - t0, left, blobs)
    if args.vacuum:
        t1 = time.time()
        with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as con:
            before = con.execute(text("PRAGMA page_count")).scalar()
            con.execute(text("PRAGMA auto_vacuum=INCREMENTAL"))
            con.execute(text("VACUUM"))
            after = con.execute(text("PRAGMA page_count")).scalar()
            mode = con.execute(text("PRAGMA auto_vacuum")).scalar()
        log.info("VACUUM: %s → %s pages (auto_vacuum=%s) in %.0fs", before, after, mode, time.time() - t1)
    return 0 if left == 0 or args.max_batches else 1


if __name__ == "__main__":
    sys.exit(main())