"""
Safe HTTP Requests Wrapper

Wraps requests library to add SSRF protection and security best practices.
"""

import datetime
import email.utils
import time
from typing import Any, Optional
from urllib.parse import urljoin

import requests
from loguru import logger

from . import ssrf_validator
from ..constants import USER_AGENT
from ..utilities.resource_utils import safe_close


# Default timeout for all HTTP requests (prevents hanging)
DEFAULT_TIMEOUT = 30  # seconds

# Maximum response size to prevent memory exhaustion (1GB)
# Set high to accommodate large documents (annual reports, PDFs, datasets).
# This is a local research tool — users intentionally download these files.
MAX_RESPONSE_SIZE = 1024 * 1024 * 1024

# HTTP status codes that indicate a redirect
_REDIRECT_STATUS_CODES = frozenset({301, 302, 303, 307, 308})

# Maximum number of redirects to follow
_MAX_REDIRECTS = 10

# Prefix of the requests library's default User-Agent
# (e.g. "python-requests/2.32.5"). Used by `SafeSession.request` to detect
# whether the session is using the upstream default — if so, we override
# it with the project `USER_AGENT` so academic API endpoints (arXiv,
# OpenAlex, PubMed, …) can identify us. Promoting the literal to a
# constant so a future requests rename only requires a one-line edit.
_DEFAULT_REQUESTS_UA_PREFIX = "python-requests"


def _install_body_guard(response: requests.Response) -> None:
    """Install a bounded reader that enforces MAX_RESPONSE_SIZE.

    Wraps response.raw.read() to track cumulative bytes and raise
    ValueError if MAX_RESPONSE_SIZE is exceeded during body consumption.
    This transparently protects both streamed (.iter_content) and
    non-streamed (.text, .json(), .content) access patterns.

    Always installs — callers (currently only _check_response_size)
    are responsible for deciding when to call this function.
    """
    original_read = response.raw.read
    bytes_read = 0

    def bounded_read(amt=None, *args, **kwargs):
        nonlocal bytes_read
        data = original_read(amt, *args, **kwargs)
        bytes_read += len(data)
        if bytes_read > MAX_RESPONSE_SIZE:
            response.close()
            raise ValueError(
                f"Response body too large: >{bytes_read} bytes "
                f"(max {MAX_RESPONSE_SIZE}, Content-Length absent or invalid)"
            )
        return data

    response.raw.read = bounded_read  # type: ignore[method-assign]


def _check_response_size(response: requests.Response) -> None:
    """Reject responses whose Content-Length exceeds MAX_RESPONSE_SIZE.

    Handles comma-separated values per RFC 7230 §3.3.2: identical
    duplicates (from proxies) are normalized; differing values are
    rejected as invalid framing. Empty parts from malformed headers
    (trailing/doubled commas) are filtered before parsing. Non-integer
    or negative values cause the header to be treated as absent.

    When Content-Length is absent, unparseable, negative, or consists
    only of commas/whitespace, installs a body guard that enforces
    the size limit during body consumption.

    Must be called before returning a response to the caller. On
    rejection the response is closed to avoid leaking the connection.

    Raises:
        ValueError: If Content-Length values conflict or exceed
            MAX_RESPONSE_SIZE.
    """
    content_length = response.headers.get("Content-Length")
    if content_length:
        try:
            # Handle comma-separated Content-Length values (RFC 7230 §3.3.2).
            # Multiple identical values may be sent by proxies; differing
            # values indicate invalid framing and must be rejected.
            raw_parts = [v.strip() for v in content_length.split(",")]
            parts = [p for p in raw_parts if p]
            if not parts:
                _install_body_guard(response)
                return  # Only commas/whitespace — treat as absent
            sizes = [int(p) for p in parts]
        except (ValueError, TypeError):
            _install_body_guard(response)
            return  # Content-Length not a valid number
        if len(set(sizes)) > 1:
            response.close()
            raise ValueError(
                f"Conflicting Content-Length values: {content_length}"
            )
        size = sizes[0]
        if size < 0:
            _install_body_guard(response)
            return  # Malformed Content-Length, treat as absent
        if size > MAX_RESPONSE_SIZE:
            response.close()
            raise ValueError(
                f"Response too large: {size} bytes (max {MAX_RESPONSE_SIZE})"
            )
        # Valid Content-Length within limit — no body guard needed
        return

    # No Content-Length header at all — install body guard
    _install_body_guard(response)


