tools.py

"""
Tools available to both agents:
  - fetch_arxiv_paper   : downloads & indexes an arXiv paper
  - fetch_pdf_paper     : loads a local / remote PDF
  - fetch_github_repo   : clones repo files into the vector store
  - fetch_github_file   : fetches a single file from GitHub
  - rag_query           : semantic search over the shared vector store
"""

from __future__ import annotations

import io
import os
import re
import tempfile
from textwrap import dedent
from typing import Optional

import requests
from bs4 import BeautifulSoup
from langchain.tools import tool
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.documents import Document

from vector_store import get_vector_store, add_documents

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

GITHUB_TOKEN = os.getenv("GITHUB_TOKEN", "")
_GH_HEADERS = {"Authorization": f"token {GITHUB_TOKEN}"} if GITHUB_TOKEN else {}

# File extensions worth indexing from a GitHub repo
CODE_EXTENSIONS = {
    ".py", ".js", ".ts", ".jsx", ".tsx", ".java", ".cpp", ".c",
    ".h", ".go", ".rs", ".rb", ".php", ".cs", ".swift", ".kt",
    ".md", ".txt", ".ipynb",
}


def _chunk_text(text: str, chunk_size: int = 1500, overlap: int = 200) -> list[str]:
    """Simple sliding-window chunker."""
    chunks, start = [], 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks


def _make_docs(chunks: list[str], metadata: dict) -> list[Document]:
    return [Document(page_content=c, metadata={**metadata, "chunk": i})
            for i, c in enumerate(chunks)]


# ---------------------------------------------------------------------------
# Paper tools
# ---------------------------------------------------------------------------

@tool
def fetch_arxiv_paper(arxiv_id_or_url: str) -> str:
    """
    Download a paper from arXiv (by ID like '2310.06825' or full URL),
    index it in the vector store, and return the abstract + title.
    """
    # Extract bare ID
    arxiv_id = re.sub(r".*arxiv\.org/(abs|pdf)/v?\d*/", "", arxiv_id_or_url)
    arxiv_id = re.sub(r".*arxiv\.org/(abs|pdf)/", "", arxiv_id)
    arxiv_id = arxiv_id.replace(".pdf", "").strip().rstrip("/")

    # Scrape title/abstract from the abs page (no API, no rate limit)
    abs_url = f"https://arxiv.org/abs/{arxiv_id}"
    title, abstract = arxiv_id, ""
    try:
        abs_resp = requests.get(abs_url, timeout=20, headers={"User-Agent": "Mozilla/5.0"})
        if abs_resp.status_code == 200:
            soup = BeautifulSoup(abs_resp.text, "html.parser")
            title_tag = soup.find("h1", class_="title")
            if title_tag:
                title = title_tag.get_text(strip=True).replace("Title:", "").strip()
            abstract_tag = soup.find("blockquote", class_="abstract")
            if abstract_tag:
                abstract = abstract_tag.get_text(strip=True).replace("Abstract:", "").strip()
    except Exception:
        pass

    # Download PDF directly (bypasses the arxiv Python library & its API rate limits)
    pdf_url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
    pdf_resp = requests.get(pdf_url, timeout=60, headers={"User-Agent": "Mozilla/5.0"})
    if pdf_resp.status_code != 200:
        return f"Could not download PDF for arXiv ID '{arxiv_id}' (HTTP {pdf_resp.status_code})."

    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
        tmp.write(pdf_resp.content)
        tmp_path = tmp.name

    loader = PyPDFLoader(tmp_path)
    pages = loader.load()
    os.unlink(tmp_path)

    full_text = "\n".join(p.page_content for p in pages)
    chunks = _chunk_text(full_text)
    metadata = {"source": "arxiv", "arxiv_id": arxiv_id, "title": title}
    add_documents(_make_docs(chunks, metadata))

    return dedent(f"""
        Indexed arXiv paper successfully.
        Title   : {title}
        arXiv ID: {arxiv_id}
        Abstract: {abstract[:800]}...
    """).strip()


@tool
def fetch_pdf_paper(pdf_path_or_url: str) -> str:
    """
    Load a PDF from a local file path or a direct URL,
    index it in the vector store, and return a short summary of page count.
    """
    if pdf_path_or_url.startswith("http"):
        resp = requests.get(pdf_path_or_url, timeout=30)
        resp.raise_for_status()
        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
            tmp.write(resp.content)
            pdf_path = tmp.name
        cleanup = True
    else:
        pdf_path = pdf_path_or_url
        cleanup = False

    loader = PyPDFLoader(pdf_path)
    pages = loader.load()
    if cleanup:
        os.unlink(pdf_path)

    full_text = "\n".join(p.page_content for p in pages)
    chunks = _chunk_text(full_text)
    metadata = {"source": "pdf", "path": pdf_path_or_url}
    add_documents(_make_docs(chunks, metadata))

    return f"Indexed PDF ({len(pages)} pages, {len(chunks)} chunks) from: {pdf_path_or_url}"


# ---------------------------------------------------------------------------
# GitHub tools
# ---------------------------------------------------------------------------

