#!/usr/bin/env python3
"""Collect recent arXiv papers related to LLM/AI agents."""

from __future__ import annotations

import argparse
import json
import re
import time
import urllib.parse
import urllib.request
import urllib.error
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",
    "arxiv": "http://arxiv.org/schemas/atom",
}


QUERIES = [
    ("llm-agent", 'all:"LLM agent" OR all:"LLM agents"'),
    ("language-agent", 'all:"language agent" OR all:"language agents"'),
    ("ai-agent", 'all:"AI agent" OR all:"AI agents"'),
    ("agentic-ai", 'all:"agentic AI" OR all:"agentic workflow"'),
    ("agent-evaluation", 'all:"agent evaluation" OR all:"agent benchmark" OR all:"agentic benchmark"'),
    ("agent-memory", 'all:"agent memory" OR all:"memory agent" OR all:"memory agents"'),
    ("tool-use", 'all:"tool use" AND all:"agent"'),
    ("function-calling", 'all:"function calling" AND all:"agent"'),
    ("coding-agent", 'all:"coding agent" OR all:"software engineering agent" OR all:"SWE-bench"'),
    ("web-gui-agent", 'all:"web agent" OR all:"browser agent" OR all:"GUI agent" OR all:"computer use"'),
    ("multi-agent-llm", 'all:"multi-agent" AND (all:"LLM" OR all:"large language model")'),
    ("agent-safety", 'all:"agent safety" OR all:"AI agent safety" OR all:"agent security"'),
    ("rag-agent", 'all:"RAG" AND all:"agent"'),
    ("planning-agent", 'all:"planning" AND all:"LLM agent"'),
    ("autonomous-agent-llm", 'all:"autonomous agent" AND (all:"LLM" OR all:"large language model")'),
]


TOPIC_RULES = [
    ("agent-evaluation", ("evaluation", "benchmark", "eval", "metric", "leaderboard", "assessment")),
    ("memory", ("memory", "remember", "stateful", "long-term", "episodic")),
    ("tool-use", ("tool", "function calling", "api", "mcp", "action")),
    ("coding-agent", ("coding", "software engineering", "swe-bench", "repository", "program repair", "code agent")),
    ("computer-use", ("computer use", "gui", "browser", "web agent", "mobile", "desktop")),
    ("multi-agent", ("multi-agent", "multiagent", "agent society", "collaboration", "debate")),
    ("agent-safety", ("safety", "security", "risk", "prompt injection", "alignment", "guardrail")),
    ("rag", ("retrieval", "rag", "knowledge base", "grounding")),
    ("planning", ("planning", "planner", "plan", "long-horizon", "task decomposition")),
    ("reasoning", ("reasoning", "reflection", "self-improvement", "verifier")),
    ("workflow-agent", ("workflow", "enterprise", "office", "productivity", "automation")),
    ("embodied-agent", ("embodied", "robot", "robotic", "navigation")),
    ("world-model", ("world model", "simulation", "environment model")),
]


HIGH_SIGNAL_TERMS = (
    "llm agent",
    "language agent",
    "ai agent",
    "agentic",
    "tool use",
    "memory",
    "benchmark",
    "evaluation",
    "coding agent",
    "swe-bench",
    "computer use",
    "gui agent",
    "multi-agent",
    "agent safety",
)


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


def slugify(text: str) -> str:
    text = text.lower()
    text = re.sub(r"[^a-z0-9]+", "-", text)
    return re.sub(r"-+", "-", text).strip("-")[:80] or "untitled"


def yaml_value(value: str | int | None) -> str:
    if value is None:
        return ""
    text = str(value).replace("\n", " ").strip()
    if not text:
        return ""
    if any(char in text for char in [":", "#", "[", "]", "{", "}", "\"", "'"]):
        return json.dumps(text, ensure_ascii=False)
    return text


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


def date_window_query(raw_query: str, from_date: str, to_date: str) -> str:
    start = from_date.replace("-", "") + "0000"
    end = to_date.replace("-", "") + "2359"
    return f"({raw_query}) AND submittedDate:[{start} TO {end}]"


