"""Trust X-Forwarded-For / X-Forwarded-Proto only from the proxy's resolved addresses."""

from __future__ import annotations

import asyncio
import ipaddress
import logging
import socket
from collections.abc import Callable, Iterable

from starlette.requests import Request
from starlette.types import ASGIApp, Receive, Scope, Send

log = logging.getLogger(__name__)

REFRESH_SECONDS = 60


def _ip(text: str) -> str | None:
    try:
        return str(ipaddress.ip_address(text.strip()))
    except ValueError:
        return None


def resolve_host(host: str) -> set[str]:
    """Every address the name resolves to (IPv4 and IPv6)."""
    infos = socket.getaddrinfo(host, None, proto=socket.IPPROTO_TCP)
    out = {ip for info in infos if (ip := _ip(str(info[4][0]).split("%")[0]))}
    return out


class TrustedProxies:
    """The current set of trusted peer addresses. A failed refresh keeps the last good set."""

    def __init__(self, host: str, resolver: Callable[[str], set[str]] = resolve_host):
        self.host = host.strip()
        self.addresses: frozenset[str] = frozenset()
        self._resolver = resolver

    def refresh(self) -> None:
        if not self.host:
            return
        try:
            found = self._resolver(self.host)
        except OSError as e:
            log.warning(
                "can't resolve TRUSTED_PROXY_HOST %s (%s); keeping %s",
                self.host,
                e,
                sorted(self.addresses) or "nothing",
            )
            return
        if found != self.addresses:
            log.info("trusted proxy %s is %s", self.host, sorted(found))
        self.addresses = frozenset(found)

    def trusts(self, peer: str | None) -> bool:
        return bool(peer) and _ip(peer or "") in self.addresses

    async def run(self) -> None:
        """Background loop: re-resolve every 60 s."""
        while True:
            await asyncio.sleep(REFRESH_SECONDS)
            await asyncio.to_thread(self.refresh)


def forwarded(
    peer: str | None, headers: Iterable[tuple[bytes, bytes]], trusted: TrustedProxies
) -> tuple[str | None, str | None]:
    """(client ip, scheme) to use; (peer, None) when the peer is not trusted or headers are unusable."""
    if not trusted.trusts(peer):
        return peer, None
    xff = proto = None
    for k, v in headers:
        name = k.lower()
        if name == b"x-forwarded-for" and xff is None:
            xff = v.decode("latin-1")
        elif name == b"x-forwarded-proto" and proto is None:
            proto = v.decode("latin-1").split(",")[0].strip().lower()
    client = peer
    if xff:
        first = _ip(xff.split(",")[0])
        if first:
            client = first
    scheme = proto if proto in ("http", "https") else None
    return client, scheme


class TrustedProxyMiddleware:
    def __init__(self, app: ASGIApp, trusted: TrustedProxies):
        self.app = app
        self.trusted = trusted

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] in ("http", "websocket"):
            client = scope.get("client")
            peer = client[0] if client else None
            ip, scheme = forwarded(peer, scope.get("headers", []), self.trusted)
            if ip != peer or scheme:
                scope = dict(scope)
                scope["client"] = (ip, client[1] if client else 0)
                if scheme:
                    scope["scheme"] = scheme if scope["type"] == "http" else ("wss" if scheme == "https" else "ws")
        await self.app(scope, receive, send)


def client_ip(request: Request) -> str | None:
    """The real client address (after trusted-proxy handling)."""
    return request.client.host if request.client else None