import asyncio
import ipaddress
import re
import socket
import ssl
from io import BytesIO
from typing import Any, Callable, Optional
from urllib.parse import ParseResult as URL
from urllib.parse import quote, urljoin, urlparse

import aiohttp
import idna
from aiohttp import FormData, abc
from tenacity import (
    RetryCallState,
    retry,
    retry_if_result,
    stop_after_attempt,
    wait_exponential_jitter,
)

from backend.util.json import loads


class HTTPClientError(Exception):
    """4xx client errors (400-499)"""

    def __init__(self, message: str, status_code: int):
        super().__init__(message)
        self.status_code = status_code


class HTTPServerError(Exception):
    """5xx server errors (500-599)"""

    def __init__(self, message: str, status_code: int):
        super().__init__(message)
        self.status_code = status_code


# Default User-Agent for all requests
DEFAULT_USER_AGENT = "AutoGPT-Platform/1.0 (https://github.com/Significant-Gravitas/AutoGPT; info@agpt.co) aiohttp"

# Retry status codes for which we will automatically retry the request
THROTTLE_RETRY_STATUS_CODES: set[int] = {429, 500, 502, 503, 504, 408}

# List of IP networks to block
BLOCKED_IP_NETWORKS = [
    # --8<-- [start:BLOCKED_IP_NETWORKS]
    # IPv4 Ranges
    ipaddress.ip_network("0.0.0.0/8"),  # "This" Network
    ipaddress.ip_network("10.0.0.0/8"),  # Private-Use
    ipaddress.ip_network("100.64.0.0/10"),  # Shared Address Space (CGNAT, RFC 6598)
    ipaddress.ip_network("127.0.0.0/8"),  # Loopback
    ipaddress.ip_network("169.254.0.0/16"),  # Link Local
    ipaddress.ip_network("172.16.0.0/12"),  # Private-Use
    ipaddress.ip_network("192.0.0.0/24"),  # IETF Protocol Assignments
    ipaddress.ip_network("192.168.0.0/16"),  # Private-Use
    ipaddress.ip_network("198.18.0.0/15"),  # Benchmarking
    ipaddress.ip_network("224.0.0.0/4"),  # Multicast
    ipaddress.ip_network("240.0.0.0/4"),  # Reserved for Future Use
    # IPv6 Ranges
    ipaddress.ip_network("::1/128"),  # Loopback
    ipaddress.ip_network("fc00::/7"),  # Unique local addresses (ULA)
    ipaddress.ip_network("fe80::/10"),  # Link-local
    ipaddress.ip_network("ff00::/8"),  # Multicast
    # --8<-- [end:BLOCKED_IP_NETWORKS]
]

ALLOWED_SCHEMES = ["http", "https"]
HOSTNAME_REGEX = re.compile(r"^[A-Za-z0-9.-]+$")  # Basic DNS-safe hostname pattern


def _is_ip_blocked(ip: str) -> bool:
    """
    Checks if the IP address is in a blocked network.

    IPv4-mapped IPv6 addresses (e.g. ``::ffff:127.0.0.1``) are normalized to
    their IPv4 equivalent before checking, so the IPv4 blocklist cannot be
    bypassed by encoding a private IPv4 address as IPv6.
    """
    ip_addr = ipaddress.ip_address(ip)

    # Normalize IPv4-mapped IPv6 → IPv4 so the IPv4 blocklist applies
    if isinstance(ip_addr, ipaddress.IPv6Address) and ip_addr.ipv4_mapped:
        ip_addr = ip_addr.ipv4_mapped

    return any(ip_addr in network for network in BLOCKED_IP_NETWORKS)


def _remove_insecure_headers(headers: dict, old_url: URL, new_url: URL) -> dict:
    """
    Removes sensitive headers (Authorization, Proxy-Authorization, Cookie)
    if the scheme/host/port of new_url differ from old_url.
    """
    if (
        (old_url.scheme != new_url.scheme)
        or (old_url.hostname != new_url.hostname)
        or (old_url.port != new_url.port)
    ):
        headers.pop("Authorization", None)
        headers.pop("Proxy-Authorization", None)
        headers.pop("Cookie", None)
    return headers


