Arxiv Reader

"""
Agent Factory — creates LangChain agents / chains.
Uses `create_tool_calling_agent` when tools are supplied,
otherwise returns a simple  prompt | llm  chain.
"""

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.agents import create_agent as create_langchain_agent
import config
from typing import Any

# ── LLM singleton cache ──────────────────────────────────────
_llm_cache: dict[str, ChatOpenAI] = {}


def get_llm(
    temperature: float | None = None,
    max_tokens: int | None = None,
    model: str | None = None,
) -> ChatOpenAI:
    """Return a (cached) ChatOpenAI instance."""
    t = temperature if temperature is not None else config.LLM_TEMPERATURE
    m = max_tokens or config.LLM_MAX_TOKENS
    mdl = model or config.LLM_MODEL
    key = f"{mdl}_{t}_{m}"
    if key not in _llm_cache:
        _llm_cache[key] = ChatOpenAI(
            base_url=config.LLM_BASE_URL,
            api_key=config.LLM_API_KEY,
            model=mdl,
            temperature=t,
            max_tokens=m,
        )
    return _llm_cache[key]


def create_agent(
    system_prompt: str,
    tools: list | None = None,
    temperature: float | None = None,
    max_tokens: int | None = None,
    checkpointer: Any | None = None,
):
    """
    Create a LangChain agent.

    - With tools  → AgentExecutor (tool-calling agent)
    - Without tools → simple  prompt | llm  chain  (invoke with {"messages": ...})
    """
    llm = get_llm(temperature=temperature, max_tokens=max_tokens)

    agent_graph = create_langchain_agent(
        model=llm,
        tools=tools,
        system_prompt=system_prompt,
        checkpointer=checkpointer
    )
    return agent_graph

"""
Classifier Agent — decides which skill category a paper belongs to.
"""

from __future__ import annotations

import json
import re
from typing import Any, Dict

import config

from agents.base_agent import create_agent
from skills.loader import get_categories_description
from utils.logger import get_logger

logger = get_logger(__name__)

PROMPTS_DIR = config.PROJECT_ROOT / "prompts"


def _load_prompt(filename: str) -> str:
    path = PROMPTS_DIR / filename
    return path.read_text(encoding="utf-8")


CLASSIFIER_SYSTEM_PROMPT = _load_prompt("classifier_system.md")


