"""SSRF-safe httpx transport with DNS resolution and IP pinning."""

import asyncio
import socket

import httpx

from langchain_core._security._exceptions import SSRFBlockedError
from langchain_core._security._policy import (
    DEFAULT_SSRF_POLICY,
    SSRFPolicy,
    _effective_allowed_hosts,
    validate_resolved_ip,
    validate_url_sync,
)

# Keys that AsyncHTTPTransport accepts (forwarded from factory kwargs).
_TRANSPORT_KWARGS = frozenset(
    {
        "verify",
        "cert",
        "trust_env",
        "http1",
        "http2",
        "limits",
        "retries",
    }
)


class SSRFSafeTransport(httpx.AsyncBaseTransport):
    """httpx async transport that validates DNS results against an SSRF policy.

    For every outgoing request the transport:
    1. Checks the URL scheme against `policy.allowed_schemes`.
    2. Validates the hostname against blocked patterns.
    3. Resolves DNS and validates **all** returned IPs.
    4. Rewrites the request to connect to the first valid IP while
       preserving the original `Host` header and TLS SNI hostname.

    Redirects are re-validated on each hop because `follow_redirects`
    is set on the *client*, causing `handle_async_request` to be called
    again for each redirect target.
    """

    def __init__(
        self,
        policy: SSRFPolicy = DEFAULT_SSRF_POLICY,
        **transport_kwargs: object,
    ) -> None:
        self._policy = policy
        self._inner = httpx.AsyncHTTPTransport(**transport_kwargs)  # type: ignore[arg-type]

    # ------------------------------------------------------------------ #
    # Core request handler
    # ------------------------------------------------------------------ #

    async def handle_async_request(
        self,
        request: httpx.Request,
    ) -> httpx.Response:
        hostname = request.url.host or ""
        scheme = request.url.scheme.lower()

        # 1-3. Scheme, hostname, and pattern checks (reuse sync validator).
        validate_url_sync(str(request.url), self._policy)

        # Allowed-hosts bypass - skip DNS/IP validation entirely.
        allowed = {h.lower() for h in _effective_allowed_hosts(self._policy)}
        if hostname.lower() in allowed:
            return await self._inner.handle_async_request(request)

        # 4. DNS resolution
        port = request.url.port or (443 if scheme == "https" else 80)
        try:
            addrinfo = await asyncio.to_thread(
                socket.getaddrinfo,
                hostname,
                port,
                type=socket.SOCK_STREAM,
            )
        except socket.gaierror as exc:
            msg = "DNS resolution failed"
            raise SSRFBlockedError(msg) from exc

        if not addrinfo:
            msg = "DNS resolution returned no results"
            raise SSRFBlockedError(msg)

        # 5. Validate ALL resolved IPs - any blocked means reject.
        for _family, _type, _proto, _canonname, sockaddr in addrinfo:
            ip_str: str = sockaddr[0]  # type: ignore[assignment]
            validate_resolved_ip(ip_str, self._policy)

        # 6. Pin to first resolved IP.
        pinned_ip = addrinfo[0][4][0]

        # 7. Rewrite URL to use pinned IP, preserving Host header and SNI.
        pinned_url = request.url.copy_with(host=pinned_ip)

        # Build extensions dict, adding sni_hostname for HTTPS so TLS
        # certificate validation uses the original hostname.
        extensions = dict(request.extensions)
        if scheme == "https":
            extensions["sni_hostname"] = hostname.encode("ascii")

        pinned_request = httpx.Request(
            method=request.method,
            url=pinned_url,
            headers=request.headers,  # Host header already set to original
            content=request.content,
            extensions=extensions,
        )

        return await self._inner.handle_async_request(pinned_request)

    # ------------------------------------------------------------------ #
    # Lifecycle
    # ------------------------------------------------------------------ #

    async def aclose(self) -> None:
        await self._inner.aclose()


# ---------------------------------------------------------------------- #
# Factory
# ---------------------------------------------------------------------- #