class HostResolver(abc.AbstractResolver):
    """
    A custom resolver that connects to specified IP addresses but still
    sets the TLS SNI to the original host name so the cert can match.
    """

    def __init__(self, ssl_hostname: str, ip_addresses: list[str]):
        self.ssl_hostname = ssl_hostname
        self.ip_addresses = ip_addresses
        self._default = aiohttp.ThreadedResolver()

    async def resolve(self, host, port=0, family=socket.AF_INET):
        if host == self.ssl_hostname:
            results = []
            for ip in self.ip_addresses:
                results.append(
                    {
                        "hostname": self.ssl_hostname,
                        "host": ip,
                        "port": port,
                        "family": family,
                        "proto": 0,
                        "flags": socket.AI_NUMERICHOST,
                    }
                )
            return results
        return await self._default.resolve(host, port, family)

    async def close(self):
        await self._default.close()


async def _resolve_host(hostname: str) -> list[str]:
    """
    Resolves the hostname to a list of IP addresses (IPv4 first, then IPv6).
    """
    loop = asyncio.get_running_loop()
    try:
        infos = await loop.getaddrinfo(hostname, None)
    except socket.gaierror:
        raise ValueError(f"Unable to resolve IP address for hostname {hostname}")

    ip_list = [info[4][0] for info in infos]
    ipv4 = [ip for ip in ip_list if ":" not in ip]
    ipv6 = [ip for ip in ip_list if ":" in ip]
    ip_addresses = ipv4 + ipv6

    if not ip_addresses:
        raise ValueError(f"No IP addresses found for {hostname}")
    return ip_addresses


async def validate_url_host(
    url: str, trusted_hostnames: Optional[list[str]] = None
) -> tuple[URL, bool, list[str]]:
    """
    Validates a (URL's) host string to prevent SSRF attacks by ensuring it does not
    point to a private, link-local, or otherwise blocked IP address — unless
    the hostname is explicitly trusted.

    Hosts in `trusted_hostnames` are permitted without checks.
    All other hosts are resolved and checked against `BLOCKED_IP_NETWORKS`.

    Params:
        url: A hostname, netloc, or URL to validate.
             If no scheme is included, `http://` is assumed.
        trusted_hostnames: A list of hostnames that don't require validation.

    Raises:
      ValueError:
        - if the URL has a disallowed URL scheme
        - if the URL/host string can't be parsed
        - if the hostname contains invalid or unsupported (non-ASCII) characters
        - if the host resolves to a blocked IP

    Returns:
        1. The validated, canonicalized, parsed host/URL,
           with hostname ASCII-safe encoding
        2. Whether the host is trusted (based on the passed `trusted_hostnames`).
        3. List of resolved IP addresses for the host; empty if the host is trusted.
    """
    parsed = parse_url(url)

    if parsed.scheme not in ALLOWED_SCHEMES:
        raise ValueError(
            f"URL scheme '{parsed.scheme}' is not allowed; allowed schemes: "
            f"{', '.join(ALLOWED_SCHEMES)}"
        )

    if not parsed.hostname:
        raise ValueError(f"Invalid host/URL; no host in parse result: {url}")

    # IDNA encode to prevent Unicode domain attacks
    try:
        ascii_hostname = idna.encode(parsed.hostname).decode("ascii")
    except idna.IDNAError:
        raise ValueError(f"Hostname '{parsed.hostname}' has unsupported characters")

    if not HOSTNAME_REGEX.match(ascii_hostname):
        raise ValueError(f"Hostname '{parsed.hostname}' has unsupported characters")

    # Re-create parsed URL object with IDNA-encoded hostname
    parsed = URL(
        parsed.scheme,
        (ascii_hostname if parsed.port is None else f"{ascii_hostname}:{parsed.port}"),
        quote(parsed.path, safe="/%:@"),
        parsed.params,
        parsed.query,
        parsed.fragment,
    )

    is_trusted = trusted_hostnames and any(
        matches_allowed_host(parsed, allowed)
        for allowed in (
            # Normalize + parse allowlist entries the same way for consistent matching
            parse_url(w)
            for w in trusted_hostnames
        )
    )

    if is_trusted:
        return parsed, True, []

    # If not allowlisted, go ahead with host resolution and IP target check
    return parsed, False, await resolve_and_check_blocked(ascii_hostname)


