#!/usr/bin/env python3
"""
arXiv adapter for AI Research Oracle.
Fetches recent AI/ML papers via arXiv API (Atom XML).
No scraping — uses the official API endpoint.

Rate limits: 1 req/3s (be polite). arXiv enforces this aggressively.

============================================================================
SCORING DESIGN PRINCIPLE (do not violate when copying this to other adapters)
============================================================================
Structural metadata (recency, author count, abstract length, category
diversity) is a WEAK signal present in EVERY paper regardless of topic. It
must NEVER dominate the final score. Content-relevance signals (AI keyword
density in the abstract, methodology-contribution detection) MUST carry the
weight. A paper that merely USES an existing AI tool as incidental methodology
scores lower than one that CONTRIBUTES a new method — detection is by verb
(creative "we propose/introduce" vs evaluative "we evaluate/using"), not by
domain. Applied-domain papers are TAGGED for filtering, NEVER penalized: for a
startup-idea oracle, novel AI applied to a vertical is desired signal.
(See _score() and _methodology_claim_score() for the implementation.)
============================================================================
"""

import json
import os
import re
import time
import urllib.request
import urllib.error
import xml.etree.ElementTree as ET
from datetime import datetime, timedelta, timezone
from html import unescape

from adapters import SourceAdapter, http_get, AdapterHTTPError

# arXiv API
ARXIV_API = "http://export.arxiv.org/api/query"