def fetch_query(
    raw_query: str,
    from_date: str,
    to_date: str,
    start: int,
    max_results: int,
    retries: int,
    retry_sleep: float,
) -> str:
    params = {
        "search_query": date_window_query(raw_query, from_date, to_date),
        "start": start,
        "max_results": max_results,
        "sortBy": "submittedDate",
        "sortOrder": "descending",
    }
    url = f"{API_URL}?{urllib.parse.urlencode(params)}"
    request = urllib.request.Request(url, headers={"User-Agent": "agent-kb-arxiv-collector/0.1"})
    for attempt in range(retries + 1):
        try:
            with urllib.request.urlopen(request, timeout=30) as response:
                return response.read().decode("utf-8", "replace")
        except urllib.error.HTTPError as exc:
            if exc.code != 429 or attempt >= retries:
                raise
            time.sleep(retry_sleep * (attempt + 1))
        except urllib.error.URLError:
            if attempt >= retries:
                raise
            time.sleep(retry_sleep * (attempt + 1))
    raise RuntimeError("unreachable fetch retry state")


def manifest_record(record: dict, topics: list[str], score: int, relevance: str, matched_queries: set[str]) -> dict:
    sorted_queries = sorted(matched_queries)
    return {
        "id": record["id"],
        "arxiv_id": record["id"],
        "source": "arxiv",
        "source_id": f"arxiv:{record['id']}",
        "title": record["title"],
        "url": record["url"],
        "pdf_url": f"https://arxiv.org/pdf/{record['id']}",
        "published": record["published"],
        "updated": record["updated"],
        "authors": record["authors"],
        "categories": record["categories"],
        "topics": topics,
        "score": score,
        "relevance": relevance,
        "primary_query": sorted_queries[0] if sorted_queries else "",
        "matched_queries": sorted_queries,
    }


def parse_feed(xml_text: str) -> list[dict]:
    root = ET.fromstring(xml_text)
    records = []
    for entry in root.findall("atom:entry", NS):
        entry_id = entry.findtext("atom:id", default="", namespaces=NS)
        title = normalize_space(entry.findtext("atom:title", default="", namespaces=NS))
        summary = normalize_space(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_space(author.findtext("atom:name", default="", namespaces=NS))
            for author in entry.findall("atom:author", NS)
        ]
        categories = [
            category.attrib.get("term", "")
            for category in entry.findall("atom:category", NS)
            if category.attrib.get("term")
        ]
        records.append(
            {
                "id": arxiv_id(entry_id),
                "url": f"https://arxiv.org/abs/{arxiv_id(entry_id)}",
                "title": title,
                "summary": summary,
                "published": published,
                "updated": updated,
                "authors": authors,
                "categories": categories,
            }
        )
    return records


def classify(record: dict, matched_queries: set[str]) -> tuple[list[str], int, str]:
    text = f"{record['title']} {record['summary']}".lower()
    title = record["title"].lower()
    topics = []
    score = 0

    if "agent" in title or "agentic" in title:
        score += 4
    elif "agent" in text or "agentic" in text:
        score += 2

    if "llm" in text or "large language model" in text or "language model" in text:
        score += 2

    for term in HIGH_SIGNAL_TERMS:
        if term in title:
            score += 3
        elif term in text:
            score += 1

    for topic, keywords in TOPIC_RULES:
        if any(keyword in text for keyword in keywords):
            topics.append(topic)
            score += 1

    if matched_queries:
        score += min(len(matched_queries), 4)

    if not topics and ("agent" in text or "agentic" in text):
        topics.append("agent")

    if score >= 13:
        relevance = "high"
    elif score >= 7:
        relevance = "medium"
    else:
        relevance = "low"

    return sorted(set(topics)), score, relevance


def existing_arxiv_ids() -> set[str]:
    ids = set()
    for path in (ROOT / "papers" / "items").glob("*.md"):
        text = path.read_text(encoding="utf-8")
        for match in re.findall(r"https://arxiv\.org/(?:abs|html)/([0-9]+\.[0-9]+)", text):
            ids.add(match)
    return ids


def item_path(record: dict) -> Path:
    year = (record.get("published") or "0000")[:4]
    title_slug = slugify(record["title"])
    id_slug = record["id"].replace(".", "-")
    return ROOT / "papers" / "items" / f"{year}-{id_slug}-{title_slug}.md"