def _resolve_redirect_method(method: str, status_code: int) -> str:
    """Determine HTTP method after redirect, per RFC 7231."""
    if status_code == 303 and method != "HEAD":
        method = "GET"
    elif status_code == 302 and method == "POST":
        method = "GET"
    elif status_code == 301 and method == "POST":
        method = "GET"
    # 307, 308: preserve original method (no change needed)
    return method


def safe_get(
    url: str,
    params: Optional[dict] = None,
    timeout: int = DEFAULT_TIMEOUT,
    allow_localhost: bool = False,
    allow_private_ips: bool = False,
    **kwargs,
) -> requests.Response:
    """
    Make a safe HTTP GET request with SSRF protection.

    Args:
        url: URL to request
        params: URL parameters
        timeout: Request timeout in seconds
        allow_localhost: Whether to allow localhost/loopback addresses.
            Set to True for trusted internal services like self-hosted
            search engines (e.g., searxng). Default False.
        allow_private_ips: Whether to allow all private/internal IPs plus localhost.
            This includes RFC1918 (10.x, 172.16-31.x, 192.168.x), CGNAT (100.64.x.x
            used by Podman/rootless containers), link-local (169.254.x.x), and IPv6
            private ranges (fc00::/7, fe80::/10). Use for trusted self-hosted services
            like SearXNG or Ollama in containerized environments.
            Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
            AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
            ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
        **kwargs: Additional arguments to pass to requests.get()

    Returns:
        Response object

    Raises:
        ValueError: If URL fails SSRF validation
        requests.RequestException: If request fails
    """
    # Validate URL to prevent SSRF
    if not ssrf_validator.validate_url(
        url,
        allow_localhost=allow_localhost,
        allow_private_ips=allow_private_ips,
    ):
        raise ValueError(
            f"URL failed security validation (possible SSRF): {url}"
        )

    # Ensure timeout is set
    if "timeout" not in kwargs:
        kwargs["timeout"] = timeout

    # Inject the project User-Agent if the caller didn't supply one.
    # Mutates a copy of any caller-supplied headers dict so we never
    # touch their object.
    headers = dict(kwargs.get("headers") or {})
    if not any(k.lower() == "user-agent" for k in headers):
        headers["User-Agent"] = USER_AGENT
        kwargs["headers"] = headers

    # Intercept allow_redirects — we handle redirects manually to validate
    # each redirect target against SSRF rules
    caller_wants_redirects = kwargs.pop("allow_redirects", True)
    kwargs["allow_redirects"] = False

    current_url = url
    try:
        response = requests.get(url, params=params, **kwargs)

        # Follow redirects manually with SSRF validation on each hop.
        # Each hop uses a fresh requests.get() call without a session,
        # so cookies set by intermediate responses are not carried
        # forward. This is acceptable for current callers (all stateless).
        # Callers needing cookie persistence across redirects should use
        # SafeSession instead, which preserves cookies via its cookie jar.
        if caller_wants_redirects:
            redirects_followed = 0
            while (
                response.status_code in _REDIRECT_STATUS_CODES
                and redirects_followed < _MAX_REDIRECTS
            ):
                redirect_url = (response.headers.get("Location") or "").strip()
                if not redirect_url:
                    break

                # Resolve relative redirects
                redirect_url = urljoin(
                    response.url or current_url, redirect_url
                )

                # Validate redirect target against SSRF rules
                if not ssrf_validator.validate_url(
                    redirect_url,
                    allow_localhost=allow_localhost,
                    allow_private_ips=allow_private_ips,
                ):
                    logger.warning(
                        f"Redirect to {redirect_url} blocked by SSRF validation "
                        f"(from {url}, hop {redirects_followed + 1})"
                    )
                    response.close()
                    raise ValueError(
                        f"Redirect target failed SSRF validation: {redirect_url}"
                    )

                current_url = redirect_url
                response.close()
                # Note: params are intentionally NOT forwarded to redirect
                # hops. Per HTTP spec, the server's Location header contains
                # the complete target URL. Re-appending original query params
                # would corrupt it.
                response = requests.get(redirect_url, **kwargs)
                redirects_followed += 1

            if (
                response.status_code in _REDIRECT_STATUS_CODES
                and redirects_followed >= _MAX_REDIRECTS
            ):
                response.close()
                # Note: raises ValueError here, while SafeSession raises
                # requests.TooManyRedirects (delegated to the base class).
                # Callers should catch ValueError for standalone functions.
                raise ValueError(
                    f"Too many redirects ({_MAX_REDIRECTS}) from {url}"
                )

        _check_response_size(response)

        return response

    except requests.Timeout:
        logger.warning(f"Request timeout after {timeout}s: {current_url}")
        raise
    except requests.RequestException:
        logger.warning(f"Request failed for {current_url}")
        raise


