# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncIterator, Callable
from urllib.parse import quote

import httpx
import websockets
from fastapi import Request, WebSocket
from starlette.responses import Response, StreamingResponse

logger = logging.getLogger(__name__)

HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
    "host",
    "content-length",
    "content-encoding",
}

RESPONSE_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
    "date",
    "server",
}

WS_HANDSHAKE_HEADERS = {
    "sec-websocket-key",
    "sec-websocket-version",
    "sec-websocket-extensions",
    "sec-websocket-protocol",
}

METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"]

SCHEME_PROTO = {"http": "http", "https": "https", "ws": "http", "wss": "https"}

DEFAULT_TIMEOUT = 300.0
OPEN_TIMEOUT = 10.0
MAX_CONNECTIONS = 200
MAX_KEEPALIVE_CONNECTIONS = 50

_client: httpx.AsyncClient | None = None


def client() -> httpx.AsyncClient:
    global _client
    if _client is None or _client.is_closed:
        _client = httpx.AsyncClient(
            follow_redirects=False,
            timeout=DEFAULT_TIMEOUT,
            limits=httpx.Limits(
                max_connections=MAX_CONNECTIONS,
                max_keepalive_connections=MAX_KEEPALIVE_CONNECTIONS,
            ),
        )
    return _client


async def close_client() -> None:
    global _client
    if _client is not None and not _client.is_closed:
        await _client.aclose()
    _client = None


def public_host(connection) -> str:
    return connection.headers.get("host", "") or connection.url.hostname or ""


def raw_query(connection) -> str:
    return connection.scope.get("query_string", b"").decode("latin-1")


def base_headers(connection, prefix: str, drop: set[str]) -> dict:
    headers = {
        name: value
        for name, value in connection.headers.items()
        if name.lower() not in HOP_HEADERS and name.lower() not in drop
    }
    if prefix:
        headers["X-Forwarded-Prefix"] = prefix
        headers["X-Script-Name"] = prefix
    origin_host = public_host(connection)
    if origin_host:
        headers["X-Forwarded-Host"] = origin_host
    headers["X-Forwarded-Proto"] = connection.headers.get(
        "x-forwarded-proto", SCHEME_PROTO.get(connection.url.scheme, "http")
    )
    return headers


def forward_headers(request: Request, prefix: str = "") -> dict:
    headers = base_headers(request, prefix, set())
    origin_host = public_host(request)
    if origin_host:
        headers["Host"] = origin_host
    headers["Accept-Encoding"] = "identity"
    return headers


def ws_headers(websocket: WebSocket, prefix: str = "") -> dict:
    return base_headers(websocket, prefix, WS_HANDSHAKE_HEADERS)


def ws_subprotocols(websocket: WebSocket) -> list[str] | None:
    offered = [
        value.strip()
        for value in websocket.headers.get("sec-websocket-protocol", "").split(",")
        if value.strip()
    ]
    return offered or None


def upstream_url(scheme: str, authority: str, path: str, query: str) -> str:
    url = f"{scheme}://{authority}/{quote(path)}"
    if query:
        url = f"{url}?{query}"
    return url


def is_root_document(path: str) -> bool:
    return not path.strip("/")


def inject_base(body: bytes, prefix: str) -> bytes:
    lowered = body.lower()
    if b"<base" in lowered:
        return body
    tag = f'<base href="{prefix}/">'.encode()
    opening = lowered.find(b"<head")
    if opening == -1:
        opening = lowered.find(b"<html")
    anchor = lowered.find(b">", opening) if opening != -1 else -1
    if anchor == -1:
        return tag + body
    return body[: anchor + 1] + tag + body[anchor + 1 :]


def rewrite_location(value: str, prefix: str) -> str:
    if not prefix:
        return value
    if value.startswith("/") and not value.startswith("//"):
        return prefix + value
    return value


def response_headers(upstream: httpx.Response, prefix: str) -> dict:
    headers = {
        name: value
        for name, value in upstream.headers.items()
        if name.lower() not in RESPONSE_HOP_HEADERS and name.lower() != "set-cookie"
    }
    if "location" in headers:
        headers["location"] = rewrite_location(headers["location"], prefix)
    return headers


