app/utils/http_utils.py

from __future__ import annotations

import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, AsyncIterator, Dict, List, Optional, Tuple
from urllib.parse import unquote, urlparse

import aiohttp
import httpx

from app.config import get_settings

_logger = logging.getLogger(__name__)
_settings = get_settings()


class DownloadError(RuntimeError):
    """Raised when a remote download fails (HTTP error, timeout, or size limit).

    ``is_size_error`` distinguishes size-limit violations so callers can map
    the failure to the appropriate domain exception.
    """

    def __init__(self, message: str, *, is_size_error: bool = False) -> None:
        super().__init__(message)
        self.is_size_error = is_size_error


class SharedAsyncClient:
    """Thread-safe, lazily-created, connection-pooled ``httpx.AsyncClient``.

    Several services (GCS, Google Maps, Google OAuth, scheduler) previously
    duplicated the same double-checked-lock initialisation and shutdown logic.
    This helper centralises that pattern so a single instance is created on
    first use and safely closed exactly once on shutdown.
    """

    def __init__(
        self,
        *,
        timeout: Any = 30.0,
        limits: Optional[httpx.Limits] = None,
        follow_redirects: bool = True,
    ) -> None:
        self._client: Optional[httpx.AsyncClient] = None
        self._lock = asyncio.Lock()
        self._timeout = timeout
        self._limits = limits or httpx.Limits(
            max_connections=100,
            max_keepalive_connections=20,
            keepalive_expiry=30,
        )
        self._follow_redirects = follow_redirects

    async def get(self) -> httpx.AsyncClient:
        """Return the shared client, creating it lazily on first use."""
        if self._client is None or self._client.is_closed:
            async with self._lock:
                if self._client is None or self._client.is_closed:
                    self._client = httpx.AsyncClient(
                        timeout=self._timeout,
                        follow_redirects=self._follow_redirects,
                        limits=self._limits,
                    )
        return self._client

    async def close(self) -> None:
        """Close the underlying client and release pooled connections."""
        async with self._lock:
            if self._client is not None and not self._client.is_closed:
                await self._client.aclose()
            self._client = None


def create_session_timeout(total_seconds: Optional[float] = None) -> aiohttp.ClientTimeout:
    """Create an aiohttp timeout using the configured request timeout."""
    if total_seconds is None:
        total_seconds = _settings.request_timeout_ms / 1000
    return aiohttp.ClientTimeout(total=total_seconds)


class SharedAiohttpSession:
    """Reuse a single ``aiohttp.ClientSession`` per event loop.

    aiohttp sessions are bound to the loop they are created on, so this caches
    one session per event loop. In the long-lived server loop this removes the
    per-request session creation (and the associated TCP/TLS handshake) from hot
    paths such as chat completions and web search.
    """

    def __init__(self) -> None:
        self._sessions: Dict[int, aiohttp.ClientSession] = {}

    async def get(self, timeout: Optional[aiohttp.ClientTimeout] = None) -> aiohttp.ClientSession:
        loop = asyncio.get_running_loop()
        loop_id = id(loop)
        session = self._sessions.get(loop_id)
        if session is None or session.closed:
            session = aiohttp.ClientSession(timeout=timeout or create_session_timeout())
            self._sessions[loop_id] = session
        return session

    async def close_all(self) -> None:
        sessions = list(self._sessions.values())
        self._sessions.clear()
        for session in sessions:
            if session is not None and not session.closed:
                await session.close()


_shared_aiohttp_session = SharedAiohttpSession()


async def get_shared_aiohttp_session(
    timeout: Optional[aiohttp.ClientTimeout] = None,
) -> aiohttp.ClientSession:
    """Return the loop-bound shared aiohttp session, created lazily on first use."""
    return await _shared_aiohttp_session.get(timeout)


@asynccontextmanager
async def shared_aiohttp_session(
    timeout: Optional[aiohttp.ClientTimeout] = None,
) -> AsyncIterator[aiohttp.ClientSession]:
    """Context manager that yields the shared session but does NOT close it on exit.

    Drop-in replacement for ``async with aiohttp.ClientSession(...)`` so call
    sites keep the same structure while reusing one pooled session per loop.
    """
    session = await get_shared_aiohttp_session(timeout)
    yield session


async def close_shared_aiohttp_sessions() -> None:
    """Close all cached sessions. Call once on application shutdown."""
    await _shared_aiohttp_session.close_all()


async def download_url(
    url: str,
    *,
    timeout_seconds: float = 120,
    max_size_bytes: Optional[int] = None,
    chunk_size: int = 256 * 1024,
) -> Tuple[bytes, Optional[str]]:
    """Download a file from a URL with streaming and size guard.

    Returns (data, filename) where filename may be None.
    Consolidates _download_file from csv_analysis_service and
    _download_from_url from dataset_metadata_service.
    """
    timeout = aiohttp.ClientTimeout(total=timeout_seconds)
    try:
        session = await get_shared_aiohttp_session()
        async with session.get(url, timeout=timeout) as resp:
            if resp.status != 200:
                raise DownloadError(f"HTTP {resp.status} when fetching {url}")

            if max_size_bytes and resp.content_length and resp.content_length > max_size_bytes:
                raise DownloadError(
                    f"Remote file advertises {resp.content_length} bytes, "
                    f"limit is {max_size_bytes}",
                    is_size_error=True,
                )

            chunks: List[bytes] = []
            total = 0
            async for chunk in resp.content.iter_chunked(chunk_size):
                total += len(chunk)
                if max_size_bytes and total > max_size_bytes:
                    raise DownloadError(
                        f"Download exceeded {max_size_bytes} bytes",
                        is_size_error=True,
                    )
                chunks.append(chunk)

            data = b"".join(chunks)
    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
        raise DownloadError(f"Download failed for {url}: {exc}") from exc

    parsed = urlparse(url)
    filename = unquote(Path(parsed.path).name) if parsed.path else None
    return data, filename