class ArxivAdapter(SourceAdapter):
    """arXiv API adapter."""

    # Default categories to scan
    DEFAULT_CATEGORIES = ["cs.AI", "cs.LG", "cs.CL"]

    def __init__(self, categories=None, rate_limit=3):
        """
        Args:
            categories: List of arXiv categories. Default: cs.AI, cs.LG, cs.CL
            rate_limit: Seconds between API calls (default 3).
        """
        self.categories = categories or self.DEFAULT_CATEGORIES
        self.rate_limit = rate_limit

    def name(self) -> str:
        return "arxiv"

    def _parse_atom(self, xml_data: str) -> list[dict]:
        """Parse arXiv Atom XML into raw paper dicts."""
        root = ET.fromstring(xml_data)

        # Handle namespaces — arXiv uses multiple
        # Walk all elements and extract what we need
        papers = []

        for entry in root.iter():
            tag = entry.tag.split("}")[-1] if "}" in entry.tag else entry.tag
            if tag == "entry":
                paper = {
                    "id": "",
                    "title": "",
                    "summary": "",
                    "published": "",
                    "updated": "",
                    "authors": [],
                    "categories": [],
                    "comment": "",
                    "journal_ref": "",
                    "doi": "",
                    "link": "",
                }

                for child in entry:
                    ctag = child.tag.split("}")[-1]

                    if ctag == "id":
                        # arXiv ID like http://arxiv.org/abs/cs.AI/2607.00123
                        paper["id"] = child.text or ""
                        # Also extract clean ID
                        if "abs/" in paper["id"]:
                            paper["arxiv_id"] = paper["id"].split("abs/")[-1]

                    elif ctag == "title":
                        paper["title"] = unescape((child.text or "").strip())

                    elif ctag == "summary":
                        paper["summary"] = unescape((child.text or "").strip())

                    elif ctag == "published":
                        paper["published"] = child.text or ""

                    elif ctag == "updated":
                        paper["updated"] = child.text or ""

                    elif ctag == "author":
                        for ac in child:
                            a_tag = ac.tag.split("}")[-1]
                            if a_tag == "name":
                                paper["authors"].append(unescape((ac.text or "").strip()))

                    elif ctag == "category":
                        term = child.get("term", "")
                        if term:
                            paper["categories"].append(term)

                    elif ctag == "arxiv":
                        # arXiv-specific: comment, journal_ref, doi
                        sub_tag = child.tag.split("}")[-1]
                        if sub_tag == "comment":
                            paper["comment"] = unescape((child.text or "").strip())
                        elif sub_tag == "journal_ref":
                            paper["journal_ref"] = unescape((child.text or "").strip())
                        elif sub_tag == "doi":
                            paper["doi"] = unescape((child.text or "").strip())

                    elif ctag == "link":
                        href = child.get("href", "")
                        # Prefer the abstract page link
                        if "abs/" in href and not paper["link"]:
                            paper["link"] = href
                        elif "pdf/" in href:
                            paper["pdf_link"] = href

                # Clean up arxiv_id if not extracted from ID field
                if "arxiv_id" not in paper and "abs/" in paper.get("id", ""):
                    paper["arxiv_id"] = paper["id"].split("abs/")[-1]

                if paper["title"] and paper["summary"]:
                    papers.append(paper)

        return papers

    def _request(self, query: str, max_results: int = 20, sort_by="submittedDate") -> list[dict]:
        """Make an arXiv API request."""
        url = (
            f"{ARXIV_API}"
            f"?search_query={urllib.parse.quote(query)}"
            f"&sortBy={sort_by}"
            f"&sortOrder=descending"
            f"&max_results={max_results}"
        )

        try:
            raw = http_get(url, headers={"User-Agent": "ai-oracle/0.1"},
                           timeout=30, max_retries=2, owner=self)
            return self._parse_atom(raw.decode("utf-8"))
        except AdapterHTTPError as e:
            print(f"    {e.failure_class}: arXiv query ({e})")
            return []
        except Exception as e:
            print(f"    arXiv request error: {e}")
            return []

    def _score(self, paper: dict, age_days: float) -> float:
        """Score based on AI-methodology relevance, not structural metadata.

        arXiv has no upvotes/stars. The key signal is whether the paper is a
        genuine AI/ML *contribution* (new architecture, framework, method) vs.
        merely *using* an existing AI tool as incidental methodology.

        Structural signals (recency, author count, category diversity) are
        weak and capped low so they cannot dominate the score — any paper has
        them regardless of topic. The relevance signals (methodology claim +
        AI keyword density) carry the weight.

        Applied-domain papers (healthcare, finance, biology) are TAGGED for
        filtering but NOT penalized — for a startup-idea oracle, "novel AI
        technique applied to a vertical" is exactly the signal we want.
        """
        # --- Structural signals (weak, capped low) ---
        # Recency: newer = slightly higher, but max 2.0 (was 5.0)
        try:
            pub_str = paper.get("published", "")
            pub_dt = datetime.fromisoformat(pub_str.replace("Z", "+00:00"))
            now = datetime.now(timezone.utc)
            age_hours = (now - pub_dt).total_seconds() / 3600
            recency = max(0, 2.0 - age_hours / 48.0)
        except (ValueError, TypeError):
            recency = max(0, 2.0 - age_days * 0.1)

        # Author count: collaborative work is a weak signal, max 0.5 (was 2.0)
        author_count = len(paper.get("authors", []))
        author_bonus = min(author_count * 0.05, 0.5)

        # Category diversity: cross-domain is mildly interesting, max 0.3
        cats = paper.get("categories", [])
        diversity_bonus = min(len(cats) * 0.1, 0.3)

        # Journal/DOI: published = validated, small bonus
        journal_bonus = 0.5 if paper.get("journal_ref") else 0.0
        doi_bonus = 0.2 if paper.get("doi") else 0.0

        # Comment field (page count etc.): tiny bonus
        comment = paper.get("comment", "")
        comment_bonus = min(len(comment) / 300.0, 0.5) if comment else 0.0

        # --- Relevance signals (carry the weight) ---
        title_lower = paper.get("title", "").lower()
        abstract = paper.get("summary", "")
        abstract_lower = abstract.lower()

        # AI keyword density in abstract (not just title)
        ai_keywords = [
            "agent", "llm", "large language", "gpt", "transformer",
            "neural", "deep learning", "machine learning", "reinforcement",
            "diffusion", "multimodal", "embedding", "attention", "rag",
            "retrieval", "fine-tun", "pretrain", "self-supervised",
            "foundation model", "reasoning", "alignment", "policy gradient",
            "graph neural", "vision-language", "vla", "vlm",
        ]
        ai_kw_hits = sum(1 for kw in ai_keywords if kw in abstract_lower)
        # Density matters more than single mention; saturate at ~6 hits
        ai_keyword_bonus = min(ai_kw_hits * 0.35, 2.0)

        # Title keyword bonus (small)
        title_kw_list = ["agent", "llm", "reasoning", "verif", "multimodal",
                         "embodied", "robot", "alignment", "autonomous"]
        title_kw_hits = sum(1 for kw in title_kw_list if kw in title_lower)
        title_kw_bonus = min(title_kw_hits * 0.2, 0.5)

        # AI methodology claim: does the abstract propose/invent something?
        methodology_bonus = self._methodology_claim_score(abstract_lower)

        # --- Domain tag (no penalty) ---
        combined_text = f"{title_lower} {abstract_lower}"
        applied_domains = {
            "applied:healthcare": ["gastric", "biopsy", "pathology", "clinical",
                "patient", "diagnosis", "medical imaging", "h. pylori",
                "cancer screening", "oncology", "neurology", "cardiology"],
            "applied:finance": ["stock market", "trading", "portfolio", "forex",
                "cryptocurrency", "fintech", "credit scoring", "fraud detection"],
            "applied:biology": ["protein folding", "gene expression", "genome",
                "molecular", "bioinformatics", "cell type", "organism"],
        }
        primary_domain = None
        for domain, domain_keywords in applied_domains.items():
            hits = sum(1 for kw in domain_keywords if kw in combined_text)
            if hits >= 2:  # need at least 2 domain keywords to trigger
                primary_domain = domain
                break
        if primary_domain:
            paper["_applied_domain"] = primary_domain

        # Total: structural (capped ~3.5) + relevance (capped ~5.0)
        structural = recency + author_bonus + diversity_bonus + journal_bonus + doi_bonus + comment_bonus
        relevance = ai_keyword_bonus + title_kw_bonus + methodology_bonus
        score = structural + relevance

        return min(round(score, 2), 10.0)

    def _methodology_claim_score(self, abstract_lower: str) -> float:
        """Detect whether the abstract proposes a NEW AI method vs. using an existing one.

        The core distinction is the VERB: creative verbs (propose/introduce/
        develop/design) signal a contribution; evaluative verbs (evaluate/
        conduct/retrospective/using/based on) signal applying an existing tool.

        Check weak (evaluative) patterns FIRST — if the paper is an evaluation
        or application of an existing system, it gets 0.5 even if it also says
        'we present'. Only if no weak pattern matches do we look for a genuine
        contribution claim.

        Returns 0.0 (no claim / unclear), 0.5 (uses existing tool),
        2.0 (novel method), or 2.5 (explicit proposal of new AI system).
        """
        # Weak FIRST: evaluative / application language
        weak_markers = ["we evaluate", "we conducted", "we conduct", "retrospective",
                        "case study", "pilot study", "using", "we leverage",
                        "we applied", "based on", "we report a", "we present a retrospective",
                        "an evaluation of", "we benchmark", "empirical study"]
        for wm in weak_markers:
            if wm in abstract_lower:
                return 0.5

        # Strong: genuinely creative verbs + AI noun within 60 chars
        proposal_verbs = ["we propose", "we introduce", "we develop", "we design",
                          "we formulate", "we construct", "we build", "we present a novel",
                          "we present a new"]
        ai_nouns = ["model", "architecture", "framework", "method", "approach",
                    "agent", "system", "network", "algorithm", "pipeline",
                    "llm", "transformer", "policy", "graph", "orchestration",
                    "harness", "memory", "solver", "planner"]
        for verb in proposal_verbs:
            idx = abstract_lower.find(verb)
            if idx != -1:
                window = abstract_lower[idx:idx + 60]
                if any(noun in window for noun in ai_nouns):
                    return 2.5

        # Strong-ish: 'novel'/'new' + AI architecture/framework/model within 40 chars
        for marker in ["novel", "new"]:
            idx = abstract_lower.find(marker)
            scan = 0
            while idx != -1 and scan < 5:
                window = abstract_lower[idx:idx + 40]
                if any(noun in window for noun in
                       ["architecture", "framework", "model", "method",
                        "approach", "agent", "system", "harness"]):
                    return 2.0
                idx = abstract_lower.find(marker, idx + 1)
                scan += 1

        return 0.0

    def _tags(self, paper: dict) -> list:
        """Generate category tags from arXiv metadata."""
        tags = ["arxiv"]

        # Add arXiv categories
        for cat in paper.get("categories", [])[:5]:
            tags.append(f"cat:{cat}")

        # Primary signal categories
        cats_lower = [c.lower() for c in paper.get("categories", [])]
        if "cs.ai" in cats_lower:
            tags.append("ai-general")
        if "cs.lg" in cats_lower:
            tags.append("machine-learning")
        if "cs.cl" in cats_lower:
            tags.append("nlp")
        if "cs.cv" in cats_lower:
            tags.append("computer-vision")
        if "cs.ro" in cats_lower:
            tags.append("robotics")
        if "cs.se" in cats_lower:
            tags.append("software-engineering")

        # Check if it's a survey/tutorial
        title_lower = paper.get("title", "").lower()
        summary_lower = paper.get("summary", "").lower()
        if any(kw in summary_lower for kw in ["survey", "tutorial", "overview of", "review of"]):
            tags.append("survey")

        # Check for preprint vs published
        if paper.get("journal_ref"):
            tags.append("published")
        else:
            tags.append("preprint")

        # Applied-domain tag
        if paper.get("_applied_domain"):
            tags.append(paper["_applied_domain"])

        return tags

    def fetch(self, query: str = "", limit: int = 20) -> list[dict]:
        """
        Fetch papers from arXiv.

        If query is empty, fetch recent papers from configured categories.
        If query is provided, search for it.
        """
        now = datetime.now(timezone.utc)
        all_papers = []

        if query:
            print(f"  Searching arXiv: '{query}'")
            papers = self._request(query, max_results=limit, sort_by="submittedDate")
            all_papers.extend(papers)
        else:
            # Fetch from each configured category
            for cat in self.categories:
                q = f"cat:{cat}"
                cat_papers = self._request(q, max_results=limit, sort_by="submittedDate")
                all_papers.extend(cat_papers)
                time.sleep(self.rate_limit)

            # Deduplicate by arxiv_id
            seen = set()
            unique = []
            for p in all_papers:
                pid = p.get("arxiv_id", p.get("id", ""))
                if pid and pid not in seen:
                    seen.add(pid)
                    unique.append(p)
            all_papers = unique

        # Sort by published date (newest first), take top limit
        all_papers.sort(
            key=lambda p: p.get("published", ""),
            reverse=True,
        )
        all_papers = all_papers[:limit]

        entries = []
        for paper in all_papers:
            # Calculate age
            published = paper.get("published", "")
            try:
                pub_dt = datetime.fromisoformat(published.replace("Z", "+00:00"))
                age_days = (now - pub_dt).days
            except (ValueError, TypeError):
                age_days = 0

            score = self._score(paper, age_days)
            tags = self._tags(paper)

            # Source ID: arxiv_id
            source_id = paper.get("arxiv_id", paper.get("id", "")).split("/")[-1]

            # URL: abstract page
            url = paper.get("link", "")
            if not url and paper.get("arxiv_id"):
                url = f"https://arxiv.org/abs/{paper['arxiv_id']}"

            # Title (clean — remove trailing category markers)
            title = re.sub(r"\s*\([A-Za-z0-9., ]*\)\s*$", "", paper.get("title", "")).strip()

            # Extracted text: summary (abstract) — already clean
            extracted_text = paper.get("summary", "")

            # Structured metadata
            raw_meta = {
                "arxiv_id": paper.get("arxiv_id", ""),
                "authors": paper.get("authors", []),
                "author_count": len(paper.get("authors", [])),
                "categories": paper.get("categories", []),
                "published": paper.get("published", ""),
                "updated": paper.get("updated", ""),
                "comment": paper.get("comment", ""),
                "journal_ref": paper.get("journal_ref", ""),
                "doi": paper.get("doi", ""),
                "pdf_link": paper.get("pdf_link", ""),
                "age_days": age_days,
                "abstract_length": len(extracted_text),
                "score_type": "estimated",  # arXiv has no upvotes
            }
            # Store applied-domain info if detected
            if paper.get("_applied_domain"):
                raw_meta["applied_domain"] = paper["_applied_domain"]

            now_str = now.strftime("%Y-%m-%dT%H:%M:%SZ")
            entries.append({
                "source": "arxiv",
                "source_id": source_id,
                "url": url,
                "title": title,
                "extracted_text": extracted_text,
                "summary": None,
                "category_tags": json.dumps(tags),
                "signal_score": score,
                "raw_metadata": json.dumps(raw_meta),
                "first_seen": now_str,
                "last_updated": now_str,
            })

        return entries


