#!/usr/bin/env python3
"""Build lazy compressed text shards for the FinanceBench Corpus tab."""

from __future__ import annotations

import argparse
import base64
import gzip
import json
import shutil
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterator


ROOT = Path(__file__).resolve().parent.parent
DEFAULT_CORPUS = Path(
    "/mnt/ramdisk/blobstore/timchen0618/data/eval/financebench/"
    "unified/corpus.unified.jsonl"
)


def iter_jsonl(path: Path) -> Iterator[dict[str, Any]]:
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                yield json.loads(line)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--corpus", type=Path, default=DEFAULT_CORPUS)
    parser.add_argument("--out", type=Path, default=ROOT / "corpus_text")
    args = parser.parse_args()

    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    metadata: dict[str, dict[str, Any]] = {}
    for row in iter_jsonl(args.corpus):
        meta = row.get("meta") or {}
        doc_name = str(meta["doc_name"])
        metadata.setdefault(doc_name, dict(meta))
        grouped[doc_name].append(
            {
                "id": str(row["id"]),
                "page_start": meta.get("page_start"),
                "page_end": meta.get("page_end"),
                "contents": str(row.get("contents") or ""),
            }
        )

    records_dir = args.out / "records"
    if records_dir.exists():
        shutil.rmtree(records_dir)
    records_dir.mkdir(parents=True, exist_ok=True)
    index_rows = []
    for doc_name in sorted(grouped):
        chunks = sorted(grouped[doc_name], key=lambda row: row["page_start"] or 0)
        record = {
            "doc_name": doc_name,
            "meta": metadata[doc_name],
            "chunks": chunks,
        }
        raw = json.dumps(record, ensure_ascii=False, separators=(",", ":")).encode()
        encoded = base64.b64encode(gzip.compress(raw, compresslevel=9)).decode("ascii")
        path = f"records/{doc_name}.json.gzip.b64"
        (args.out / path).write_text(encoded, encoding="ascii")
        index_rows.append(
            {
                "doc_name": doc_name,
                "n_chunks": len(chunks),
                "n_pages": max(chunk["page_end"] or 0 for chunk in chunks),
                "path": path,
            }
        )

    args.out.mkdir(parents=True, exist_ok=True)
    (args.out / "index.json").write_text(
        json.dumps(
            {
                "meta": {
                    "n_documents": len(index_rows),
                    "n_chunks": sum(row["n_chunks"] for row in index_rows),
                    "source": str(args.corpus),
                },
                "rows": index_rows,
            },
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )
    print(f"Wrote {len(index_rows)} documents to {args.out}")


if __name__ == "__main__":
    main()