"""Pull the astro-ph text slice out of ArxivCap via DuckDB (no image download)."""
import argparse
import os
import time
from collections import Counter
from pathlib import Path

import duckdb

SRC = "hf://datasets/MMInstruction/ArxivCap/data/*.parquet"

EXTRACT = """
SELECT
    arxiv_id,
    title,
    abstract,
    meta.meta_from_kaggle.categories AS categories,
    meta.meta_from_s2.citationCount AS citations,
    list_transform(caption_images, x -> x.caption) AS captions,
    list_transform(
        caption_images,
        x -> substr(
            array_to_string(
                flatten(list_transform(x.cil_pairs, y -> y.image_ocr)), ' '
            ), 1, 500)
    ) AS ocr
FROM read_parquet('{src}')
WHERE contains(meta.meta_from_kaggle.categories, 'astro-ph')
{limit}
"""


def year_from_arxiv_id(arxiv_id):
    try:
        yy = int(arxiv_id.split("/")[1][:2]) if "/" in arxiv_id else int(arxiv_id[:2])
        return 1900 + yy if yy > 80 else 2000 + yy
    except (ValueError, IndexError):
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=SRC)
    ap.add_argument("--out", default="astro_captions.parquet")
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--parts-dir", default="census_parts")
    args = ap.parse_args()

    con = duckdb.connect()
    for setting in ["SET enable_progress_bar = true", "SET http_retries = 8",
                    "SET http_retry_wait_ms = 2000", "SET http_retry_backoff = 2",
                    "SET threads = 2"]:
        try:
            con.execute(setting)
        except duckdb.Error:
            pass

    token = os.environ.get("HF_TOKEN", "")
    if token:
        con.execute(f"CREATE SECRET hf (TYPE HUGGINGFACE, TOKEN '{token}')")
    else:
        print("WARNING: no HF_TOKEN set, anonymous requests get rate limited fast.")

    print("Extracting astro-ph slice (text columns only)...")
    if args.limit or not args.src.startswith("hf://"):
        limit = f"LIMIT {args.limit}" if args.limit else ""
        con.execute(f"COPY ({EXTRACT.format(src=args.src, limit=limit)}) "
                    f"TO '{args.out}' (FORMAT PARQUET)")
    else:
        parts = Path(args.parts_dir)
        parts.mkdir(exist_ok=True)
        chunks = [f"{y:02d}" for y in list(range(91, 100)) + list(range(0, 24))]
        for yy in chunks:
            part_file = parts / f"astro_{yy}.parquet"
            if part_file.exists():
                continue
            chunk_src = args.src.replace("*", f"arXiv_src_{yy}*")
            query = EXTRACT.format(src=chunk_src, limit="")
            for attempt in range(5):
                try:
                    con.execute(f"COPY ({query}) TO '{part_file}' (FORMAT PARQUET)")
                    print(f"  {yy}: done")
                    break
                except duckdb.Error as e:
                    msg = str(e)
                    if "No files found" in msg:
                        print(f"  {yy}: no files, skipping")
                        break
                    if attempt == 4:
                        raise
                    wait = 30 * (attempt + 1)
                    print(f"  {yy}: {msg[:80]} -- retrying in {wait}s")
                    time.sleep(wait)
        con.execute(f"COPY (SELECT * FROM read_parquet('{parts}/astro_*.parquet')) "
                    f"TO '{args.out}' (FORMAT PARQUET)")

    papers, figures, usable, fpp, cap_chars = con.execute(f"""
        SELECT count(*),
               sum(len(captions)),
               sum(len(list_filter(captions, c -> length(c) >= 20))),
               avg(len(captions)),
               round(avg(list_aggregate(list_transform(captions, c -> length(c)), 'avg')), 0)
        FROM read_parquet('{args.out}')
    """).fetchone()

    print(f"\nPapers: {papers:,}  Figures: {int(figures):,}  "
          f"Usable captions: {int(usable):,}  Figs/paper: {fpp:.1f}  "
          f"Mean caption: {int(cap_chars)} chars")

    ids = [r[0] for r in con.execute(
        f"SELECT arxiv_id FROM read_parquet('{args.out}')").fetchall()]
    by_year = Counter(y for y in (year_from_arxiv_id(i) for i in ids) if y)
    for y in sorted(by_year):
        print(f"  {y}: {by_year[y]:,}")
    print(f"\nSlice written to {args.out}")


if __name__ == "__main__":
    main()