if __name__ == "__main__":
    import argparse
    import sqlite3
    import urllib.parse

    parser = argparse.ArgumentParser(description="arXiv adapter for AI Research Oracle")
    parser.add_argument("--query", default="", help="Search query (empty = recent categories)")
    parser.add_argument("--limit", type=int, default=20, help="Max entries")
    parser.add_argument("--db", default=os.path.join(os.path.dirname(__file__), "..", "oracle.db"), help="SQLite DB")
    parser.add_argument("--schema", default=os.path.join(os.path.dirname(__file__), "..", "schema.sql"), help="Schema file")
    parser.add_argument("--dry-run", action="store_true", help="Don't store in DB")
    args = parser.parse_args()

    print(f"=== arXiv Adapter ===")
    print(f"  Query: {args.query or '(recent AI categories)'}")
    print(f"  Limit: {args.limit}")
    print()

    adapter = ArxivAdapter()
    entries = adapter.fetch(query=args.query, limit=args.limit)

    print(f"  Fetched {len(entries)} entries")

    if not args.dry_run:
        conn = sqlite3.connect(args.db)
        if os.path.exists(args.schema):
            with open(args.schema) as f:
                conn.executescript(f.read())
            conn.commit()

        cur = conn.cursor()
        stored = 0
        for entry in entries:
            try:
                cur.execute("""
                    INSERT OR REPLACE INTO entries
                    (source, source_id, url, title, extracted_text, summary,
                     category_tags, signal_score, raw_metadata, first_seen, last_updated)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    entry["source"], entry["source_id"], entry["url"], entry["title"],
                    entry["extracted_text"], entry["summary"],
                    entry["category_tags"], entry["signal_score"],
                    entry["raw_metadata"], entry["first_seen"], entry["last_updated"],
                ))
                stored += 1
            except Exception as e:
                print(f"  DB error: {e}")

        conn.commit()
        conn.close()
        print(f"  Stored {stored} entries")

    # Print top 5
    print(f"\n  Top entries:")
    for i, e in enumerate(entries[:5]):
        meta = json.loads(e["raw_metadata"]) if isinstance(e["raw_metadata"], str) else e["raw_metadata"]
        authors = meta.get("authors", [])
        author_str = f"{authors[0]} et al." if len(authors) > 2 else ", ".join(authors[:2])
        print(f"    [{i+1}] score={e['signal_score']:.2f} authors={author_str}")
        print(f"        {e['title'][:90]}")
        print(f"        {e['url']}")
        print(f"        abstract={meta.get('abstract_length', 0)}ch")

    print(f"\n  Done.")