def matches_allowed_host(url: URL, allowed: URL) -> bool:
    if url.hostname != allowed.hostname:
        return False

    # Allow any port if not explicitly specified in the allowlist
    if allowed.port is None:
        return True

    return url.port == allowed.port


async def resolve_and_check_blocked(hostname: str) -> list[str]:
    """
    Resolves hostname to IPs and raises ValueError if any resolve to
    a blocked network. Returns the list of resolved IP addresses.
    """
    ip_addresses = await _resolve_host(hostname)
    for ip_str in ip_addresses:
        if _is_ip_blocked(ip_str):
            raise ValueError(
                f"Access to blocked or private IP address {ip_str} "
                f"for hostname {hostname} is not allowed."
            )
    return ip_addresses


def parse_url(url: str) -> URL:
    """Canonicalizes and parses a URL string."""
    url = url.strip("/ ").replace("\\", "/")

    # Ensure scheme is present for proper parsing.
    # Avoid regex to sidestep CodeQL py/polynomial-redos on user-controlled
    # input.  We only need to detect "scheme://"; urlparse() and the
    # ALLOWED_SCHEMES check downstream handle full validation.
    if "://" not in url:
        url = f"http://{url}"

    return urlparse(url)


def pin_url(url: URL, ip_addresses: Optional[list[str]] = None) -> URL:
    """
    Pins a URL to a specific IP address to prevent DNS rebinding attacks.

    Args:
        url: The original URL
        ip_addresses: List of IP addresses corresponding to the URL's host

    Returns:
        pinned_url: The URL with hostname replaced with IP address
    """
    if not url.hostname:
        raise ValueError(f"URL has no hostname: {url}")

    if not ip_addresses:
        # Resolve all IP addresses for the hostname
        # (This call is blocking; ensure to call async _resolve_host before if possible)
        ip_addresses = []
        # You may choose to raise or call synchronous resolve here; for simplicity, leave empty.

    # Pin to the first valid IP (for SSRF defense)
    pinned_ip = ip_addresses[0]

    # If it's IPv6, bracket it
    if ":" in pinned_ip:
        pinned_netloc = f"[{pinned_ip}]"
    else:
        pinned_netloc = pinned_ip

    if url.port:
        pinned_netloc += f":{url.port}"

    return URL(
        url.scheme,
        pinned_netloc,
        url.path,
        url.params,
        url.query,
        url.fragment,
    )


ClientResponse = aiohttp.ClientResponse
ClientResponseError = aiohttp.ClientResponseError


class Response:
    """
    Buffered wrapper around aiohttp.ClientResponse that does *not* require
    callers to manage connection or session lifetimes.
    """

    def __init__(
        self,
        *,
        response: ClientResponse,
        url: str,
        body: bytes,
    ):
        self.status: int = response.status
        self.headers = response.headers
        self.reason: str | None = response.reason
        self.request_info = response.request_info
        self.url: str = url
        self.content: bytes = body  # raw bytes

    def json(self, encoding: str | None = None, **kwargs) -> dict:
        """
        Parse the body as JSON and return the resulting Python object.
        """
        return loads(
            self.content.decode(encoding or "utf-8", errors="replace"), **kwargs
        )

    def text(self, encoding: str | None = None) -> str:
        """
        Decode the body to a string.  Encoding is guessed from the
        Content-Type header if not supplied.
        """
        if encoding is None:
            # Try to extract charset from headers; fall back to UTF-8
            ctype = self.headers.get("content-type", "")
            match = re.search(r"charset=([^\s;]+)", ctype, flags=re.I)
            encoding = match.group(1) if match else None
        return self.content.decode(encoding or "utf-8", errors="replace")

    @property
    def ok(self) -> bool:
        return 200 <= self.status < 300


def _return_last_result(retry_state: RetryCallState) -> "Response":
    """
    Ensure the final attempt's response is returned when retrying stops.
    """
    if retry_state.outcome is None:
        raise RuntimeError("Retry state is missing an outcome.")

    exception = retry_state.outcome.exception()
    if exception is not None:
        raise exception

    return retry_state.outcome.result()


