#!/usr/bin/env python3
"""Paper research helper for ML Intern Codex.

Emulates the useful parts of upstream ml-intern's hf_papers tool with public
HTTP APIs: Hugging Face Papers, arXiv/ar5iv HTML, and Semantic Scholar.
"""

from __future__ import annotations

import argparse
import html.parser
import json
import os
import re
import sys
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any


HF_API = "https://huggingface.co/api"
ARXIV_HTML = "https://arxiv.org/html"
AR5IV_HTML = "https://ar5iv.labs.arxiv.org/html"
S2_API = "https://api.semanticscholar.org"
MAX_SECTION_TEXT_LEN = 8000


def request_json(url: str, params: dict[str, Any] | None = None, method: str = "GET", body: dict[str, Any] | None = None) -> Any:
    if params:
        url = f"{url}?{urllib.parse.urlencode({k: v for k, v in params.items() if v is not None})}"
    data = None
    headers = {"User-Agent": "ml-intern-codex/0.1"}
    if body is not None:
        data = json.dumps(body).encode("utf-8")
        headers["Content-Type"] = "application/json"
    s2_key = os.environ.get("S2_API_KEY")
    if s2_key and url.startswith(S2_API):
        headers["x-api-key"] = s2_key
    request = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(request, timeout=30) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        text = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {text[:500]}") from exc


def request_text(url: str) -> str:
    request = urllib.request.Request(url, headers={"User-Agent": "ml-intern-codex/0.1"})
    with urllib.request.urlopen(request, timeout=30) as response:
        return response.read().decode("utf-8", errors="replace")


def arxiv_s2_id(arxiv_id: str) -> str:
    return f"ARXIV:{arxiv_id}"


def truncate(text: str, limit: int) -> str:
    text = re.sub(r"\s+", " ", text).strip()
    return text if len(text) <= limit else text[:limit].rstrip() + "..."


def paper_arxiv_id(paper: dict[str, Any]) -> str:
    external = paper.get("externalIds") or paper.get("external_ids") or {}
    return external.get("ArXiv") or paper.get("arxiv_id") or paper.get("id", "")


def format_hf_paper(paper: dict[str, Any], idx: int) -> str:
    nested = paper.get("paper") if isinstance(paper.get("paper"), dict) else paper
    title = nested.get("title") or paper.get("title") or "(untitled)"
    arxiv_id = nested.get("id") or nested.get("arxivId") or paper.get("id") or ""
    summary = nested.get("summary") or nested.get("abstract") or ""
    lines = [f"### {idx}. {title}"]
    if arxiv_id:
        lines.append(f"arxiv_id: {arxiv_id}")
        lines.append(f"https://arxiv.org/abs/{arxiv_id}")
    if nested.get("publishedAt"):
        lines.append(f"Published: {nested['publishedAt']}")
    if nested.get("githubUrl"):
        lines.append(f"GitHub: {nested['githubUrl']}")
    if summary:
        lines.append(truncate(summary, 500))
    return "\n".join(lines)


def format_s2_paper(paper: dict[str, Any], idx: int) -> str:
    title = paper.get("title") or "(untitled)"
    year = paper.get("year") or "?"
    cites = paper.get("citationCount", 0)
    venue = paper.get("venue") or ""
    arxiv_id = paper_arxiv_id(paper)
    tldr = (paper.get("tldr") or {}).get("text", "")
    parts = [f"Year: {year}", f"Citations: {cites}"]
    if venue:
        parts.append(f"Venue: {venue}")
    if arxiv_id:
        parts.append(f"arxiv_id: {arxiv_id}")
    lines = [f"### {idx}. {title}", " | ".join(parts)]
    if arxiv_id:
        lines.append(f"https://arxiv.org/abs/{arxiv_id}")
    if tldr:
        lines.append(f"TL;DR: {tldr}")
    return "\n".join(lines)


