"""arXiv + GitHub 检索工具（支持 HTTP / SOCKS5 代理）

用法：
    # 1. 直连
    python3 research/search_info.py

    # 2. HTTP 代理（推荐：6984 同时支持 http/socks5）
    HTTPS_PROXY=http://127.0.0.1:6984 HTTP_PROXY=http://127.0.0.1:6984 \
        python3 research/search_info.py

    # 3. SOCKS5 代理（需 pip install pysocks）
    ALL_PROXY=socks5h://127.0.0.1:6984 python3 research/search_info.py

    # 4. 命令行指定（覆盖环境变量）
    python3 research/search_info.py --proxy http://127.0.0.1:6984

输出：
    research/search_results.json — 结构化 JSON（按主题分组）
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from typing import Any


# ──────────────────────────────────────────────────────────────────────
# 主题化查询（与 plans/camera/zed2i_iterative_framework.md 五条主线对齐）
# ──────────────────────────────────────────────────────────────────────
TOPICS: dict[str, str] = {
    "stereo_matching": 'all:"stereo matching" AND (all:"foundation" OR all:"zero-shot" OR all:"transformer")',
    "visual_inertial_slam": 'all:("visual inertial" OR "VIO" OR "VI-SLAM") AND cat:cs.CV',
    "gaussian_splatting_slam": 'all:"gaussian splatting" AND all:"SLAM"',
    "monocular_depth_foundation": 'all:("depth anything" OR "metric3d" OR "marigold" OR "UniDepth")',
    "indoor_rgbd_dataset": 'all:("RGB-D dataset" OR "indoor dataset") AND all:"reconstruction"',
    "world_model_video": 'all:"world model" AND (all:"video" OR all:"prediction" OR all:"generative")',
    "zed_camera": 'all:("ZED 2i" OR "ZED stereo" OR "ZED-2i" OR "Stereolabs")',
    "orbbec_gemini": 'all:("Orbbec" OR "Femto Mega" OR "Femto Bolt" OR "Azure Kinect") AND all:("depth" OR "RGB-D" OR "SLAM")',
    "rgbd_indoor_reconstruction": 'all:("RGB-D" OR "RGBD") AND all:("indoor" OR "room") AND (all:"reconstruction" OR all:"SLAM")',
    "neural_stereo_depth": 'all:("foundation stereo" OR "RAFT-Stereo" OR "IGEV") AND all:"depth"',
}

# GitHub 主题化查询
GITHUB_TOPICS: dict[str, str] = {
    "stereo_matching": "stereo matching depth in:name,description,readme stars:>500",
    "vio_slam": "visual inertial SLAM in:name,description,readme stars:>500",
    "gaussian_splatting": "gaussian splatting in:name,description,readme stars:>1000",
    "monocular_depth": "monocular depth estimation in:name,description,readme stars:>1000",
    "world_model": "world model in:name,description,readme stars:>500",
}


# ──────────────────────────────────────────────────────────────────────
# 代理配置
# ──────────────────────────────────────────────────────────────────────
def build_opener(proxy: str | None) -> urllib.request.OpenerDirector:
    """构建带代理的 urllib opener。

    proxy 形式：
      - http://127.0.0.1:6984
      - https://127.0.0.1:6984
      - socks5h://127.0.0.1:6984 （需 pysocks）
      - None  → 走环境变量 (HTTP_PROXY/HTTPS_PROXY/ALL_PROXY) 或直连
    """
    if proxy and proxy.startswith(("socks5://", "socks5h://", "socks4://")):
        try:
            import socks  # type: ignore
            import socket

            scheme, rest = proxy.split("://", 1)
            host, port_str = rest.split(":")
            port = int(port_str)
            socks_type = {
                "socks5": socks.SOCKS5,
                "socks5h": socks.SOCKS5,
                "socks4": socks.SOCKS4,
            }[scheme]
            socks.set_default_proxy(socks_type, host, port, rdns=(scheme == "socks5h"))
            socket.socket = socks.socksocket  # type: ignore[assignment]
            print(f"[proxy] SOCKS5 启用: {proxy}", file=sys.stderr)
            return urllib.request.build_opener()
        except ImportError:
            print(
                "[proxy] 警告: 需要 SOCKS5 但未安装 pysocks，将回退到环境变量。\n"
                "        修复: pip install pysocks",
                file=sys.stderr,
            )
            return urllib.request.build_opener()

    if proxy:
        handler = urllib.request.ProxyHandler({"http": proxy, "https": proxy})
        print(f"[proxy] HTTP(S) 代理启用: {proxy}", file=sys.stderr)
        return urllib.request.build_opener(handler)

    # 走环境变量
    env_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("HTTP_PROXY") or os.environ.get("ALL_PROXY")
    if env_proxy:
        print(f"[proxy] 使用环境变量代理: {env_proxy}", file=sys.stderr)
    else:
        print("[proxy] 直连模式（无代理）", file=sys.stderr)
    return urllib.request.build_opener()


# ──────────────────────────────────────────────────────────────────────
# arXiv
# ──────────────────────────────────────────────────────────────────────
def search_arxiv(
    opener: urllib.request.OpenerDirector,
    query: str,
    max_results: int = 20,
    sort_by: str = "submittedDate",
    sort_order: str = "descending",
) -> list[dict[str, Any]]:
    """调用 arXiv API 并解析返回的 ATOM XML。"""
    base = "https://export.arxiv.org/api/query"
    params = urllib.parse.urlencode(
        {
            "search_query": query,
            "start": 0,
            "max_results": max_results,
            "sortBy": sort_by,
            "sortOrder": sort_order,
        }
    )
    url = f"{base}?{params}"
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 (research-bot)"})

    # arXiv 限速重试（429 → 指数退避）
    data = ""
    for attempt in range(4):
        try:
            with opener.open(req, timeout=30) as resp:
                data = resp.read().decode("utf-8")
            break
        except urllib.error.HTTPError as e:
            if e.code == 429:
                wait = 10 * (attempt + 1)
                print(
                    f"[arxiv][{query[:40]}...] 429 限速，等待 {wait}s 重试 ({attempt+1}/4)",
                    file=sys.stderr,
                )
                time.sleep(wait)
                continue
            print(f"[arxiv][{query[:40]}...] HTTPError {e.code}: {e.reason}", file=sys.stderr)
            return []
        except (urllib.error.URLError, TimeoutError) as e:
            print(f"[arxiv][{query[:40]}...] ERROR: {e}", file=sys.stderr)
            return []
    else:
        print(f"[arxiv][{query[:40]}...] 持续 429，放弃", file=sys.stderr)
        return []

    try:
        root = ET.fromstring(data)
    except ET.ParseError as e:
        print(f"[arxiv][{query[:40]}...] XML parse error: {e}", file=sys.stderr)
        return []

    ns = {"atom": "http://www.w3.org/2005/Atom", "arxiv": "http://arxiv.org/schemas/atom"}
    papers: list[dict[str, Any]] = []

    for entry in root.findall("atom:entry", ns):

        def _text(elem_path: str) -> str:
            elem = entry.find(elem_path, ns)
            return (elem.text or "").replace("\n", " ").strip() if elem is not None else ""

        arxiv_id_url = _text("atom:id")
        arxiv_id = arxiv_id_url.rsplit("/", 1)[-1] if arxiv_id_url else ""
        # 去掉版本号 e.g. 2501.09898v2 → 2501.09898
        arxiv_id_clean = arxiv_id.split("v")[0] if arxiv_id else ""

        authors: list[str] = []
        for a in entry.findall("atom:author", ns):
            name_elem = a.find("atom:name", ns)
            if name_elem is not None and name_elem.text:
                authors.append(name_elem.text.strip())

        categories = [
            c.attrib.get("term", "") for c in entry.findall("atom:category", ns)
        ]

        summary = _text("atom:summary")
        papers.append(
            {
                "arxiv_id": arxiv_id_clean,
                "url": f"https://arxiv.org/abs/{arxiv_id_clean}" if arxiv_id_clean else arxiv_id_url,
                "title": _text("atom:title"),
                "published": _text("atom:published"),
                "updated": _text("atom:updated"),
                "authors": authors,
                "categories": categories,
                "summary": summary[:400] + ("..." if len(summary) > 400 else ""),
            }
        )

    return papers


# ──────────────────────────────────────────────────────────────────────
# GitHub
# ──────────────────────────────────────────────────────────────────────
def search_github(
    opener: urllib.request.OpenerDirector,
    query: str,
    per_page: int = 10,
) -> list[dict[str, Any]]:
    base = "https://api.github.com/search/repositories"
    params = urllib.parse.urlencode({"q": query, "sort": "stars", "order": "desc", "per_page": per_page})
    url = f"{base}?{params}"

    headers = {"User-Agent": "Mozilla/5.0 (research-bot)", "Accept": "application/vnd.github+json"}
    token = os.environ.get("GITHUB_TOKEN")
    if token:
        headers["Authorization"] = f"Bearer {token}"

    req = urllib.request.Request(url, headers=headers)

    try:
        with opener.open(req, timeout=30) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        # 60 req/hr 限速时 GitHub 返回 403
        print(f"[github][{query[:40]}...] HTTPError {e.code}: {e.reason}", file=sys.stderr)
        return []
    except (urllib.error.URLError, TimeoutError) as e:
        print(f"[github][{query[:40]}...] ERROR: {e}", file=sys.stderr)
        return []

    repos: list[dict[str, Any]] = []
    for item in data.get("items", []):
        repos.append(
            {
                "name": item.get("full_name", ""),
                "description": item.get("description", "") or "",
                "stars": item.get("stargazers_count", 0),
                "forks": item.get("forks_count", 0),
                "language": item.get("language", "") or "",
                "url": item.get("html_url", ""),
                "updated_at": item.get("updated_at", ""),
            }
        )
    return repos


# ──────────────────────────────────────────────────────────────────────
# 主入口
# ──────────────────────────────────────────────────────────────────────
def main() -> int:
    parser = argparse.ArgumentParser(description="arXiv + GitHub 主题化检索")
    parser.add_argument("--proxy", default=None, help="代理 URL，如 http://127.0.0.1:6984")
    parser.add_argument("--max-results", type=int, default=15, help="每个主题最多论文数")
    parser.add_argument("--out", default="research/data/search_results.json", help="输出 JSON 路径")
    parser.add_argument(
        "--topics",
        nargs="+",
        default=None,
        help="只跑指定主题（默认全部），如：--topics stereo_matching world_model_video",
    )
    parser.add_argument("--no-github", action="store_true", help="跳过 GitHub 检索")
    parser.add_argument("--delay", type=float, default=3.0, help="主题间延迟（秒），arXiv 建议 ≥3s")
    args = parser.parse_args()

    opener = build_opener(args.proxy)

    selected = args.topics or list(TOPICS.keys())
    results: dict[str, Any] = {
        "meta": {
            "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
            "proxy": args.proxy or os.environ.get("HTTPS_PROXY") or "direct",
            "max_results_per_topic": args.max_results,
        },
        "arxiv": {},
        "github": {},
    }

    # arXiv
    for topic in selected:
        if topic not in TOPICS:
            print(f"[skip] 未知主题: {topic}", file=sys.stderr)
            continue
        print(f"[arxiv] 检索主题: {topic}", file=sys.stderr)
        papers = search_arxiv(opener, TOPICS[topic], max_results=args.max_results)
        results["arxiv"][topic] = papers
        print(f"        → {len(papers)} 篇", file=sys.stderr)
        time.sleep(args.delay)  # arXiv API rate limit

    # GitHub
    if not args.no_github:
        for topic, q in GITHUB_TOPICS.items():
            if args.topics and topic not in args.topics:
                continue
            print(f"[github] 检索主题: {topic}", file=sys.stderr)
            repos = search_github(opener, q, per_page=10)
            results["github"][topic] = repos
            print(f"         → {len(repos)} 个仓库", file=sys.stderr)
            time.sleep(1.0)

    out_path = args.out
    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    total_papers = sum(len(v) for v in results["arxiv"].values())
    total_repos = sum(len(v) for v in results["github"].values())
    print(f"\n✅ 完成：{total_papers} 篇 arXiv + {total_repos} 个 GitHub 仓库 → {out_path}", file=sys.stderr)
    return 0


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