import httpx
from fastapi import FastAPI, Request
from fastapi.responses import Response
from starlette.responses import JSONResponse

from client import AsyncCustomHost, NameSolver
from const import OFFICIAL_API

app = FastAPI()

# Configure upstream client
official_client = httpx.AsyncClient(
    verify=True,
    timeout=httpx.Timeout(30.0, connect=10.0),
    limits=httpx.Limits(max_keepalive_connections=50, max_connections=100),
    transport=AsyncCustomHost(NameSolver())
)

# Hop-by-hop headers we must strip
HOP_BY_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailer",
    "transfer-encoding",
    "upgrade",
    "content-length",  # let httpx calculate
    "host",  # must be rewritten
}


@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
async def catch_all(request: Request, path: str):
    # Copy request body
    body = await request.body()

    # Copy headers, filtering hop-by-hop
    req_headers = {
        k: v
        for k, v in request.headers.items()
        if k.lower() not in HOP_BY_HOP_HEADERS
    }

    # Rewrite Host header to match upstream
    req_headers["host"] = httpx.URL(OFFICIAL_API).host

    # Forward request upstream
    upstream_response = await official_client.request(
        method=request.method,
        url=f"{OFFICIAL_API}/{path}",
        headers=req_headers,
        content=body,
        params=request.query_params,
        cookies=request.cookies,
    )

    # Copy response headers, filtering hop-by-hop
    resp_headers = {
        k: v
        for k, v in upstream_response.headers.items()
        if k.lower() not in HOP_BY_HOP_HEADERS
    }

    content_type = upstream_response.headers.get("content-type", "")

    resp_headers.pop("content-encoding", None)  # remove, since body is decoded
    if content_type.startswith("application/json"):
        response = JSONResponse(
            content=upstream_response.json(),
            status_code=upstream_response.status_code,
            headers=resp_headers
        )
    else:
        response = Response(
            content=upstream_response.content,
            status_code=upstream_response.status_code,
            headers=resp_headers,
            media_type=content_type
        )

    return response