def safe_post(
    url: str,
    data: Optional[Any] = None,
    json: Optional[dict] = None,
    timeout: int = DEFAULT_TIMEOUT,
    allow_localhost: bool = False,
    allow_private_ips: bool = False,
    **kwargs,
) -> requests.Response:
    """
    Make a safe HTTP POST request with SSRF protection.

    Args:
        url: URL to request
        data: Data to send in request body
        json: JSON data to send in request body
        timeout: Request timeout in seconds
        allow_localhost: Whether to allow localhost/loopback addresses.
            Set to True for trusted internal services like self-hosted
            search engines (e.g., searxng). Default False.
        allow_private_ips: Whether to allow all private/internal IPs plus localhost.
            This includes RFC1918 (10.x, 172.16-31.x, 192.168.x), CGNAT (100.64.x.x
            used by Podman/rootless containers), link-local (169.254.x.x), and IPv6
            private ranges (fc00::/7, fe80::/10). Use for trusted self-hosted services
            like SearXNG or Ollama in containerized environments.
            Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
            AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
            ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
        **kwargs: Additional arguments to pass to requests.post()

    Returns:
        Response object

    Raises:
        ValueError: If URL fails SSRF validation
        requests.RequestException: If request fails
    """
    # Validate URL to prevent SSRF
    if not ssrf_validator.validate_url(
        url,
        allow_localhost=allow_localhost,
        allow_private_ips=allow_private_ips,
    ):
        raise ValueError(
            f"URL failed security validation (possible SSRF): {url}"
        )

    # Ensure timeout is set
    if "timeout" not in kwargs:
        kwargs["timeout"] = timeout

    # Inject the project User-Agent if the caller didn't supply one.
    # Mutates a copy of any caller-supplied headers dict so we never
    # touch their object.
    headers = dict(kwargs.get("headers") or {})
    if not any(k.lower() == "user-agent" for k in headers):
        headers["User-Agent"] = USER_AGENT
        kwargs["headers"] = headers

    # Intercept allow_redirects — we handle redirects manually to validate
    # each redirect target against SSRF rules
    caller_wants_redirects = kwargs.pop("allow_redirects", True)
    kwargs["allow_redirects"] = False

    current_url = url
    try:
        response = requests.post(url, data=data, json=json, **kwargs)

        # Follow redirects manually with SSRF validation on each hop.
        # Each hop uses a fresh request without a session, so cookies
        # set by intermediate responses are not carried forward. Callers
        # needing cookie persistence should use SafeSession instead.
        if caller_wants_redirects:
            redirect_method = "POST"
            redirects_followed = 0
            while (
                response.status_code in _REDIRECT_STATUS_CODES
                and redirects_followed < _MAX_REDIRECTS
            ):
                redirect_url = (response.headers.get("Location") or "").strip()
                if not redirect_url:
                    break

                # Resolve relative redirects
                redirect_url = urljoin(
                    response.url or current_url, redirect_url
                )

                # Validate redirect target against SSRF rules
                if not ssrf_validator.validate_url(
                    redirect_url,
                    allow_localhost=allow_localhost,
                    allow_private_ips=allow_private_ips,
                ):
                    logger.warning(
                        f"Redirect to {redirect_url} blocked by SSRF validation "
                        f"(from {url}, hop {redirects_followed + 1})"
                    )
                    response.close()
                    raise ValueError(
                        f"Redirect target failed SSRF validation: {redirect_url}"
                    )

                redirect_method = _resolve_redirect_method(
                    redirect_method, response.status_code
                )
                current_url = redirect_url
                response.close()

                if redirect_method == "GET":
                    # 301/302/303: convert to GET, drop body
                    data = None
                    json = None
                    response = requests.get(redirect_url, **kwargs)
                else:
                    # 307/308: preserve current method and body
                    response = requests.post(
                        redirect_url, data=data, json=json, **kwargs
                    )
                redirects_followed += 1

            if (
                response.status_code in _REDIRECT_STATUS_CODES
                and redirects_followed >= _MAX_REDIRECTS
            ):
                response.close()
                # Note: raises ValueError here, while SafeSession raises
                # requests.TooManyRedirects (delegated to the base class).
                # Callers should catch ValueError for standalone functions.
                raise ValueError(
                    f"Too many redirects ({_MAX_REDIRECTS}) from {url}"
                )

        _check_response_size(response)

        return response

    except requests.Timeout:
        logger.warning(f"Request timeout after {timeout}s: {current_url}")
        raise
    except requests.RequestException:
        logger.warning(f"Request failed for {current_url}")
        raise


