# Batched full-corpus LLM extraction driver (v2, 2026-10-27).
# Sections are loaded per CIK-batch (bounds memory) and every processed section is
# recorded in the `extracted_sections` ledger, so a restart never re-pays for done
# work and newly parsed sections of already-covered companies are picked up.
#
#   PYTHONPATH=. ../.venv/bin/python run_full_extract.py [--workers 24] [--batch 20] [--dry-run]
from __future__ import annotations

import argparse
import sys

from spid.config import get_logger
from spid.pipeline import run_extraction
from spid.store import Store

log = get_logger("driver")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--batch", type=int, default=20, help="CIKs per batch")
    ap.add_argument("--workers", type=int, default=24)
    ap.add_argument("--backend", default="llm")
    ap.add_argument("--max-sections", type=int, default=None, help="stop after roughly this many sections (budget cap)")
    ap.add_argument("--dry-run", action="store_true", help="only count pending sections per CIK")
    args = ap.parse_args()

    store = Store()
    store.con.execute("SET arrow_large_buffer_size=true")
    pending = store.con.execute("""
        SELECT cik, COUNT(*) FROM v_sections
        WHERE word_count >= 30 AND section_id NOT IN (SELECT section_id FROM extracted_sections)
        GROUP BY cik ORDER BY cik""").fetchall()
    total = sum(n for _, n in pending)
    log.info("pending: %d sections over %d CIKs | batch=%d workers=%d backend=%s",
             total, len(pending), args.batch, args.workers, args.backend)
    if args.dry_run:
        return 0

    ciks = [c for c, _ in pending]
    n_batches = (len(ciks) + args.batch - 1) // args.batch
    cum_sec = cum_men = 0
    for bi in range(n_batches):
        batch = ciks[bi * args.batch:(bi + 1) * args.batch]
        try:
            res = run_extraction(store, backend=args.backend, ciks=batch, forms=None, workers=args.workers)
        except Exception as e:  # noqa: BLE001
            log.error("BATCH %d FAILED (%s) ciks=%s - continuing", bi + 1, e, batch)
            continue
        cum_sec += res["sections"]
        cum_men += res["mentions"]
        log.info("BATCH %d/%d done | batch_sec=%d batch_men=%d | CUM sections=%d/%d mentions=%d",
                 bi + 1, n_batches, res["sections"], res["mentions"], cum_sec, total, cum_men)
        if args.max_sections and cum_sec >= args.max_sections:
            log.info("budget cap reached (%d sections) - stopping", cum_sec)
            break
    log.info("EXTRACT RUN COMPLETE | sections=%d mentions=%d | TOTAL mentions in DB=%d",
             cum_sec, cum_men, store.count("project_mentions"))
    store.close()
    return 0


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