"""Load the FinanceBench benchmark JSONL files into the `financebench` schema.

Input: the two files from the FinanceBench open-source release —
`financebench_document_information.jsonl` (one row per filing) and
`financebench_open_source.jsonl` (one row per question, with the gold answer
and its evidence passages). The loader applies `schema.sql` (drop + recreate,
so a reload is a clean replace) and inserts everything in one pass.

Every benchmark PDF lives at s3://qubera-docs/financebench/<doc_name>.pdf,
recorded in `documents.s3_key`.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, LiteralString, Optional, cast

from loguru import logger

from quber.playground import db

SCHEMA_FILE = Path(__file__).with_name("schema.sql")
DEFAULT_DATA_DIR = Path(__file__).parents[5] / "financebench" / "data"

S3_PREFIX = "financebench"


def doc_key_for(doc_name: str) -> str:
    """A benchmark doc_name as a slug: lowercased, dashes for underscores.
    '3M_2022_10K' -> '3m-2022-10k'. Nothing calls it. A playground document's
    storage key is the first 16 hex chars of its content hash."""
    return doc_name.lower().replace("_", "-")


def _pg_text(s: Optional[str]) -> Optional[str]:
    """Strip NUL characters: Postgres text fields reject 0x00, and evidence
    full-page text comes from the same PDF text layers that produce them."""
    return s.replace("\x00", "") if s is not None else None


def _rows(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]


def load(data_dir: Path = DEFAULT_DATA_DIR) -> tuple[int, int, int]:
    """Rebuild the financebench schema from the JSONL files.

    Returns (documents, questions, evidence) row counts.
    """
    docs = _rows(data_dir / "financebench_document_information.jsonl")
    questions = _rows(data_dir / "financebench_open_source.jsonl")

    # The source metadata file carries one duplicated doc_name
    # (FOOTLOCKER_2023_annualreport, identical rows except doc_period). Keep
    # the first row per doc_name and report every drop — never fail silently.
    seen: set[str] = set()
    deduped = []
    for d in docs:
        if d["doc_name"] in seen:
            logger.warning("Duplicate doc_name in document information, keeping first: {}", d)
            continue
        seen.add(d["doc_name"])
        deduped.append(d)
    docs = deduped

    with db.connect() as conn:
        # Same trust level as a query literal: a schema file we author.
        conn.execute(cast(LiteralString, SCHEMA_FILE.read_text()))

        for d in docs:
            conn.execute(
                """INSERT INTO financebench.documents
                   (doc_name, company, gics_sector, doc_type, doc_period, doc_link, s3_key)
                   VALUES (%s, %s, %s, %s, %s, %s, %s)""",
                (
                    d["doc_name"],
                    d.get("company"),
                    d.get("gics_sector"),
                    d.get("doc_type"),
                    str(d.get("doc_period") or ""),
                    d.get("doc_link"),
                    f"{S3_PREFIX}/{d['doc_name']}.pdf",
                ),
            )

        n_evidence = 0
        for q in questions:
            conn.execute(
                """INSERT INTO financebench.questions
                   (financebench_id, doc_name, question_type, question_reasoning,
                    domain_question_num, question, answer, justification, dataset_subset_label)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)""",
                (
                    q["financebench_id"],
                    q["doc_name"],
                    q.get("question_type"),
                    q.get("question_reasoning"),
                    str(q.get("domain_question_num") or ""),
                    _pg_text(q["question"]),
                    _pg_text(q["answer"]),
                    _pg_text(q.get("justification")),
                    q.get("dataset_subset_label"),
                ),
            )
            for ev in q.get("evidence") or []:
                conn.execute(
                    """INSERT INTO financebench.evidence
                       (financebench_id, doc_name, evidence_page_num, evidence_text, evidence_text_full_page)
                       VALUES (%s, %s, %s, %s, %s)""",
                    (
                        q["financebench_id"],
                        ev.get("doc_name") or q["doc_name"],
                        ev.get("evidence_page_num"),
                        _pg_text(ev.get("evidence_text")),
                        _pg_text(ev.get("evidence_text_full_page")),
                    ),
                )
                n_evidence += 1

    logger.success(
        "Loaded financebench: {d} documents, {q} questions, {e} evidence passages",
        d=len(docs),
        q=len(questions),
        e=n_evidence,
    )
    return len(docs), len(questions), n_evidence


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

    p = argparse.ArgumentParser(description="Load the FinanceBench JSONL files into Postgres.")
    p.add_argument("--data-dir", default=str(DEFAULT_DATA_DIR), help="directory holding the two JSONL files")
    args = p.parse_args(argv)
    load(Path(args.data_dir))


if __name__ == "__main__":
    main()