def _parse_github_url(url: str) -> tuple[str, str, Optional[str]]:
    """Return (owner, repo, subpath_or_None) from a GitHub URL."""
    m = re.match(
        r"https?://github\.com/([^/]+)/([^/]+)(?:/(?:tree|blob)/[^/]+/?(.*)?)?",
        url,
    )
    if not m:
        raise ValueError(f"Cannot parse GitHub URL: {url}")
    return m.group(1), m.group(2), m.group(3) or ""


@tool
def fetch_github_repo(github_url: str, max_files: int = 30) -> str:
    """
    Fetch source files from a GitHub repository (up to max_files),
    index them in the vector store. Accepts a repo URL like
    'https://github.com/owner/repo' or a sub-folder URL.
    """
    owner, repo, subpath = _parse_github_url(github_url)
    api_base = f"https://api.github.com/repos/{owner}/{repo}/contents/{subpath}"

    def _recurse(api_url: str, collected: list[dict], depth: int = 0) -> None:
        if depth > 4 or len(collected) >= max_files:
            return
        resp = requests.get(api_url, headers=_GH_HEADERS, timeout=15)
        if resp.status_code != 200:
            return
        items = resp.json()
        if isinstance(items, dict):       # single file
            items = [items]
        for item in items:
            if len(collected) >= max_files:
                break
            if item["type"] == "file":
                ext = os.path.splitext(item["name"])[1].lower()
                if ext in CODE_EXTENSIONS:
                    collected.append(item)
            elif item["type"] == "dir":
                _recurse(item["url"], collected, depth + 1)

    file_items: list[dict] = []
    _recurse(api_base, file_items)

    indexed = 0
    for item in file_items:
        raw = requests.get(item["download_url"], headers=_GH_HEADERS, timeout=15)
        if raw.status_code != 200:
            continue
        content = raw.text
        chunks = _chunk_text(content)
        metadata = {
            "source": "github",
            "repo": f"{owner}/{repo}",
            "file_path": item["path"],
            "html_url": item["html_url"],
        }
        add_documents(_make_docs(chunks, metadata))
        indexed += 1

    return (f"Indexed {indexed} files from {owner}/{repo} "
            f"(subpath='{subpath or '/'}') into the vector store.")


@tool
def fetch_github_file(github_file_url: str) -> str:
    """
    Fetch and index a single file from GitHub.
    Accepts a blob URL like
    'https://github.com/owner/repo/blob/main/path/to/file.py'
    """
    # Convert blob URL to raw URL
    raw_url = github_file_url.replace("github.com", "raw.githubusercontent.com")
    raw_url = re.sub(r"/blob/", "/", raw_url)

    resp = requests.get(raw_url, headers=_GH_HEADERS, timeout=15)
    resp.raise_for_status()
    content = resp.text

    owner, repo, file_path = _parse_github_url(github_file_url)
    file_path = re.sub(r"^blob/[^/]+/", "", file_path)

    chunks = _chunk_text(content)
    metadata = {
        "source": "github",
        "repo": f"{owner}/{repo}",
        "file_path": file_path,
        "html_url": github_file_url,
    }
    add_documents(_make_docs(chunks, metadata))
    return f"Indexed file '{file_path}' from {owner}/{repo} ({len(chunks)} chunks)."


# ---------------------------------------------------------------------------
# GitHub link finder (used by Code Explainer)
# ---------------------------------------------------------------------------

@tool
def find_github_links(dummy: str = "") -> str:
    """
    Search the indexed paper for any GitHub repository URLs mentioned by the authors.
    Returns a list of GitHub links found. Call this before fetch_github_repo.
    The dummy parameter is unused — just pass an empty string.
    """
    vs = get_vector_store()
    results = vs.similarity_search("github.com repository code implementation", k=15)

    github_pattern = re.compile(
        r"https?://github\.com/[\w\-\.]+/[\w\-\.]+", re.IGNORECASE
    )
    links: set[str] = set()
    for doc in results:
        found = github_pattern.findall(doc.page_content)
        links.update(found)

    # Filter out obviously wrong hits (e.g. github.com/user only)
    links = {l for l in links if l.count("/") >= 4}

    if not links:
        return "No GitHub repository links found in the indexed paper."

    return "GitHub links found in paper:\n" + "\n".join(f"- {l}" for l in sorted(links))


# ---------------------------------------------------------------------------
# RAG query tool (shared by both agents)
# ---------------------------------------------------------------------------

@tool
def rag_query(query: str, k: int = 6) -> str:
    """
    Semantic search over all indexed documents (papers + code).
    Returns the top-k most relevant chunks with their source metadata.
    Use this to retrieve specific sections before summarising or generating code.
    """
    vs = get_vector_store()
    results = vs.similarity_search(query, k=k)
    if not results:
        return "No relevant content found. Make sure to index a paper or repo first."

    output_parts = []
    for i, doc in enumerate(results, 1):
        meta = doc.metadata
        source_label = (
            f"[{meta.get('source','?')}] "
            f"{meta.get('title') or meta.get('file_path') or meta.get('path', '')}"
            f" (chunk {meta.get('chunk', '?')})"
        )
        output_parts.append(f"--- Result {i}: {source_label} ---\n{doc.page_content}")

    return "\n\n".join(output_parts)