"""arXiv API client.

Provides:
- fetch_source: download the .tar.gz source of a paper and extract the primary .tex file.
- fetch_pdf_url: return the canonical PDF URL for a given arXiv ID.
"""

from __future__ import annotations

import io
import logging
import tarfile

import httpx

logger = logging.getLogger(__name__)

_BASE = "https://arxiv.org"


def fetch_source(arxiv_id: str) -> str | None:
    """Download and extract the primary LaTeX source for an arXiv paper.

    Downloads the .tar.gz source bundle from ``https://arxiv.org/src/{arxiv_id}``,
    locates the primary .tex file (preferring any file containing ``\\documentclass``,
    falling back to the largest .tex by size), and returns its contents as a UTF-8
    string.

    Parameters
    ----------
    arxiv_id:
        The arXiv identifier (e.g. ``"2301.07041"``).

    Returns
    -------
    str | None
        Raw LaTeX source string, or None if the paper is not found or no .tex
        file is present (signals Nougat fallback).
    """
    url = f"{_BASE}/src/{arxiv_id}"
    try:
        response = httpx.get(url, timeout=60, follow_redirects=True)
    except httpx.RequestError:
        raise
    if response.status_code == 404:
        logger.debug("arXiv 404 for source id=%s", arxiv_id)
        return None
    if response.status_code != 200:
        response.raise_for_status()

    raw = response.content
    try:
        with tarfile.open(fileobj=io.BytesIO(raw), mode="r:gz") as tf:
            tex_members = [m for m in tf.getmembers() if m.name.endswith(".tex")]
            if not tex_members:
                logger.debug("No .tex files found in arXiv source for %s", arxiv_id)
                return None

            # Prefer the file containing \documentclass (primary document)
            primary: tarfile.TarInfo | None = None
            for member in tex_members:
                f = tf.extractfile(member)
                if f is None:
                    continue
                content_bytes = f.read()
                if b"\\documentclass" in content_bytes:
                    primary = member
                    # Decode and return immediately — first match wins
                    return content_bytes.decode("utf-8", errors="replace")

            if primary is None:
                # Fallback: largest .tex by size
                largest = max(tex_members, key=lambda m: m.size)
                f = tf.extractfile(largest)
                if f is None:
                    return None
                return f.read().decode("utf-8", errors="replace")
    except tarfile.TarError as exc:
        logger.warning("Failed to open tar archive for %s: %s", arxiv_id, exc)
        return None

    return None  # unreachable but satisfies type checker


def fetch_pdf_url(arxiv_id: str) -> str:
    """Return the canonical PDF URL for an arXiv paper.

    This is a pure computation — no HTTP request is made.

    Parameters
    ----------
    arxiv_id:
        The arXiv identifier (e.g. ``"2301.07041"``).

    Returns
    -------
    str
        The full URL of the PDF file.
    """
    return f"{_BASE}/pdf/{arxiv_id}.pdf"