class SSRFSafeSyncTransport(httpx.BaseTransport):
    """httpx sync transport that validates DNS results against an SSRF policy.

    Sync mirror of `SSRFSafeTransport`. See that class for full documentation.
    """

    def __init__(
        self,
        policy: SSRFPolicy = DEFAULT_SSRF_POLICY,
        **transport_kwargs: object,
    ) -> None:
        self._policy = policy
        self._inner = httpx.HTTPTransport(**transport_kwargs)  # type: ignore[arg-type]

    def handle_request(
        self,
        request: httpx.Request,
    ) -> httpx.Response:
        hostname = request.url.host or ""
        scheme = request.url.scheme.lower()

        validate_url_sync(str(request.url), self._policy)

        allowed = {h.lower() for h in _effective_allowed_hosts(self._policy)}
        if hostname.lower() in allowed:
            return self._inner.handle_request(request)

        port = request.url.port or (443 if scheme == "https" else 80)
        try:
            addrinfo = socket.getaddrinfo(
                hostname,
                port,
                type=socket.SOCK_STREAM,
            )
        except socket.gaierror as exc:
            msg = "DNS resolution failed"
            raise SSRFBlockedError(msg) from exc

        if not addrinfo:
            msg = "DNS resolution returned no results"
            raise SSRFBlockedError(msg)

        for _family, _type, _proto, _canonname, sockaddr in addrinfo:
            ip_str: str = sockaddr[0]  # type: ignore[assignment]
            validate_resolved_ip(ip_str, self._policy)

        pinned_ip = addrinfo[0][4][0]
        pinned_url = request.url.copy_with(host=pinned_ip)

        extensions = dict(request.extensions)
        if scheme == "https":
            extensions["sni_hostname"] = hostname.encode("ascii")

        pinned_request = httpx.Request(
            method=request.method,
            url=pinned_url,
            headers=request.headers,
            content=request.content,
            extensions=extensions,
        )

        return self._inner.handle_request(pinned_request)

    def close(self) -> None:
        self._inner.close()


# ---------------------------------------------------------------------- #
# Factories
# ---------------------------------------------------------------------- #


def ssrf_safe_client(
    policy: SSRFPolicy = DEFAULT_SSRF_POLICY,
    **kwargs: object,
) -> httpx.Client:
    """Create an `httpx.Client` with SSRF protection."""
    transport_kwargs: dict[str, object] = {}
    client_kwargs: dict[str, object] = {}
    for key, value in kwargs.items():
        if key in _TRANSPORT_KWARGS:
            transport_kwargs[key] = value
        else:
            client_kwargs[key] = value

    transport = SSRFSafeSyncTransport(policy=policy, **transport_kwargs)

    client_kwargs.setdefault("follow_redirects", True)
    client_kwargs.setdefault("max_redirects", 10)

    return httpx.Client(
        transport=transport,
        **client_kwargs,  # type: ignore[arg-type]
    )


def ssrf_safe_async_client(
    policy: SSRFPolicy = DEFAULT_SSRF_POLICY,
    **kwargs: object,
) -> httpx.AsyncClient:
    """Create an `httpx.AsyncClient` with SSRF protection.

    Drop-in replacement for `httpx.AsyncClient(...)` - callers just swap
    the constructor call.  Transport-specific kwargs (`verify`, `cert`,
    `retries`, etc.) are forwarded to the inner `AsyncHTTPTransport`;
    everything else goes to the `AsyncClient`.
    """
    transport_kwargs: dict[str, object] = {}
    client_kwargs: dict[str, object] = {}
    for key, value in kwargs.items():
        if key in _TRANSPORT_KWARGS:
            transport_kwargs[key] = value
        else:
            client_kwargs[key] = value

    transport = SSRFSafeTransport(policy=policy, **transport_kwargs)

    # Apply defaults only if not overridden by caller.
    client_kwargs.setdefault("follow_redirects", True)
    client_kwargs.setdefault("max_redirects", 10)

    return httpx.AsyncClient(
        transport=transport,
        **client_kwargs,  # type: ignore[arg-type]
    )