def apply_cookies(response: Response, upstream: httpx.Response) -> Response:
    for cookie in upstream.headers.get_list("set-cookie"):
        response.headers.append("set-cookie", cookie)
    return response


async def stream_body(
    upstream: httpx.Response, on_complete: Callable[[int], None] | None
) -> AsyncIterator[bytes]:
    sent = 0
    try:
        async for chunk in upstream.aiter_raw():
            sent += len(chunk)
            yield chunk
    finally:
        await upstream.aclose()
        if on_complete is not None:
            on_complete(sent)


async def proxy_http(
    request: Request,
    host: str,
    port: int,
    path: str,
    *,
    prefix: str = "",
    timeout: float = DEFAULT_TIMEOUT,
    rewrite_html: bool = True,
    on_complete: Callable[[int], None] | None = None,
) -> Response:
    upstream_request = client().build_request(
        request.method,
        upstream_url("http", f"{host}:{port}", path, raw_query(request)),
        headers=forward_headers(request, prefix),
        content=await request.body(),
        timeout=timeout,
    )
    try:
        upstream = await client().send(upstream_request, stream=True)
    except httpx.HTTPError as error:
        return Response(f"upstream error: {error}", status_code=502)
    content_type = upstream.headers.get("content-type", "")
    headers = response_headers(upstream, prefix)
    if (
        prefix
        and rewrite_html
        and is_root_document(path)
        and "text/html" in content_type.lower()
    ):
        body = await upstream.aread()
        await upstream.aclose()
        content = inject_base(body, prefix)
        headers.pop("content-length", None)
        headers.pop("content-encoding", None)
        if on_complete is not None:
            on_complete(len(content))
        return apply_cookies(
            Response(
                content=content,
                status_code=upstream.status_code,
                headers=headers,
                media_type=content_type or None,
            ),
            upstream,
        )
    return apply_cookies(
        StreamingResponse(
            stream_body(upstream, on_complete),
            status_code=upstream.status_code,
            headers=headers,
            media_type=content_type or None,
        ),
        upstream,
    )


async def close_quietly(websocket: WebSocket) -> None:
    try:
        await websocket.close(code=1011)
    except Exception:
        pass


async def proxy_ws(
    websocket: WebSocket,
    host: str,
    port: int,
    path: str,
    *,
    prefix: str = "",
    timeout: float = OPEN_TIMEOUT,
) -> None:
    authority = public_host(websocket) or f"{host}:{port}"
    target = upstream_url("ws", authority, path, raw_query(websocket))
    try:
        upstream = await websockets.connect(
            target,
            host=host,
            port=port,
            open_timeout=timeout,
            max_size=None,
            additional_headers=ws_headers(websocket, prefix),
            subprotocols=ws_subprotocols(websocket),
        )
    except Exception as error:
        logger.debug("ws proxy to %s:%s failed: %s", host, port, error)
        await close_quietly(websocket)
        return
    try:
        await websocket.accept(subprotocol=upstream.subprotocol)
        await pump(websocket, upstream)
    finally:
        await upstream.close()


async def pump(client_ws: WebSocket, upstream) -> None:
    async def client_to_upstream():
        try:
            while True:
                message = await client_ws.receive()
                if message["type"] == "websocket.disconnect":
                    break
                if message.get("text") is not None:
                    await upstream.send(message["text"])
                elif message.get("bytes") is not None:
                    await upstream.send(message["bytes"])
        except Exception:
            pass
        finally:
            await upstream.close()

    async def upstream_to_client():
        try:
            async for message in upstream:
                if isinstance(message, (bytes, bytearray)):
                    await client_ws.send_bytes(bytes(message))
                else:
                    await client_ws.send_text(message)
        except Exception:
            pass
        finally:
            try:
                await client_ws.close()
            except Exception:
                pass

    await asyncio.gather(client_to_upstream(), upstream_to_client())
