#!/usr/bin/env python3
"""Collect a paginated, abstract-level Agent-memory research corpus from arXiv."""

from __future__ import annotations

import argparse
import json
import re
import time
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from collections import defaultdict
from datetime import date
from pathlib import Path


ROOT = Path(__file__).resolve().parents[2]
API_URL = "https://export.arxiv.org/api/query"
NS = {
    "atom": "http://www.w3.org/2005/Atom",
    "opensearch": "http://a9.com/-/spec/opensearch/1.1/",
}

QUERIES = [
    ("agent-memory", 'all:"agent memory" OR all:"memory agent" OR all:"memory agents"'),
    (
        "long-term-memory",
        '(all:"long-term memory" OR all:"persistent memory") AND '
        '(all:"LLM agent" OR all:"language agent" OR all:"large language model agent")',
    ),
    (
        "episodic-memory",
        'all:"episodic memory" AND (all:"LLM" OR all:"large language model") AND all:"agent"',
    ),
    (
        "procedural-memory",
        '(all:"procedural memory" OR all:"skill memory") AND '
        '(all:"LLM agent" OR all:"language agent")',
    ),
    (
        "memory-evaluation",
        '(all:"memory benchmark" OR all:"memory evaluation") AND '
        '(all:"LLM agent" OR all:"language agent" OR all:"agent memory")',
    ),
]


def normalize(value: str) -> str:
    return re.sub(r"\s+", " ", value).strip()


def arxiv_id(raw: str) -> str:
    return raw.rsplit("/", 1)[-1].split("v", 1)[0]


def date_query(query: str, start: str, end: str) -> str:
    return (
        f"({query}) AND submittedDate:[{start.replace('-', '')}0000 "
        f"TO {end.replace('-', '')}2359]"
    )


def fetch_page(
    query: str,
    start_date: str,
    end_date: str,
    offset: int,
    page_size: int,
    retries: int,
    timeout: int,
) -> bytes:
    params = {
        "search_query": date_query(query, start_date, end_date),
        "start": offset,
        "max_results": page_size,
        "sortBy": "submittedDate",
        "sortOrder": "ascending",
    }
    request = urllib.request.Request(
        f"{API_URL}?{urllib.parse.urlencode(params)}",
        headers={"User-Agent": "agent-kb-research/0.1"},
    )
    for attempt in range(retries + 1):
        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                return response.read()
        except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError):
            if attempt >= retries:
                raise
            time.sleep(2 ** attempt)
    raise RuntimeError("unreachable retry state")


def parse_page(payload: bytes) -> tuple[int, list[dict]]:
    root = ET.fromstring(payload)
    total = int(root.findtext("opensearch:totalResults", default="0", namespaces=NS))
    records = []
    for entry in root.findall("atom:entry", NS):
        paper_id = arxiv_id(entry.findtext("atom:id", default="", namespaces=NS))
        records.append(
            {
                "arxiv_id": paper_id,
                "title": normalize(entry.findtext("atom:title", default="", namespaces=NS)),
                "abstract": normalize(entry.findtext("atom:summary", default="", namespaces=NS)),
                "published": entry.findtext("atom:published", default="", namespaces=NS)[:10],
                "updated": entry.findtext("atom:updated", default="", namespaces=NS)[:10],
                "authors": [
                    normalize(author.findtext("atom:name", default="", namespaces=NS))
                    for author in entry.findall("atom:author", NS)
                ],
                "categories": [
                    node.attrib["term"]
                    for node in entry.findall("atom:category", NS)
                    if node.attrib.get("term")
                ],
                "url": f"https://arxiv.org/abs/{paper_id}",
                "pdf_url": f"https://arxiv.org/pdf/{paper_id}",
            }
        )
    return total, records


def write_json(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    tmp.replace(path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--from-date", default="2025-01-01")
    parser.add_argument("--to-date", default=date.today().isoformat())
    parser.add_argument("--page-size", type=int, default=100)
    parser.add_argument("--max-per-query", type=int, default=2000)
    parser.add_argument("--sleep", type=float, default=1.0)
    parser.add_argument("--retries", type=int, default=4)
    parser.add_argument("--timeout", type=int, default=45)
    parser.add_argument(
        "--output",
        type=Path,
        default=ROOT / "data" / "research" / "agent-memory-expanded.json",
    )
    args = parser.parse_args()

    records: dict[str, dict] = {}
    matched_queries: defaultdict[str, set[str]] = defaultdict(set)
    query_totals = {}
    for name, query in QUERIES:
        offset = 0
        while offset < args.max_per_query:
            payload = fetch_page(
                query,
                args.from_date,
                args.to_date,
                offset,
                args.page_size,
                args.retries,
                args.timeout,
            )
            total, page = parse_page(payload)
            query_totals[name] = total
            for record in page:
                records[record["arxiv_id"]] = record
                matched_queries[record["arxiv_id"]].add(name)
            offset += len(page)
            print(f"{name}: {min(offset, total)}/{total}", flush=True)
            if not page or offset >= total:
                break
            time.sleep(args.sleep)
        time.sleep(args.sleep)

    papers = [
        {**record, "matched_queries": sorted(matched_queries[paper_id])}
        for paper_id, record in records.items()
    ]
    papers.sort(key=lambda item: (item["published"], item["arxiv_id"]))
    output = {
        "generated_at": date.today().isoformat(),
        "from_date": args.from_date,
        "to_date": args.to_date,
        "query_totals": query_totals,
        "unique_papers": len(papers),
        "papers": papers,
    }
    write_json(args.output, output)
    print(json.dumps({"query_totals": query_totals, "unique_papers": len(papers)}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())