# Create a safe session class
class SafeSession(requests.Session):
    """
    Session with built-in SSRF protection.

    Redirect validation relies on ``requests.Session.resolve_redirects()``
    calling ``self.send()`` for each hop — an internal implementation detail
    of the ``requests`` library.  This is simpler than re-implementing the
    redirect loop (as ``safe_get``/``safe_post`` do) and keeps session-level
    features (cookies, auth) working.  The trade-off is coupling to the
    ``requests`` internals; if a future version stops routing hops through
    ``send()``, redirect targets would no longer be validated.

    Usage:
        with SafeSession() as session:
            response = session.get(url)

        # For trusted internal services (e.g., searxng on localhost):
        with SafeSession(allow_localhost=True) as session:
            response = session.get(url)

        # For trusted internal services on any private network IP:
        with SafeSession(allow_private_ips=True) as session:
            response = session.get(url)

    Raises:
        ValueError: If a URL (initial or redirect target) fails SSRF
            validation, or if the response Content-Length exceeds
            MAX_RESPONSE_SIZE.  Note: ``safe_get``/``safe_post`` also raise
            ``ValueError`` for too-many-redirects, but ``SafeSession`` raises
            ``requests.TooManyRedirects`` for that case since it delegates
            redirect counting to the ``requests`` library.
        requests.RequestException: On transport-level failures.
    """

    def __init__(
        self, allow_localhost: bool = False, allow_private_ips: bool = False
    ):
        """
        Initialize SafeSession.

        Args:
            allow_localhost: Whether to allow localhost/loopback addresses.
            allow_private_ips: Whether to allow all private/internal IPs plus localhost.
                This includes RFC1918, CGNAT (100.64.x.x used by Podman), link-local, and
                IPv6 private ranges. Use for trusted self-hosted services like SearXNG or
                Ollama in containerized environments.
                Note: cloud metadata endpoints (AWS / Azure / OCI / DigitalOcean /
                AlibabaCloud / Tencent / ECS) are ALWAYS blocked — see
                ``ssrf_validator.ALWAYS_BLOCKED_METADATA_IPS``.
        """
        super().__init__()
        self.max_redirects = _MAX_REDIRECTS
        self.allow_localhost = allow_localhost
        self.allow_private_ips = allow_private_ips

    def request(self, method: str, url: str, **kwargs) -> requests.Response:  # type: ignore[override]
        """Override request method to add SSRF validation."""
        # Validate URL
        if not ssrf_validator.validate_url(
            url,
            allow_localhost=self.allow_localhost,
            allow_private_ips=self.allow_private_ips,
        ):
            raise ValueError(
                f"URL failed security validation (possible SSRF): {url}"
            )

        # Ensure timeout is set
        if "timeout" not in kwargs:
            kwargs["timeout"] = DEFAULT_TIMEOUT

        # Inject project User-Agent if the caller didn't already set one.
        # Session-level User-Agent (self.headers) is left alone — only
        # per-request headers are copied so we never mutate caller state.
        headers = dict(kwargs.get("headers") or {})
        session_ua = self.headers.get("User-Agent", "")
        has_per_request_ua = any(k.lower() == "user-agent" for k in headers)
        if not has_per_request_ua and (
            not session_ua or session_ua.startswith(_DEFAULT_REQUESTS_UA_PREFIX)
        ):
            headers["User-Agent"] = USER_AGENT
            kwargs["headers"] = headers

        return super().request(method, url, **kwargs)

    def send(
        self, request: requests.PreparedRequest, **kwargs
    ) -> requests.Response:
        """Override send to validate every outgoing request against SSRF.

        This runs on **all** calls — both the initial request (routed
        here by ``requests.Session.request()``) and each redirect hop
        (routed here by ``resolve_redirects()``).  The initial URL is
        therefore validated twice (once in ``request()``, once here);
        this is intentional defense-in-depth.
        """
        if request.url and not ssrf_validator.validate_url(
            request.url,
            allow_localhost=self.allow_localhost,
            allow_private_ips=self.allow_private_ips,
        ):
            logger.warning(
                f"Request to {request.url} blocked by SSRF validation"
            )
            # Note: This error says "security validation" while safe_get/
            # safe_post say "SSRF validation". The difference indicates the
            # source (session vs standalone function) in logs.
            raise ValueError(
                f"Redirect target failed security validation (possible SSRF): {request.url}"
            )

        response = super().send(request, **kwargs)
        _check_response_size(response)
        return response


