"""Backfill `financebench.documents.page_count` from the PDFs themselves.

The FinanceBench release files carry no page counts, so this reads each
document's PDF from the bucket (through the quber S3 cache) and records its
page count. Run it after `load.py` — the loader rebuilds the schema, which
clears the column. A missing PDF logs a warning and leaves the count NULL;
it is never a silent skip.
"""

from __future__ import annotations

from typing import Optional

import pymupdf
from loguru import logger

from quber.playground import db

BUCKET = "qubera-docs"


def backfill() -> tuple[int, int]:
    """Set page_count for every document whose PDF is in the bucket.

    Returns (updated, missing) counts.
    """
    from quber.files.cache import resolve_document

    with db.connect() as conn:
        rows = conn.execute(
            "SELECT doc_name, s3_key FROM financebench.documents ORDER BY doc_name"
        ).fetchall()
        updated = missing = 0
        for doc_name, s3_key in rows:
            try:
                local = resolve_document(f"s3://{BUCKET}/{s3_key}")
                with pymupdf.open(local) as pdf:
                    n = pdf.page_count
            except Exception as exc:
                logger.warning("No page count for {} ({}): {}", doc_name, s3_key, exc)
                missing += 1
                continue
            conn.execute(
                "UPDATE financebench.documents SET page_count = %s WHERE doc_name = %s",
                (n, doc_name),
            )
            updated += 1

    logger.success("Page counts: {} documents updated, {} missing PDFs", updated, missing)
    return updated, missing


def main(argv: Optional[list[str]] = None) -> None:
    backfill()


if __name__ == "__main__":
    main()