class ArxivHTMLParser(html.parser.HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.capture_title = False
        self.capture_abstract = False
        self.capture_heading = False
        self.capture_paragraph = False
        self.title_parts: list[str] = []
        self.abstract_parts: list[str] = []
        self.sections: list[dict[str, Any]] = []
        self.current_heading: list[str] = []
        self.current_paragraph: list[str] = []
        self.current_section: dict[str, Any] | None = None

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        classes = dict(attrs).get("class", "") or ""
        if tag == "h1" and "ltx_title" in classes:
            self.capture_title = True
        elif tag == "div" and "ltx_abstract" in classes:
            self.capture_abstract = True
        elif tag in {"h2", "h3"} and "ltx_title" in classes:
            self.capture_heading = True
            self.current_heading = []
        elif tag == "p":
            self.capture_paragraph = True
            self.current_paragraph = []

    def handle_endtag(self, tag: str) -> None:
        if tag == "h1" and self.capture_title:
            self.capture_title = False
        elif tag == "div" and self.capture_abstract:
            self.capture_abstract = False
        elif tag in {"h2", "h3"} and self.capture_heading:
            heading = truncate(" ".join(self.current_heading), 500)
            section_id = ""
            match = re.match(r"^([A-Z]?\d+(?:\.\d+)*)\s", heading)
            if match:
                section_id = match.group(1)
            self.current_section = {"id": section_id, "title": heading, "text": ""}
            self.sections.append(self.current_section)
            self.capture_heading = False
        elif tag == "p" and self.capture_paragraph:
            paragraph = truncate(" ".join(self.current_paragraph), 4000)
            if paragraph:
                if self.capture_abstract:
                    self.abstract_parts.append(paragraph)
                elif self.current_section is not None:
                    existing = self.current_section["text"]
                    self.current_section["text"] = (existing + "\n\n" + paragraph).strip()
            self.capture_paragraph = False

    def handle_data(self, data: str) -> None:
        text = data.strip()
        if not text:
            return
        if self.capture_title:
            self.title_parts.append(text.removeprefix("Title:"))
        if self.capture_heading:
            self.current_heading.append(text)
        if self.capture_paragraph:
            self.current_paragraph.append(text)


def parse_arxiv_html(html_text: str) -> dict[str, Any]:
    parser = ArxivHTMLParser()
    parser.feed(html_text)
    return {
        "title": truncate(" ".join(parser.title_parts), 500),
        "abstract": truncate(" ".join(parser.abstract_parts), 2000),
        "sections": parser.sections,
    }


def op_trending(args: argparse.Namespace) -> str:
    params: dict[str, Any] = {"limit": args.limit * 3 if args.query else args.limit}
    if args.date:
        params["date"] = args.date
    papers = request_json(f"{HF_API}/daily_papers", params)
    if args.query:
        needle = args.query.lower()
        papers = [
            paper
            for paper in papers
            if needle in json.dumps(paper, ensure_ascii=False).lower()
        ]
    lines = ["# Trending Papers"]
    for idx, paper in enumerate(papers[: args.limit], 1):
        lines.append(format_hf_paper(paper, idx))
        lines.append("")
    return "\n".join(lines)


def op_search(args: argparse.Namespace) -> str:
    if not args.query:
        raise SystemExit("search requires --query")
    use_s2 = any([args.date_from, args.date_to, args.categories, args.min_citations, args.sort_by != "relevance"])
    if use_s2:
        params: dict[str, Any] = {
            "query": args.query,
            "limit": args.limit,
            "fields": "title,externalIds,year,citationCount,tldr,venue,publicationDate",
        }
        if args.date_from or args.date_to:
            params["publicationDateOrYear"] = f"{args.date_from or ''}:{args.date_to or ''}"
        if args.categories:
            params["fieldsOfStudy"] = args.categories
        if args.min_citations:
            params["minCitationCount"] = str(args.min_citations)
        if args.sort_by != "relevance":
            params["sort"] = f"{args.sort_by}:desc"
        data = request_json(f"{S2_API}/graph/v1/paper/search/bulk", params)
        papers = data.get("data", [])
        lines = [f"# Papers matching '{args.query}' (Semantic Scholar)"]
        for idx, paper in enumerate(papers[: args.limit], 1):
            lines.append(format_s2_paper(paper, idx))
            lines.append("")
        return "\n".join(lines)
    papers = request_json(f"{HF_API}/papers/search", {"q": args.query, "limit": args.limit})
    lines = [f"# Papers matching '{args.query}' (Hugging Face Papers)"]
    for idx, paper in enumerate(papers[: args.limit], 1):
        lines.append(format_hf_paper(paper, idx))
        lines.append("")
    return "\n".join(lines)


def op_paper_details(args: argparse.Namespace) -> str:
    if not args.arxiv_id:
        raise SystemExit("paper_details requires --arxiv-id")
    paper = request_json(f"{HF_API}/papers/{args.arxiv_id}")
    lines = [f"# {paper.get('title', args.arxiv_id)}", f"https://huggingface.co/papers/{args.arxiv_id}", f"https://arxiv.org/abs/{args.arxiv_id}"]
    for key in ("publishedAt", "submittedOnDailyAt", "githubUrl"):
        if paper.get(key):
            lines.append(f"{key}: {paper[key]}")
    if paper.get("summary"):
        lines.append("")
        lines.append("## Abstract")
        lines.append(paper["summary"])
    if paper.get("ai_summary"):
        lines.append("")
        lines.append("## AI Summary")
        lines.append(str(paper["ai_summary"]))
    return "\n".join(lines)


def op_read_paper(args: argparse.Namespace) -> str:
    if not args.arxiv_id:
        raise SystemExit("read_paper requires --arxiv-id")
    parsed = None
    for base in (ARXIV_HTML, AR5IV_HTML):
        try:
            parsed = parse_arxiv_html(request_text(f"{base}/{args.arxiv_id}"))
            if parsed["sections"]:
                break
        except Exception:
            parsed = None
    if not parsed or not parsed["sections"]:
        return op_paper_details(args) + f"\n\nHTML sections unavailable. PDF: https://arxiv.org/pdf/{args.arxiv_id}"
    if not args.section:
        lines = [f"# {parsed['title'] or args.arxiv_id}", f"https://arxiv.org/abs/{args.arxiv_id}", "", "## Abstract", parsed["abstract"], "", "## Sections"]
        for section in parsed["sections"]:
            preview = truncate(section.get("text", ""), 280)
            lines.append(f"- {section['title']}: {preview}")
        return "\n".join(lines)
    wanted = args.section.lower()
    selected = None
    for section in parsed["sections"]:
        if section["id"].lower() == wanted or wanted in section["title"].lower():
            selected = section
            break
    if not selected:
        available = "\n".join(f"- {section['title']}" for section in parsed["sections"])
        raise SystemExit(f"section not found. Available sections:\n{available}")
    return "\n".join([
        f"# {selected['title']}",
        f"https://arxiv.org/abs/{args.arxiv_id}",
        "",
        truncate(selected.get("text", ""), MAX_SECTION_TEXT_LEN),
    ])


def format_citation(entry: dict[str, Any]) -> str:
    paper = entry.get("citingPaper") or entry.get("citedPaper") or {}
    title = paper.get("title") or "(untitled)"
    year = paper.get("year") or "?"
    cites = paper.get("citationCount", 0)
    arxiv_id = paper_arxiv_id(paper)
    line = f"- {title} ({year}, {cites} cites)"
    if arxiv_id:
        line += f" arxiv:{arxiv_id}"
    if entry.get("isInfluential"):
        line += " [influential]"
    contexts = entry.get("contexts") or []
    if contexts:
        line += f"\n  > {truncate(contexts[0], 220)}"
    return line


def op_citation_graph(args: argparse.Namespace) -> str:
    if not args.arxiv_id:
        raise SystemExit("citation_graph requires --arxiv-id")
    fields = "title,externalIds,year,citationCount,influentialCitationCount,contexts,intents,isInfluential"
    params = {"fields": fields, "limit": args.limit}
    paths: dict[str, str] = {}
    if args.direction in {"references", "both"}:
        paths["References"] = f"/graph/v1/paper/{arxiv_s2_id(args.arxiv_id)}/references"
    if args.direction in {"citations", "both"}:
        paths["Citations"] = f"/graph/v1/paper/{arxiv_s2_id(args.arxiv_id)}/citations"
    lines = [f"# Citation Graph for {args.arxiv_id}", f"https://arxiv.org/abs/{args.arxiv_id}"]
    with ThreadPoolExecutor(max_workers=2) as pool:
        futures = {pool.submit(request_json, f"{S2_API}{path}", params): name for name, path in paths.items()}
        for future in as_completed(futures):
            name = futures[future]
            lines.append("")
            lines.append(f"## {name}")
            try:
                data = future.result()
                for entry in data.get("data", []):
                    lines.append(format_citation(entry))
            except Exception as exc:
                lines.append(f"Error: {exc}")
    return "\n".join(lines)


def op_resources(args: argparse.Namespace) -> str:
    if not args.arxiv_id:
        raise SystemExit(f"{args.operation} requires --arxiv-id")
    sort = {"downloads": "downloads", "likes": "likes", "trending": "trendingScore"}[args.sort]
    calls: dict[str, tuple[str, dict[str, Any]]] = {}
    if args.operation in {"find_datasets", "find_all_resources"}:
        calls["Datasets"] = (f"{HF_API}/datasets", {"filter": f"arxiv:{args.arxiv_id}", "limit": args.limit, "sort": sort, "direction": -1})
    if args.operation in {"find_models", "find_all_resources"}:
        calls["Models"] = (f"{HF_API}/models", {"filter": f"arxiv:{args.arxiv_id}", "limit": args.limit, "sort": sort, "direction": -1})
    if args.operation in {"find_collections", "find_all_resources"}:
        calls["Collections"] = (f"{HF_API}/collections", {"paper": args.arxiv_id})
    lines = [f"# Resources linked to paper {args.arxiv_id}", f"https://huggingface.co/papers/{args.arxiv_id}"]
    with ThreadPoolExecutor(max_workers=3) as pool:
        futures = {pool.submit(request_json, url, params): name for name, (url, params) in calls.items()}
        for future in as_completed(futures):
            name = futures[future]
            lines.append("")
            lines.append(f"## {name}")
            try:
                items = future.result()
                for item in items[: args.limit]:
                    repo_id = item.get("id") or item.get("modelId") or item.get("slug") or item.get("title")
                    likes = item.get("likes")
                    downloads = item.get("downloads")
                    meta = []
                    if downloads is not None:
                        meta.append(f"downloads={downloads}")
                    if likes is not None:
                        meta.append(f"likes={likes}")
                    lines.append(f"- {repo_id}" + (f" ({', '.join(meta)})" if meta else ""))
            except Exception as exc:
                lines.append(f"Error: {exc}")
    return "\n".join(lines)


def op_snippet_search(args: argparse.Namespace) -> str:
    if not args.query:
        raise SystemExit("snippet_search requires --query")
    params: dict[str, Any] = {"query": args.query, "limit": args.limit, "fields": "title,externalIds,year,citationCount"}
    if args.date_from or args.date_to:
        params["publicationDateOrYear"] = f"{args.date_from or ''}:{args.date_to or ''}"
    if args.categories:
        params["fieldsOfStudy"] = args.categories
    if args.min_citations:
        params["minCitationCount"] = str(args.min_citations)
    data = request_json(f"{S2_API}/graph/v1/snippet/search", params)
    lines = [f"# Snippet Search: {args.query}"]
    for idx, item in enumerate(data.get("data", [])[: args.limit], 1):
        paper = item.get("paper") or {}
        snippet = item.get("snippet") or {}
        lines.append(f"### {idx}. {paper.get('title', '(untitled)')}")
        arxiv_id = paper_arxiv_id(paper)
        if arxiv_id:
            lines.append(f"arxiv:{arxiv_id}")
        if snippet.get("section"):
            lines.append(f"Section: {snippet['section']}")
        if snippet.get("text"):
            lines.append(f"> {truncate(snippet['text'], 400)}")
        lines.append("")
    return "\n".join(lines)


def op_recommend(args: argparse.Namespace) -> str:
    if not args.arxiv_id and not args.positive_ids:
        raise SystemExit("recommend requires --arxiv-id or --positive-ids")
    fields = "title,externalIds,year,citationCount,tldr,venue"
    if args.positive_ids and not args.arxiv_id:
        positive = [arxiv_s2_id(item.strip()) for item in args.positive_ids.split(",") if item.strip()]
        negative = [arxiv_s2_id(item.strip()) for item in args.negative_ids.split(",") if item.strip()]
        data = request_json(
            f"{S2_API}/recommendations/v1/papers/",
            {"fields": fields, "limit": args.limit},
            method="POST",
            body={"positivePaperIds": positive, "negativePaperIds": negative},
        )
    else:
        data = request_json(
            f"{S2_API}/recommendations/v1/papers/forpaper/{arxiv_s2_id(args.arxiv_id)}",
            {"fields": fields, "limit": args.limit, "from": "recent"},
        )
    papers = data.get("recommendedPapers", [])
    lines = ["# Recommended Papers"]
    for idx, paper in enumerate(papers[: args.limit], 1):
        lines.append(format_s2_paper(paper, idx))
        lines.append("")
    return "\n".join(lines)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("operation", choices=[
        "trending",
        "search",
        "paper_details",
        "read_paper",
        "citation_graph",
        "snippet_search",
        "recommend",
        "find_datasets",
        "find_models",
        "find_collections",
        "find_all_resources",
    ])
    parser.add_argument("--query")
    parser.add_argument("--arxiv-id")
    parser.add_argument("--section")
    parser.add_argument("--direction", choices=["citations", "references", "both"], default="both")
    parser.add_argument("--date")
    parser.add_argument("--date-from", default="")
    parser.add_argument("--date-to", default="")
    parser.add_argument("--categories")
    parser.add_argument("--min-citations", type=int)
    parser.add_argument("--sort-by", choices=["relevance", "citationCount", "publicationDate"], default="relevance")
    parser.add_argument("--positive-ids", default="")
    parser.add_argument("--negative-ids", default="")
    parser.add_argument("--sort", choices=["downloads", "likes", "trending"], default="downloads")
    parser.add_argument("--limit", type=int, default=10)
    return parser


def main() -> int:
    args = build_parser().parse_args()
    args.limit = min(max(args.limit, 1), 50)
    handlers = {
        "trending": op_trending,
        "search": op_search,
        "paper_details": op_paper_details,
        "read_paper": op_read_paper,
        "citation_graph": op_citation_graph,
        "snippet_search": op_snippet_search,
        "recommend": op_recommend,
        "find_datasets": op_resources,
        "find_models": op_resources,
        "find_collections": op_resources,
        "find_all_resources": op_resources,
    }
    try:
        print(handlers[args.operation](args))
    except Exception as exc:
        print(f"Error running papers {args.operation}: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())