"""
Safe HTTP Requests Wrapper

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

import requests
from typing import Any, Optional
from urllib.parse import urljoin

from loguru import logger

from . import ssrf_validator


# 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


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: AWS metadata endpoint (169.254.169.254) is ALWAYS blocked.
        **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

    # 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: AWS metadata endpoint (169.254.169.254) is ALWAYS blocked.
        **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

    # 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: AWS metadata endpoint (169.254.169.254) is ALWAYS blocked.
        """
        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

        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