def write_item(record: dict, topics: list[str], score: int, relevance: str, matched_queries: set[str], today: str) -> Path:
    path = item_path(record)
    authors = ", ".join(record["authors"][:8])
    if len(record["authors"]) > 8:
        authors += ", et al."
    categories = ", ".join(record["categories"])
    query_list = ", ".join(sorted(matched_queries))
    topic_block = "\n".join(f"  - {topic}" for topic in topics) or "  - agent"
    category_block = "\n".join(f"  - {category}" for category in record["categories"]) or "  -"
    content = f"""# Paper: {record['title']}

---
type: paper
title: {yaml_value(record['title'])}
authors: {yaml_value(authors)}
year: {yaml_value((record.get('published') or '')[:4])}
venue: arXiv
url: {record['url']}
code_url:
source: arxiv
collected_at: {today}
published_at: {record.get('published') or ''}
updated_at: {record.get('updated') or ''}
status: queued
relevance: {relevance}
topics:
{topic_block}
methods:
  -
benchmarks:
  -
models:
  -
datasets:
{category_block}
related_concepts:
  -
related_jobs:
  -
related_experiments:
  -
related_projects:
  -
collection_score: {score}
collection_queries: {yaml_value(query_list)}
---

## One-line Takeaway

Auto-collected from arXiv because it matched the Agent collection queries. Needs human skim.

## Why Collected

- matched queries: {query_list or 'agent-related arXiv sweep'}
- inferred topics: {', '.join(topics) or 'agent'}
- arXiv categories: {categories or 'unknown'}
- collection score: {score}

## Review Checklist

- Does this paper directly inform Agent architecture, evaluation, memory, tools, safety, coding agents, GUI/browser agents, or multi-agent workflows?
- Does it include a benchmark, dataset, code, or reproducible experimental setup?
- Should it be promoted from `queued` to `skimmed` or `summarized`?

## Links

- arXiv: {record['url']}
"""
    path.write_text(content, encoding="utf-8")
    return path


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--from-date", default="2025-07-08")
    parser.add_argument("--to-date", default=date.today().isoformat())
    parser.add_argument("--per-query", type=int, default=120)
    parser.add_argument("--page-size", type=int, default=60)
    parser.add_argument("--max-items", type=int, default=180)
    parser.add_argument("--min-score", type=int, default=6)
    parser.add_argument("--sleep", type=float, default=1.0)
    parser.add_argument("--retries", type=int, default=4)
    parser.add_argument("--retry-sleep", type=float, default=10.0)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    by_id: dict[str, dict] = {}
    matches: dict[str, set[str]] = defaultdict(set)

    for label, raw_query in QUERIES:
        fetched = 0
        start = 0
        while fetched < args.per_query:
            page_size = min(args.page_size, args.per_query - fetched)
            xml_text = fetch_query(
                raw_query,
                args.from_date,
                args.to_date,
                start,
                page_size,
                args.retries,
                args.retry_sleep,
            )
            records = parse_feed(xml_text)
            if not records:
                break
            for record in records:
                by_id.setdefault(record["id"], record)
                matches[record["id"]].add(label)
            fetched += len(records)
            start += len(records)
            if len(records) < page_size:
                break
            time.sleep(args.sleep)

    existing = existing_arxiv_ids()
    scored = []
    all_candidates = []
    for record_id, record in by_id.items():
        topics, score, relevance = classify(record, matches[record_id])
        all_candidates.append(manifest_record(record, topics, score, relevance, matches[record["id"]]))
        if score < args.min_score:
            continue
        scored.append((score, relevance, topics, record))

    scored.sort(key=lambda item: (item[0], item[3].get("published") or ""), reverse=True)
    selected = scored[: args.max_items]

    data_dir = ROOT / "data"
    data_dir.mkdir(exist_ok=True)
    manifest = []
    written = []
    skipped_existing = 0
    today = date.today().isoformat()
    for score, relevance, topics, record in selected:
        manifest.append(manifest_record(record, topics, score, relevance, matches[record["id"]]))
        if record["id"] in existing:
            skipped_existing += 1
            continue
        if not args.dry_run:
            path = write_item(record, topics, score, relevance, matches[record["id"]], today)
            written.append(str(path.relative_to(ROOT)))

    all_candidates.sort(key=lambda item: (item["score"], item["published"]), reverse=True)
    candidates_path = data_dir / f"arxiv-agent-candidates-{args.from_date}-to-{args.to_date}.json"
    candidates_path.write_text(json.dumps(all_candidates, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    manifest_path = data_dir / f"arxiv-agent-papers-{args.from_date}-to-{args.to_date}.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    summary = {
        "from_date": args.from_date,
        "to_date": args.to_date,
        "queries": [label for label, _ in QUERIES],
        "unique_seen": len(by_id),
        "selected": len(selected),
        "written": len(written),
        "skipped_existing": skipped_existing,
        "candidates_manifest": str(candidates_path.relative_to(ROOT)),
        "manifest": str(manifest_path.relative_to(ROOT)),
        "written_paths": written,
    }
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


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