class Requests:
    """
    A wrapper around an aiohttp ClientSession that validates URLs before
    making requests, preventing SSRF by blocking private networks and
    other disallowed address spaces.
    """

    def __init__(
        self,
        trusted_origins: list[str] | None = None,
        raise_for_status: bool = True,
        extra_url_validator: Callable[[URL], URL] | None = None,
        extra_headers: dict[str, str] | None = None,
        retry_max_wait: float = 300.0,
        retry_max_attempts: int | None = None,
    ):
        self.trusted_origins = []
        for url in trusted_origins or []:
            hostname = parse_url(url).netloc  # {host}[:{port}]
            if not hostname:
                raise ValueError(f"Invalid URL: Unable to determine hostname of {url}")
            self.trusted_origins.append(hostname)

        self.raise_for_status = raise_for_status
        self.extra_url_validator = extra_url_validator
        self.extra_headers = extra_headers
        self.retry_max_wait = retry_max_wait
        if retry_max_attempts is not None and retry_max_attempts < 1:
            raise ValueError("retry_max_attempts must be None or >= 1")
        self.retry_max_attempts = retry_max_attempts

    async def request(
        self,
        method: str,
        url: str,
        *,
        headers: Optional[dict] = None,
        files: list[tuple[str, tuple[str, BytesIO, str]]] | None = None,
        data: Any | None = None,
        json: Any | None = None,
        allow_redirects: bool = True,
        max_redirects: int = 10,
        **kwargs,
    ) -> Response:
        retry_kwargs: dict[str, Any] = {
            "wait": wait_exponential_jitter(max=self.retry_max_wait),
            "retry": retry_if_result(lambda r: r.status in THROTTLE_RETRY_STATUS_CODES),
            "reraise": True,
        }

        if self.retry_max_attempts is not None:
            retry_kwargs["stop"] = stop_after_attempt(self.retry_max_attempts)
            retry_kwargs["retry_error_callback"] = _return_last_result

        @retry(**retry_kwargs)
        async def _make_request() -> Response:
            return await self._request(
                method=method,
                url=url,
                headers=headers,
                files=files,
                data=data,
                json=json,
                allow_redirects=allow_redirects,
                max_redirects=max_redirects,
                **kwargs,
            )

        return await _make_request()

    async def _request(
        self,
        method: str,
        url: str,
        *,
        headers: Optional[dict] = None,
        files: list[tuple[str, tuple[str, BytesIO, str]]] | None = None,
        data: Any | None = None,
        json: Any | None = None,
        allow_redirects: bool = True,
        max_redirects: int = 10,
        **kwargs,
    ) -> Response:
        # Convert auth tuple to aiohttp.BasicAuth if necessary
        if "auth" in kwargs and isinstance(kwargs["auth"], tuple):
            kwargs["auth"] = aiohttp.BasicAuth(*kwargs["auth"])

        if files is not None:
            if json is not None:
                raise ValueError(
                    "Cannot mix file uploads with JSON body; "
                    "use 'data' for extra form fields instead."
                )

            form = FormData(quote_fields=False)
            # add normal form fields first
            if isinstance(data, dict):
                for k, v in data.items():
                    form.add_field(k, str(v))
            elif data is not None:
                raise ValueError(
                    "When uploading files, 'data' must be a dict of form fields."
                )

            # add the file parts
            for field_name, (filename, fh, content_type) in files:
                form.add_field(
                    name=field_name,
                    value=fh,
                    filename=filename,
                    content_type=content_type or "application/octet-stream",
                )

            data = form

        # Validate URL and get trust status
        parsed_url, is_trusted, ip_addresses = await validate_url_host(
            url, self.trusted_origins
        )

        # Apply any extra user-defined validation/transformation
        if self.extra_url_validator is not None:
            parsed_url = self.extra_url_validator(parsed_url)

        # Pin the URL if untrusted
        hostname = parsed_url.hostname
        if hostname is None:
            raise ValueError(f"Invalid URL: Unable to determine hostname of {url}")

        original_url = parsed_url.geturl()
        connector: Optional[aiohttp.TCPConnector] = None
        if not is_trusted:
            # Replace hostname with IP for connection but preserve SNI via resolver
            resolver = HostResolver(ssl_hostname=hostname, ip_addresses=ip_addresses)
            ssl_context = ssl.create_default_context()
            connector = aiohttp.TCPConnector(resolver=resolver, ssl=ssl_context)
        session_kwargs: dict = {}
        if connector:
            session_kwargs["connector"] = connector

        # Merge any extra headers
        req_headers = dict(headers) if headers else {}
        if self.extra_headers is not None:
            req_headers.update(self.extra_headers)

        # Set default User-Agent if not provided
        if "User-Agent" not in req_headers and "user-agent" not in req_headers:
            req_headers["User-Agent"] = DEFAULT_USER_AGENT

        # Override Host header if using IP connection
        if connector:
            req_headers["Host"] = hostname

        # Override data if files are provided
        # Set max_field_size to handle servers with large headers (e.g., long CSP headers)
        # Default is 8190 bytes, we increase to 16KB to accommodate legitimate large headers
        session_kwargs["max_field_size"] = 16384

        async with aiohttp.ClientSession(**session_kwargs) as session:
            # Perform the request with redirects disabled for manual handling
            async with session.request(
                method,
                parsed_url.geturl(),
                headers=req_headers,
                allow_redirects=False,
                data=data,
                json=json,
                **kwargs,
            ) as response:
                if self.raise_for_status:
                    try:
                        response.raise_for_status()
                    except ClientResponseError as e:
                        body = await response.read()
                        error_message = f"HTTP {response.status} Error: {response.reason}, Body: {body.decode(errors='replace')}"

                        # Raise specific exceptions based on status code range
                        if 400 <= response.status <= 499:
                            raise HTTPClientError(error_message, response.status) from e
                        elif 500 <= response.status <= 599:
                            raise HTTPServerError(error_message, response.status) from e
                        else:
                            # Generic fallback for other HTTP errors
                            raise Exception(error_message) from e

                # If allowed and a redirect is received, follow the redirect manually
                if allow_redirects and response.status in (301, 302, 303, 307, 308):
                    if max_redirects <= 0:
                        raise Exception("Too many redirects.")

                    location = response.headers.get("Location")
                    if not location:
                        return Response(
                            response=response,
                            url=original_url,
                            body=await response.read(),
                        )

                    # The base URL is the pinned_url we just used
                    # so that relative redirects resolve correctly.
                    redirect_url = urlparse(urljoin(parsed_url.geturl(), location))
                    # Carry forward the same headers but update Host
                    new_headers = _remove_insecure_headers(
                        req_headers, parsed_url, redirect_url
                    )

                    return await self.request(
                        method,
                        redirect_url.geturl(),
                        headers=new_headers,
                        allow_redirects=allow_redirects,
                        max_redirects=max_redirects - 1,
                        files=files,
                        data=data,
                        json=json,
                        **kwargs,
                    )

                # Reset response URL to original host for clarity
                if parsed_url.hostname != hostname:
                    try:
                        response.url = original_url  # type: ignore
                    except Exception:
                        pass

                return Response(
                    response=response,
                    url=original_url,
                    body=await response.read(),
                )

    async def get(self, url: str, *args, **kwargs) -> Response:
        return await self.request("GET", url, *args, **kwargs)

    async def post(self, url: str, *args, **kwargs) -> Response:
        return await self.request("POST", url, *args, **kwargs)

    async def put(self, url: str, *args, **kwargs) -> Response:
        return await self.request("PUT", url, *args, **kwargs)

    async def delete(self, url: str, *args, **kwargs) -> Response:
        return await self.request("DELETE", url, *args, **kwargs)

    async def head(self, url: str, *args, **kwargs) -> Response:
        return await self.request("HEAD", url, *args, **kwargs)

    async def options(self, url: str, *args, **kwargs) -> Response:
        return await self.request("OPTIONS", url, *args, **kwargs)

    async def patch(self, url: str, *args, **kwargs) -> Response:
        return await self.request("PATCH", url, *args, **kwargs)