# Exponential backoff schedule (seconds). Kept short: journal-quality
# downloads are run from a user request or a scheduled job, not from a
# time-sensitive hot path, so three retries over ~7 seconds is plenty
# without adding real latency.
_RETRY_BACKOFF_SECONDS = (1, 2, 4)

# HTTP status codes worth retrying (transient server / rate-limit errors).
_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504})

# Upper bound on a honored Retry-After (seconds). RFC 7231 puts no
# ceiling on the header, so a hostile or misconfigured upstream could
# pin a worker via an arbitrarily large value. Cap here to bound the
# damage; legitimate waits (seconds to low minutes) pass through.
_MAX_RETRY_AFTER_SECONDS = 300


def _parse_retry_after(retry_after_raw: Optional[str]) -> Optional[int]:
    """Parse a ``Retry-After`` header value, clamped to ``[0, MAX]``.

    Returns ``None`` if the header is missing or unparseable, so the
    caller can fall back to the exponential-backoff schedule. Accepts
    both RFC 7231 forms: delay-seconds (integer) and HTTP-date.
    """
    if retry_after_raw is None:
        return None
    try:
        seconds = int(retry_after_raw)
    except ValueError:
        try:
            retry_dt = email.utils.parsedate_to_datetime(retry_after_raw)
        except (ValueError, TypeError):
            logger.debug(
                f"Unparseable Retry-After {retry_after_raw!r}; "
                f"using backoff schedule"
            )
            return None
        now_utc = datetime.datetime.now(datetime.timezone.utc)
        seconds = int((retry_dt - now_utc).total_seconds())
    return max(0, min(seconds, _MAX_RETRY_AFTER_SECONDS))