class ClassifierAgent:
    """Classify papers into skill categories."""

    def __init__(self, skills: Dict[str, Any]):
        self.skills = skills
        categories_desc = get_categories_description(skills)
        self.category_names = [
            name for name in skills.keys() if name != "general"
        ]
        system_prompt = CLASSIFIER_SYSTEM_PROMPT.format(
            categories=categories_desc
        )
        self.chain = create_agent(system_prompt, temperature=0.1)

    def classify(self, title: str, abstract: str) -> str:
        """Return the category name for a single paper."""
        user_input = (
            f"请对以下论文进行分类：\n\n"
            f"**标题**: {title}\n\n"
            f"**摘要**: {abstract}"
        )
        try:
            result = self.chain.invoke({"messages": [{"role": "user", "content": user_input}]})
            text = result["messages"][-1].content
            parsed = self._parse_json(text)
            category = parsed.get("category", "general")
            confidence = float(parsed.get("confidence", 0))

            if category not in self.category_names or confidence < 0.4:
                category = "general"

            logger.info(
                f"  [{category}] (conf={confidence:.2f}) {title[:60]}..."
            )
            return category

        except Exception as e:
            logger.warning(f"Classification failed for '{title[:50]}': {e}")
            return "general"

    # ------------------------------------------------------------------
    def classify_batch(self, papers: list[dict]) -> dict[str, str]:
        """Classify a list of papers. Returns {arxiv_id: category}."""
        results: dict[str, str] = {}
        for paper in papers:
            cat = self.classify(paper["title"], paper["abstract"])
            results[paper["arxiv_id"]] = cat
        return results

    @staticmethod
    def _parse_json(text: str) -> dict:
        """Robustly extract JSON from LLM output."""
        # Try direct parse
        text = text.strip()
        # Remove possible ```json ... ``` wrapping
        text = re.sub(r"^```(?:json)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            # Try to find JSON object in text
            match = re.search(r"\{.*\}", text, re.DOTALL)
            if match:
                return json.loads(match.group())
            return {"category": "general", "confidence": 0, "reasoning": "parse error"}

"""
Reader Agent — performs two-pass deep reading of a full paper.

Pass 1: Abstract + Introduction + Preliminaries + Contributions + Limitations
         → initial summary
Pass 2: Initial summary + main body (excluding refs/appendix/Pass-1 sections)
         → detailed notes, plus a decision on whether to read the appendix
Pass 3 (optional): Previous notes + appendix → updated notes
"""

from __future__ import annotations

import re
from pathlib import Path
from typing import Any, Dict

from agents.base_agent import create_agent
import config
from paper_reader.latex_parser import ParsedPaper
from utils.helpers import truncate_text
from utils.logger import get_logger

logger = get_logger(__name__)

PROMPTS_DIR = config.PROJECT_ROOT / "prompts"


def _load_prompt(filename: str) -> str:
    path = PROMPTS_DIR / filename
    return path.read_text(encoding="utf-8")


# ── Prompt templates ──────────────────────────────────────────

READER_SYSTEM_PROMPT = _load_prompt("reader_system.md")
FIRST_PASS_USER = _load_prompt("reader_first_pass_user.md")
SECOND_PASS_USER = _load_prompt("reader_second_pass_user.md")
APPENDIX_PASS_USER = _load_prompt("reader_appendix_pass_user.md")


class ReaderAgent:
    """Two-pass deep reading agent for a specific skill category."""

    def __init__(self, skill_name: str, skill_config: Dict[str, Any]):
        self.skill_name = skill_name
        self.reading_prompt = skill_config.get("reading_prompt", "")
        system_prompt = READER_SYSTEM_PROMPT.format(
            reading_prompt=self.reading_prompt
        )
        self.chain = create_agent(system_prompt, max_tokens=16000)

    # ------------------------------------------------------------------

    def read_paper(
        self, paper_info: dict, parsed_paper: ParsedPaper
    ) -> str:
        """
        Execute the full reading pipeline and return Markdown notes.
        """
        title = paper_info["title"]
        authors = ", ".join(paper_info.get("authors", []))
        arxiv_id = paper_info["arxiv_id"]

        # ── Pass 1 ─────────────────────────────────────────
        logger.info(f"  [Pass 1] {title[:60]}...")
        first_pass_text = truncate_text(parsed_paper.first_pass_text, 30000)
        user_msg_1 = FIRST_PASS_USER.format(
            title=title,
            authors=authors,
            arxiv_id=arxiv_id,
            first_pass_content=first_pass_text,
        )
        result_1 = self.chain.invoke({"messages": [{"role": "user", "content": user_msg_1}]})
        initial_summary = result_1["messages"][-1].content

        # ── Pass 2 ─────────────────────────────────────────
        logger.info(f"  [Pass 2] {title[:60]}...")
        main_body = truncate_text(parsed_paper.main_body_text, 50000)
        if not main_body.strip():
            # Paper has no extractable main body → return Pass 1 result
            return self._format_final_notes(paper_info, initial_summary)

        user_msg_2 = SECOND_PASS_USER.format(
            initial_summary=initial_summary,
            main_body=main_body,
        )
        result_2 = self.chain.invoke({"messages": [{"role": "user", "content": user_msg_2}]})
        detailed_notes = result_2["messages"][-1].content

        # ── Pass 3 (optional: appendix) ────────────────────
        if self._should_read_appendix(detailed_notes) and parsed_paper.has_appendix:
            logger.info(f"  [Pass 3 - Appendix] {title[:60]}...")
            appendix = truncate_text(parsed_paper.appendix_text, 30000)
            user_msg_3 = APPENDIX_PASS_USER.format(
                current_notes=detailed_notes,
                appendix_content=appendix,
            )
            result_3 = self.chain.invoke({"messages": [{"role": "user", "content": user_msg_3}]})
            appendix_notes = result_3["messages"][-1].content.strip()
            # 将增量内容追加到第二次笔记后（跳过无补充的情况）
            if appendix_notes and "附录无重要补充内容" not in appendix_notes:
                detailed_notes = detailed_notes + "\n\n" + appendix_notes
            return self._format_final_notes(paper_info, detailed_notes)

        return self._format_final_notes(paper_info, detailed_notes)

    # ------------------------------------------------------------------

    @staticmethod
    def _should_read_appendix(notes: str) -> bool:
        """Check if the model decided to read the appendix."""
        match = re.search(r"APPENDIX_NEEDED:\s*(YES|NO)", notes, re.IGNORECASE)
        if match:
            return match.group(1).upper() == "YES"
        return False

    @staticmethod
    def _format_final_notes(paper_info: dict, notes_body: str) -> str:
        """Wrap notes with YAML front-matter for Obsidian."""
        notes_body = notes_body.strip()

        authors = ", ".join(paper_info.get("authors", [])[:5])
        categories = ", ".join(paper_info.get("categories", []))
        arxiv_id = paper_info["arxiv_id"]

        header = f"""\
---
title: "{paper_info['title']}"
arxiv_id: "{arxiv_id}"
authors: "{authors}"
categories: "{categories}"
date_read: "{paper_info.get('date_read', '')}"
skill_category: "{paper_info.get('skill_category', '')}"
pdf_url: "{paper_info.get('pdf_url', '')}"
---

# {paper_info['title']}

> **ArXiv**: [{arxiv_id}](https://arxiv.org/abs/{arxiv_id}) | **Authors**: {authors}

"""
        return header + notes_body

"""
Summary Agent — quick abstract-only summary for uncategorized papers.
"""

from __future__ import annotations

from typing import Any, Dict

import config

from agents.base_agent import create_agent
from utils.logger import get_logger

logger = get_logger(__name__)

PROMPTS_DIR = config.PROJECT_ROOT / "prompts"


def _load_prompt(filename: str) -> str:
    path = PROMPTS_DIR / filename
    return path.read_text(encoding="utf-8")


SUMMARY_SYSTEM_PROMPT = _load_prompt("summary_system.md")


class SummaryAgent:
    """Generate quick summaries from title + abstract only."""

    def __init__(self, skill_config: Dict[str, Any]):
        reading_prompt = skill_config.get("reading_prompt", "")
        system_prompt = SUMMARY_SYSTEM_PROMPT.format(
            reading_prompt=reading_prompt
        )
        self.chain = create_agent(system_prompt, temperature=0.2, max_tokens=4000)

    def summarize(self, paper_info: dict) -> str:
        """Generate a summary note for one paper."""
        title = paper_info["title"]
        abstract = paper_info["abstract"]
        authors = ", ".join(paper_info.get("authors", [])[:5])
        arxiv_id = paper_info["arxiv_id"]

        user_input = (
            f"请对以下论文进行快速总结：\n\n"
            f"**标题**: {title}\n"
            f"**作者**: {authors}\n"
            f"**ArXiv ID**: {arxiv_id}\n\n"
            f"**摘要**:\n{abstract}"
        )

        try:
            result = self.chain.invoke({"messages": [{"role": "user", "content": user_input}]})
            body = result["messages"][-1].content
        except Exception as e:
            logger.error(f"Summary failed for '{title[:50]}': {e}")
            body = f"*自动总结失败: {e}*\n\n## 摘要\n\n{abstract}"

        return self._format_note(paper_info, body)

    # ------------------------------------------------------------------

    @staticmethod
    def _format_note(paper_info: dict, body: str) -> str:
        authors = ", ".join(paper_info.get("authors", [])[:5])
        categories = ", ".join(paper_info.get("categories", []))
        arxiv_id = paper_info["arxiv_id"]

        header = f"""\
---
title: "{paper_info['title']}"
arxiv_id: "{arxiv_id}"
authors: "{authors}"
categories: "{categories}"
date_read: "{paper_info.get('date_read', '')}"
skill_category: "general"
pdf_url: "{paper_info.get('pdf_url', '')}"
---

# {paper_info['title']}

> **ArXiv**: [{arxiv_id}](https://arxiv.org/abs/{arxiv_id}) | **Authors**: {authors}

"""
        return header + body

"""
ArXiv Fetcher — fetch recent papers and their LaTeX source.

Strategy (in order):
  1. RSS feed   – fast, gives today's new listings directly
  2. arxiv API  – fallback, get the most recent N papers (no date filter)

Uses direct HTTP download for LaTeX source files.
"""

from __future__ import annotations

import gzip
import io
import re
import tarfile
import time
import xml.etree.ElementTree as ET
from datetime import date, datetime, timedelta, timezone
from typing import Callable, Optional

import arxiv
import requests

import config
from utils.helpers import extract_arxiv_id
from utils.logger import get_logger

logger = get_logger(__name__)

# arXiv RSS namespaces
_RSS_NS = {
    "rdf": "http://www.w3.org/1999/02/22-rdf-syntax-ns#",
    "rss": "http://purl.org/rss/1.0/",
    "dc": "http://purl.org/dc/elements/1.1/",
    "taxo": "http://purl.org/rss/1.0/modules/taxonomy/",
}


class ArxivFetcher:
    """Fetch paper listings and LaTeX source from arXiv."""

    def __init__(self):
        self.client = arxiv.Client(
            page_size=100,
            delay_seconds=3.0,
            num_retries=3,
        )
        self.categories = config.ARXIV_CATEGORIES
        self.max_results = config.ARXIV_MAX_RESULTS

    # ── Paper listing ─────────────────────────────────────────

    def fetch_single_paper(self, arxiv_id_or_url: str) -> Optional[dict]:
        """
        Fetch a single paper by arxiv_id or URL.
        
        Parameters
        ----------
        arxiv_id_or_url : str
            Can be:
            - arxiv_id: "2401.12345"
            - abs URL: "https://arxiv.org/abs/2401.12345"
            - pdf URL: "https://arxiv.org/pdf/2401.12345.pdf"
        
        Returns
        -------
        dict or None
            Paper metadata dict, or None if not found.
        """
        # Extract arxiv_id from URL if needed
        arxiv_id = extract_arxiv_id(arxiv_id_or_url)
        if not arxiv_id:
            logger.error(f"Invalid arxiv_id or URL: {arxiv_id_or_url}")
            return None
        
        logger.info(f"Fetching single paper: {arxiv_id}")
        
        try:
            search = arxiv.Search(id_list=[arxiv_id])
            results = list(self.client.results(search))
            
            if not results:
                logger.error(f"Paper not found: {arxiv_id}")
                return None
            
            paper = self._result_to_dict(results[0])
            logger.info(f"  Found: {paper['title'][:60]}...")
            return paper
            
        except Exception as e:
            logger.error(f"Failed to fetch {arxiv_id}: {e}")
            return None

    def fetch_papers(
        self,
        target_date: Optional[date] = None,
        is_known_fn: Optional[Callable[[str], bool]] = None,
    ) -> list[dict]:
        """
        Fetch recent papers from arXiv.

        Strategy:
        1. Try RSS feeds for each category (today's new listings).
        2. If RSS yields nothing, fall back to the arxiv API.

        Stop conditions (both RSS and API):
          - Paper already in DB (is_known_fn returns True) → stop / filter out
          - Paper older than (today - 1 day) → stop

        Parameters
        ----------
        target_date : date, optional
            Unused currently but kept for interface compatibility.
        is_known_fn : callable, optional
            Function that takes an arxiv_id and returns True if the
            paper is already tracked in the local DB.
        """
        # ── Method 1: RSS feeds ──
        logger.info(
            f"Fetching papers via RSS from [{', '.join(self.categories)}] ..."
        )
        papers = self._fetch_via_rss()
        if papers:
            # Dedup by arxiv_id (a paper can appear in multiple category feeds)
            papers = self._dedup_papers(papers)
            logger.info(f"RSS: got {len(papers)} unique papers")

            # Filter out papers already in DB
            if is_known_fn:
                before = len(papers)
                papers = [p for p in papers if not is_known_fn(p["arxiv_id"])]
                skipped = before - len(papers)
                if skipped:
                    logger.info(
                        f"  RSS: filtered out {skipped} papers already in DB, "
                        f"{len(papers)} new papers remaining"
                    )

            if self.max_results and len(papers) > self.max_results:
                papers = papers[: self.max_results]
            return papers

        # ── Method 2: arxiv API (fallback) ──
        logger.info("RSS returned nothing, falling back to arxiv API ...")
        return self._fetch_via_api(is_known_fn=is_known_fn)

    # ── RSS fetching ──────────────────────────────────────────

    def _fetch_via_rss(self) -> list[dict]:
        """Fetch today's new papers from arXiv RSS feeds."""
        all_papers: list[dict] = []

        for cat in self.categories:
            url = f"https://rss.arxiv.org/rss/{cat}"
            logger.info(f"  RSS: {url}")
            try:
                resp = requests.get(url, timeout=30)
                resp.raise_for_status()
                papers = self._parse_rss(resp.text, cat)
                all_papers.extend(papers)
                logger.info(f"    → {len(papers)} papers from {cat}")
            except Exception as e:
                logger.warning(f"  RSS failed for {cat}: {e}")

        return all_papers

    def _parse_rss(self, xml_text: str, category: str) -> list[dict]:
        """Parse an arXiv RSS feed XML into paper dicts."""
        papers: list[dict] = []
        try:
            root = ET.fromstring(xml_text)
        except ET.ParseError as e:
            logger.warning(f"  RSS XML parse error: {e}")
            return papers

        # arXiv RSS uses RDF/RSS 1.0 format
        items = root.findall(".//rss:item", _RSS_NS)
        if not items:
            # Try plain RSS 2.0 fallback
            items = root.findall(".//item")

        for item in items:
            paper = self._rss_item_to_dict(item, category)
            if paper:
                papers.append(paper)

        return papers

    def _rss_item_to_dict(self, item: ET.Element, default_cat: str) -> Optional[dict]:
        """Convert a single RSS <item> to our paper dict."""
        # Try RDF namespace first, then plain
        title_el = item.find("rss:title", _RSS_NS) or item.find("title")
        link_el = item.find("rss:link", _RSS_NS) or item.find("link")
        desc_el = item.find("rss:description", _RSS_NS) or item.find("description")
        creator_el = item.find("dc:creator", _RSS_NS)

        if title_el is None or link_el is None:
            return None

        raw_title = (title_el.text or "").strip()
        link = (link_el.text or "").strip()
        raw_desc = (desc_el.text or "").strip() if desc_el is not None else ""

        # Skip "UPDATED" entries, focus on new submissions
        # arXiv RSS titles look like: "Title (arXiv:2401.12345v1 [cs.AI])"
        # or sometimes: "Title. (arXiv:2401.12345v1 [cs.AI] UPDATED)"
        is_updated = "UPDATED" in raw_title

        # Extract arxiv ID from link: https://arxiv.org/abs/2401.12345
        arxiv_id = extract_arxiv_id(link) if link else ""
        if not arxiv_id:
            return None

        # Clean title: remove the trailing "(arXiv:...)" part
        title = re.sub(r"\s*\(arXiv:[^)]+\)\s*$", "", raw_title).strip()
        title = re.sub(r"\.\s*$", "", title)  # remove trailing period

        # Parse abstract from description (may contain HTML)
        abstract = self._clean_html(raw_desc)
        # arXiv RSS description often starts with "<p>Abstract: ..."
        abstract = re.sub(r"^Abstract:\s*", "", abstract, flags=re.IGNORECASE).strip()

        # Parse authors
        authors: list[str] = []
        if creator_el is not None and creator_el.text:
            # Format: "<a href='...'>Author1</a>, <a href='...'>Author2</a>"
            author_text = self._clean_html(creator_el.text)
            authors = [a.strip() for a in author_text.split(",") if a.strip()]

        # Extract categories from title bracket part
        cat_match = re.search(r"\[([^\]]+)\]", raw_title)
        categories = (
            [c.strip() for c in cat_match.group(1).split(",")]
            if cat_match
            else [default_cat]
        )

        return {
            "arxiv_id": arxiv_id,
            "title": title,
            "abstract": abstract,
            "authors": authors,
            "categories": categories,
            "published": date.today(),
            "pdf_url": f"https://arxiv.org/pdf/{arxiv_id}",
            "is_updated": is_updated,
        }

    # ── API fallback ──────────────────────────────────────────

    def _fetch_via_api(
        self, is_known_fn: Optional[Callable[[str], bool]] = None,
    ) -> list[dict]:
        """
        Fallback: fetch most recent papers via the arxiv API.

        Papers are sorted newest-first. Early-stop when:
          1. Paper already in local DB → all older papers should also be known
          2. Paper published before yesterday → we only care about today + yesterday
        """
        cat_query = " OR ".join(f"cat:{cat}" for cat in self.categories)
        search = arxiv.Search(
            query=cat_query,
            max_results=self.max_results,
            sort_by=arxiv.SortCriterion.SubmittedDate,
            sort_order=arxiv.SortOrder.Descending,
        )

        # Only look at today and yesterday (1 day lookback)
        cutoff_date = date.today() - timedelta(days=config.FETCH_LOOKBACK_DAYS)

        papers: list[dict] = []
        total_scanned = 0
        logger.info(
            f"  API: fetching up to {self.max_results} papers "
            f"(cutoff date: {cutoff_date}, stop on first known paper) ..."
        )

        for result in self.client.results(search):
            total_scanned += 1
            paper = self._result_to_dict(result)

            # ── Early stop 1: paper older than yesterday ──
            if paper["published"] < cutoff_date:
                logger.info(
                    f"  ■ Stop: reached paper from {paper['published']} "
                    f"(cutoff {cutoff_date}). Scanned {total_scanned}."
                )
                break

            # ── Early stop 2: paper already in DB → we've fetched up to here ──
            if is_known_fn and is_known_fn(paper["arxiv_id"]):
                logger.info(
                    f"  ■ Stop: paper {paper['arxiv_id']} already in DB. "
                    f"All older papers should be known too. Scanned {total_scanned}."
                )
                break

            papers.append(paper)
            if len(papers) >= self.max_results:
                break

        logger.info(f"  API: got {len(papers)} new papers (scanned {total_scanned})")
        return papers

    # ── LaTeX source fetching ─────────────────────────────────

    def fetch_latex_source(self, arxiv_id: str) -> Optional[str]:
        """
        Fetch the LaTeX source of a paper.

        1. Try `arxiv_to_prompt` library (if installed).
        2. Fall back to downloading the e-print tar/gz from arXiv.
        """
        # ── Method 1: arxiv_to_prompt ──
        try:
            from arxiv_to_prompt import process_latex_source

            logger.info(f"  Fetching source via arxiv_to_prompt: {arxiv_id}")
            return process_latex_source(arxiv_id, keep_comments=False)
        except ImportError:
            pass
        except Exception as e:
            logger.warning(f"arxiv_to_prompt failed for {arxiv_id}: {e}")
        
        return None

    @staticmethod
    def _find_main_tex(tar: tarfile.TarFile) -> Optional[str]:
        """Find the main .tex file in a tar archive."""
        tex_files: list[tuple[str, str]] = []

        for member in tar.getmembers():
            if member.name.endswith(".tex") and not member.name.startswith("."):
                f = tar.extractfile(member)
                if f:
                    content = f.read().decode("utf-8", errors="ignore")
                    tex_files.append((member.name, content))

        if not tex_files:
            return None

        # Prefer the file containing \begin{document}
        for name, content in tex_files:
            if "\\begin{document}" in content:
                return content

        # Fall back to largest tex file
        tex_files.sort(key=lambda x: len(x[1]), reverse=True)
        return tex_files[0][1]

    @staticmethod
    def _result_to_dict(result: arxiv.Result) -> dict:
        """Convert an arxiv.Result to a plain dict."""
        return {
            "arxiv_id": extract_arxiv_id(result.entry_id),
            "title": result.title.replace("\n", " ").strip(),
            "abstract": result.summary.replace("\n", " ").strip(),
            "authors": [a.name for a in result.authors],
            "categories": list(result.categories),
            "published": result.published.date(),
            "pdf_url": result.pdf_url or "",
        }

    @staticmethod
    def _dedup_papers(papers: list[dict]) -> list[dict]:
        """Remove duplicate papers by arxiv_id, keeping the first."""
        seen: set[str] = set()
        unique: list[dict] = []
        for p in papers:
            if p["arxiv_id"] not in seen:
                seen.add(p["arxiv_id"])
                unique.append(p)
        return unique

    @staticmethod
    def _clean_html(text: str) -> str:
        """Strip HTML tags from text."""
        text = re.sub(r"<[^>]+>", "", text)
        text = text.replace("&amp;", "&").replace("&lt;", "<").replace("&gt;", ">")
        text = text.replace("&quot;", '"').replace("&#39;", "'")
        return text.strip()

"""
Global configuration module.
Loads all settings from .env file at project root.
"""

import os
from pathlib import Path
from dotenv import load_dotenv

# ── Project Root ──────────────────────────────────────────────
PROJECT_ROOT = Path(__file__).parent.resolve()
load_dotenv(PROJECT_ROOT / ".env")

# ── LLM ──────────────────────────────────────────────────────
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://api.openai.com/v1")
LLM_API_KEY = os.getenv("LLM_API_KEY", "")
LLM_MODEL = os.getenv("LLM_MODEL", "gpt-4o")
LLM_TEMPERATURE = float(os.getenv("LLM_TEMPERATURE", "0.3"))
LLM_MAX_TOKENS = int(os.getenv("LLM_MAX_TOKENS", "16000"))

# ── Skills ───────────────────────────────────────────────────
SKILLS_DIR = PROJECT_ROOT / "skills"

#!/usr/bin/env python3
"""Single arXiv paper reader entrypoint.

Usage:
    python main.py --list
    python main.py --arxiv-id 2401.12345
    python main.py --arxiv-id https://arxiv.org/abs/2401.12345
    python main.py --arxiv-id 2401.12345 --category rag_and_retrieval
"""

from __future__ import annotations

import argparse
import sys

# Ensure project root is on sys.path
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from arxiv_fetcher.fetcher import ArxivFetcher
from skills.loader import load_all_skills
from agents.classifier_agent import ClassifierAgent
from agents.reader_agent import ReaderAgent
from agents.summary_agent import SummaryAgent
from paper_reader.latex_parser import parse_latex
from utils.logger import get_logger

logger = get_logger("daily-arxiv")


def process_single_paper(arxiv_id_or_url: str, category: str | None = None) -> str:
    """Read one paper and return raw markdown text."""
    logger.info("=" * 60)
    logger.info("Single Paper Mode")
    logger.info("=" * 60)

    skills = load_all_skills()
    if not skills:
        raise RuntimeError("No skills loaded. Please check the skills/ directory.")

    fetcher = ArxivFetcher()

    logger.info(f"Fetching paper: {arxiv_id_or_url}")
    paper = fetcher.fetch_single_paper(arxiv_id_or_url)
    if not paper:
        raise RuntimeError("Failed to fetch paper.")

    selected_category = category
    if selected_category is None:
        classifier = ClassifierAgent(skills)
        selected_category = classifier.classify(paper["title"], paper.get("abstract", ""))

    if selected_category not in skills:
        valid = ", ".join(sorted(skills.keys()))
        raise ValueError(f"Unknown category '{selected_category}'. Available: {valid}")

    # Build agents
    reader_agents: dict[str, ReaderAgent] = {}
    for skill_name, skill_config in skills.items():
        if skill_name != "general":
            reader_agents[skill_name] = ReaderAgent(skill_name, skill_config)
    summary_agent = SummaryAgent(skills.get("general", {"reading_prompt": ""}))

    logger.info(f"Processing as category: {selected_category}")

    if selected_category != "general" and selected_category in reader_agents:
        notes = _deep_read(
            fetcher,
            reader_agents[selected_category],
            paper,
            summary_agent,
        )
        return notes

    return summary_agent.summarize(paper)


def _deep_read(
    fetcher: ArxivFetcher,
    reader: ReaderAgent,
    paper: dict,
    fallback_agent: SummaryAgent,
 ) -> str:
    """
    Attempt deep two-pass reading.
    """
    arxiv_id = paper["arxiv_id"]

    # Fetch LaTeX source
    latex_source = fetcher.fetch_latex_source(arxiv_id)
    if not latex_source:
        logger.warning(f"  No LaTeX source for {arxiv_id}, falling back to summary.")
        return fallback_agent.summarize(paper)

    # Parse LaTeX
    parsed = parse_latex(latex_source)

    if not parsed.abstract and not parsed.sections:
        logger.warning(f"  LaTeX parsing yielded nothing for {arxiv_id}, falling back.")
        return fallback_agent.summarize(paper)

    # Two-pass reading
    return reader.read_paper(paper, parsed)


# ── CLI ───────────────────────────────────────────────────────


def main():
    parser = argparse.ArgumentParser(
        description="Single arXiv paper reader"
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List available paper categories (skills)",
    )
    parser.add_argument(
        "--arxiv-id",
        type=str,
        default=None,
        help="Read a single paper by arXiv ID or URL",
    )
    parser.add_argument(
        "--category",
        type=str,
        default=None,
        help="Optional: force a specific category (skill folder name)",
    )
    args = parser.parse_args()

    skills = load_all_skills()

    if args.list:
        for name in sorted(skills.keys()):
            print(name)
        return

    if not args.arxiv_id:
        parser.error("--arxiv-id is required unless --list is used")

    notes = process_single_paper(args.arxiv_id, category=args.category)
    print(notes)


if __name__ == "__main__":
    main()

"""
LaTeX Parser — parse a raw .tex file into structured sections
so the Reader Agent can perform multi-pass reading.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field


@dataclass
class Section:
    """A single section extracted from a LaTeX document."""

    title: str
    content: str
    level: int = 1  # 1 = \section, 2 = \subsection, 3 = \subsubsection
    is_appendix: bool = False


@dataclass
class ParsedPaper:
    """Structured representation of a parsed LaTeX paper."""

    raw: str = ""
    abstract: str = ""
    sections: list[Section] = field(default_factory=list)
    appendix_sections: list[Section] = field(default_factory=list)

    # ── Convenience properties for the reading pipeline ──────

    # Names (case-insensitive) that belong to the first-pass reading
    _FIRST_PASS_KEYWORDS: tuple[str, ...] = (
        "introduction",
        "preliminar",
        "background",
        "contribution",
        "limitation",
        "related work",
        "overview",
        "motivation",
    )

    @property
    def first_pass_text(self) -> str:
        """Abstract + Introduction + Preliminaries + Contributions + Limitations."""
        parts: list[str] = []
        if self.abstract:
            parts.append(f"## Abstract\n\n{self.abstract}")

        for sec in self.sections:
            title_lower = sec.title.lower()
            if any(kw in title_lower for kw in self._FIRST_PASS_KEYWORDS):
                parts.append(f"## {sec.title}\n\n{sec.content}")

        return "\n\n---\n\n".join(parts) if parts else self.abstract

    @property
    def main_body_text(self) -> str:
        """Everything except first-pass sections, references, and appendix."""
        parts: list[str] = []
        for sec in self.sections:
            title_lower = sec.title.lower()
            # Skip first-pass sections
            if any(kw in title_lower for kw in self._FIRST_PASS_KEYWORDS):
                continue
            # Skip references
            if "reference" in title_lower or "bibliograph" in title_lower:
                continue
            parts.append(f"## {sec.title}\n\n{sec.content}")

        return "\n\n---\n\n".join(parts)

    @property
    def appendix_text(self) -> str:
        """Appendix sections concatenated."""
        if not self.appendix_sections:
            return ""
        parts = [f"## {sec.title}\n\n{sec.content}" for sec in self.appendix_sections]
        return "\n\n---\n\n".join(parts)

    @property
    def has_appendix(self) -> bool:
        return bool(self.appendix_sections)


# ── Public API ────────────────────────────────────────────────


def parse_latex(tex: str) -> ParsedPaper:
    """
    Parse raw LaTeX source into a ``ParsedPaper``.

    Handles:
    - \\begin{abstract} ... \\end{abstract}
    - \\section{}, \\subsection{}, \\subsubsection{}
    - \\appendix marker
    - \\bibliography / \\begin{thebibliography}
    """
    paper = ParsedPaper(raw=tex)

    # ── 1. Extract abstract ──
    abs_match = re.search(
        r"\\begin\{abstract\}(.*?)\\end\{abstract\}", tex, re.DOTALL
    )
    if abs_match:
        paper.abstract = _clean_latex(abs_match.group(1).strip())

    # ── 2. Locate key markers ──
    appendix_pos = _find_marker(tex, r"\\appendix(?:\b|[^a-zA-Z])")
    bib_pos = _find_marker(
        tex,
        r"\\(?:bibliography\{|begin\{thebibliography\}|printbibliography)",
    )
    doc_end_pos = tex.find("\\end{document}")
    if doc_end_pos < 0:
        doc_end_pos = len(tex)

    # ── 3. Extract sections ──
    section_pattern = re.compile(
        r"\\(section|subsection|subsubsection)\{([^}]+)\}"
    )
    matches = list(section_pattern.finditer(tex))

    for i, m in enumerate(matches):
        level_str, title = m.group(1), m.group(2).strip()
        level = {"section": 1, "subsection": 2, "subsubsection": 3}[level_str]
        start = m.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else doc_end_pos

        # Trim if content extends past bibliography/appendix/end
        for boundary in [bib_pos, doc_end_pos]:
            if boundary and start < boundary < end:
                end = boundary

        content = _clean_latex(tex[start:end].strip())

        sec = Section(title=title, content=content, level=level)

        # Determine if this section is in the appendix region
        if appendix_pos is not None and m.start() >= appendix_pos:
            sec.is_appendix = True
            paper.appendix_sections.append(sec)
        else:
            paper.sections.append(sec)

    return paper


# ── Helpers ───────────────────────────────────────────────────


def _find_marker(tex: str, pattern: str) -> int | None:
    """Return the position of the first regex match, or None."""
    m = re.search(pattern, tex)
    return m.start() if m else None


def _clean_latex(text: str) -> str:
    """Light cleanup of LaTeX content for readability.

    We do NOT strip all LaTeX commands — LLMs read LaTeX fine.
    We only remove very noisy elements.
    """
    # Remove \label{...}
    text = re.sub(r"\\label\{[^}]*\}", "", text)
    # Remove \vspace{...}, \hspace{...}
    text = re.sub(r"\\[vh]space\{[^}]*\}", "", text)
    # Collapse multiple blank lines
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()

langchain>=1.2.9
langchain-openai>=1.1.7
requests>=2.31.0
python-dotenv>=1.0.0
arxiv>=2.1.0
arxiv-to-prompt

"""
Skill Loader — scans the skills/ directory structure and loads
category metadata + reading prompts for each sub-folder.
"""

from pathlib import Path
from typing import Dict, Any

import config
from utils.logger import get_logger

logger = get_logger(__name__)

METADATA_FILE = "_metadata.md"
READING_PROMPT_FILE = "reading_prompt.md"


def _read_md(path: Path) -> str:
    """Read a markdown file, return empty string if missing."""
    if path.exists():
        return path.read_text(encoding="utf-8")
    return ""


def load_all_skills() -> Dict[str, Dict[str, Any]]:
    """
    Load every skill category from the skills/ directory.

    Returns:
        {
            "agent_systems": {
                "name": "agent_systems",
                "metadata": "<content of _metadata.md>",
                "reading_prompt": "<content of reading_prompt.md>",
            },
            ...
        }
    """
    skills_dir = config.SKILLS_DIR
    skills: Dict[str, Dict[str, Any]] = {}

    if not skills_dir.exists():
        logger.warning(f"Skills directory not found: {skills_dir}")
        return skills

    for child in sorted(skills_dir.iterdir()):
        if child.is_dir() and not child.name.startswith("_"):
            metadata = _read_md(child / METADATA_FILE)
            reading_prompt = _read_md(child / READING_PROMPT_FILE)

            if not metadata:
                logger.warning(
                    f"Skill '{child.name}' has no {METADATA_FILE}, skipping."
                )
                continue

            skills[child.name] = {
                "name": child.name,
                "metadata": metadata,
                "reading_prompt": reading_prompt,
            }
            logger.info(f"Loaded skill: {child.name}")

    return skills


def get_categories_description(skills: Dict[str, Dict[str, Any]]) -> str:
    """
    Build a textual description of all categories (used by the classifier).
    """
    lines: list[str] = []
    for name, info in skills.items():
        if name == "general":
            continue  # general is the fallback, not a classification target
        lines.append(f"### {name}\n{info['metadata']}\n")
    return "\n".join(lines)

"""Utility helpers for text processing, truncation, etc."""

import re
from datetime import date, timedelta


def sanitize_filename(title: str, max_len: int = 80) -> str:
    """Convert a paper title into a safe filename slug."""
    safe = re.sub(r"[^\w\s-]", "", title)
    safe = re.sub(r"\s+", "-", safe).strip("-")
    return safe[:max_len]


def truncate_text(text: str, max_chars: int = 60000) -> str:
    """Truncate text to fit within a character budget, keeping the beginning."""
    if len(text) <= max_chars:
        return text
    return text[:max_chars] + "\n\n...[内容因超出长度限制被截断]..."


def get_target_dates() -> list[date]:
    """Return today and yesterday as candidate dates."""
    today = date.today()
    return [today, today - timedelta(days=1)]


def extract_arxiv_id(entry_id: str) -> str:
    """Extract clean arxiv id from a full entry URL.
    e.g. 'http://arxiv.org/abs/2401.12345v1' -> '2401.12345'
    """
    raw = entry_id.rstrip("/").split("/")[-1]
    # Remove version suffix
    return re.sub(r"v\d+$", "", raw)

"""Logging configuration for Daily Arxiv Reader."""

import logging
import sys


def get_logger(name: str, level: int = logging.INFO) -> logging.Logger:
    """Get a configured logger instance."""
    logger = logging.getLogger(name)
    if not logger.handlers:
        logger.setLevel(level)
        handler = logging.StreamHandler(sys.stdout)
        handler.setLevel(level)
        fmt = logging.Formatter(
            "[%(asctime)s] %(levelname)-7s %(name)s — %(message)s",
            datefmt="%H:%M:%S",
        )
        handler.setFormatter(fmt)
        logger.addHandler(handler)
    return logger