import argparse
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from pathlib import Path

import pandas as pd


DEFAULT_CATEGORIES = ["cs.CL", "math.NT", "astro-ph.GA", "q-bio.NC", "physics.optics", "econ.EM"]
ATOM = "{http://www.w3.org/2005/Atom}"
ARXIV_API = "https://export.arxiv.org/api/query"


def clean_text(text: str | None) -> str:
    return " ".join((text or "").split())


def fetch_category(category: str, max_results: int, pause: float) -> list[dict[str, str]]:
    rows = []
    for start in range(0, max_results, 100):
        batch_size = min(100, max_results - start)
        query = urllib.parse.urlencode(
            {
                "search_query": f"cat:{category}",
                "start": start,
                "max_results": batch_size,
                "sortBy": "submittedDate",
                "sortOrder": "descending",
            }
        )
        with urllib.request.urlopen(f"{ARXIV_API}?{query}", timeout=60) as response:
            root = ET.fromstring(response.read())

        for entry in root.findall(f"{ATOM}entry"):
            title = clean_text(entry.findtext(f"{ATOM}title"))
            abstract = clean_text(entry.findtext(f"{ATOM}summary"))
            if title and abstract:
                rows.append({"title": title, "abstract": abstract, "tag": category})
        time.sleep(pause)
    return rows


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("data/arxiv_broad_sample.csv"))
    parser.add_argument("--per-category", type=int, default=120)
    parser.add_argument("--pause", type=float, default=3.0, help="Пауза между запросами, чтобы уважать arXiv API.")
    parser.add_argument("--categories", nargs="+", default=DEFAULT_CATEGORIES)
    args = parser.parse_args()

    all_rows = []
    for category in args.categories:
        print(f"Fetching {category}...")
        all_rows.extend(fetch_category(category, args.per_category, args.pause))

    df = pd.DataFrame(all_rows).drop_duplicates(subset=["title"])
    args.output.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(args.output, index=False)
    print(f"Saved {len(df)} rows to {args.output}")


if __name__ == "__main__":
    main()