def safe_get_with_retries(
    url: str,
    params: Optional[dict] = None,
    timeout: int = DEFAULT_TIMEOUT,
    allow_localhost: bool = False,
    allow_private_ips: bool = False,
    max_retries: int = 3,
    backoff_times: tuple = _RETRY_BACKOFF_SECONDS,
    consume_body: bool = False,
    **kwargs,
) -> requests.Response:
    """`safe_get` plus exponential-backoff retry on transient errors.

    Retries on:
      * ``requests.ConnectionError``
      * ``requests.Timeout``
      * HTTP ``429`` (rate limit) and ``5xx`` (server error)
      * (when ``consume_body=True``) body-read failures —
        ``ChunkedEncodingError``, ``ReadTimeout``, mid-stream
        ``ConnectionError``

    Honors the ``Retry-After`` header when present (falls back to the
    backoff schedule otherwise). SSRF-validation errors (``ValueError``)
    and non-retryable HTTP 4xx responses are not retried.

    Without ``consume_body``, only failures raised inside ``safe_get``
    itself (DNS, connect, header timeout, retryable status) trigger a
    retry. The body isn't read until the caller touches ``.content`` /
    ``.text`` / ``.json()``, by which point this wrapper has already
    returned — so a mid-stream S3 hiccup (``ChunkedEncodingError``)
    propagates uncaught. ``consume_body=True`` reads the body inside
    the retry loop so those transient body-read failures are also
    retried. The cached body is still available to the caller via
    ``response.content`` after the wrapper returns.

    Args:
        url: Target URL.
        params: Query parameters.
        timeout: Per-attempt socket timeout.
        allow_localhost: Forwarded to ``safe_get``.
        allow_private_ips: Forwarded to ``safe_get``.
        max_retries: Maximum retry attempts after the initial try.
        backoff_times: Per-attempt sleep seconds.
        consume_body: If True, read ``response.content`` inside the
            retry loop so body-read transients are retried. Use for
            large or chunk-transferred bodies (~MB+) where mid-stream
            disconnects are realistic. The body-guard's ``ValueError``
            (oversized body) is NOT retried — it propagates immediately.
        **kwargs: Forwarded to ``safe_get``.

    Returns:
        The first successful (or final-attempt) ``requests.Response``.
        When ``consume_body=True``, the body has already been read and
        is cached on the response.

    Raises:
        ValueError: If SSRF validation fails or, with
            ``consume_body=True``, the body-guard rejects an oversized
            response. Retries do not help in either case.
        requests.RequestException: If every attempt fails.
    """
    attempt = 0
    while True:
        try:
            response = safe_get(
                url,
                params=params,
                timeout=timeout,
                allow_localhost=allow_localhost,
                allow_private_ips=allow_private_ips,
                **kwargs,
            )
        except (requests.ConnectionError, requests.Timeout) as exc:
            if attempt >= max_retries:
                raise
            wait = backoff_times[min(attempt, len(backoff_times) - 1)]
            logger.warning(
                f"{exc.__class__.__name__} on {url}; "
                f"retrying in {wait}s "
                f"(attempt {attempt + 1}/{max_retries})"
            )
            time.sleep(wait)
            attempt += 1
            continue

        if response.status_code in _RETRYABLE_STATUS_CODES:
            if attempt >= max_retries:
                return response
            parsed = _parse_retry_after(response.headers.get("Retry-After"))
            wait = (
                parsed
                if parsed is not None
                else backoff_times[min(attempt, len(backoff_times) - 1)]
            )
            logger.warning(
                f"HTTP {response.status_code} on {url}; "
                f"retrying in {wait}s "
                f"(attempt {attempt + 1}/{max_retries})"
            )
            response.close()
            time.sleep(wait)
            attempt += 1
            continue

        if consume_body:
            try:
                # Force body read while still inside the retry loop.
                # ChunkedEncodingError / ReadTimeout / mid-stream
                # ConnectionError can fire here on large responses
                # from flaky upstreams. ReadTimeout is a Timeout
                # subclass but NOT a ConnectionError subclass, so the
                # except must list both ConnectionError and Timeout
                # — listing Timeout alone would miss ConnectError, and
                # listing ConnectionError alone would miss ReadTimeout.
                _ = response.content
            except (
                requests.exceptions.ChunkedEncodingError,
                requests.exceptions.ConnectionError,
                requests.exceptions.Timeout,
            ) as exc:
                # safe_close instead of bare close: a close() that
                # raises here would mask the original body-read error
                # we actually want to surface / retry on.
                safe_close(response, "response")
                if attempt >= max_retries:
                    raise
                wait = backoff_times[min(attempt, len(backoff_times) - 1)]
                logger.warning(
                    f"{exc.__class__.__name__} reading body of {url}; "
                    f"retrying in {wait}s "
                    f"(attempt {attempt + 1}/{max_retries})"
                )
                time.sleep(wait)
                attempt += 1
                continue

        return response
