import httpx
from fastapi import FastAPI, Request, HTTPException
from starlette.responses import StreamingResponse, JSONResponse
from starlette.background import BackgroundTask
import os
import random
import logging
import time
import uvicorn
from contextlib import asynccontextmanager

# --- Production-Ready Configuration ---
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
logging.basicConfig(
    level=LOG_LEVEL,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# --- Simplified Target Configuration ---
# The proxy will now forward the request path exactly as it receives it.
# We only need to define the target host.
TARGET_HOST = os.getenv("TARGET_HOST", "https://api.gmi-serving.com")
logger.info(f"Proxying all request paths to host: {TARGET_HOST}")

# --- Retry Logic Configuration ---
MAX_RETRIES = int(os.getenv("MAX_RETRIES", "15"))
DEFAULT_RETRY_CODES = "429,500,502,503,504"
RETRY_CODES_STR = os.getenv("RETRY_CODES", DEFAULT_RETRY_CODES)
try:
    RETRY_STATUS_CODES = {int(code.strip()) for code in RETRY_CODES_STR.split(',')}
    logger.info(f"Will retry on the following status codes: {RETRY_STATUS_CODES}")
except ValueError:
    logger.error(f"Invalid RETRY_CODES format: '{RETRY_CODES_STR}'. Falling back to default: {DEFAULT_RETRY_CODES}")
    RETRY_STATUS_CODES = {int(code.strip()) for code in DEFAULT_RETRY_CODES.split(',')}


# --- Helper Function ---
def generate_random_ip():
    """Generates a random, valid-looking IPv4 address for X-Forwarded-For header."""
    return ".".join(str(random.randint(1, 254)) for _ in range(4))


# --- HTTPX Client Lifecycle Management ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manages the lifecycle of the HTTPX client."""
    async with httpx.AsyncClient(base_url=TARGET_HOST, timeout=None) as client:
        logger.info(f"HTTPX client created for target: {TARGET_HOST}")
        app.state.http_client = client
        yield
    logger.info("HTTPX client closed gracefully.")


# Initialize the FastAPI app with the lifespan manager and disabled docs
app = FastAPI(docs_url=None, redoc_url=None, lifespan=lifespan)


# --- API Endpoints ---

@app.get("/", include_in_schema=False)
async def health_check():
    """Provides a basic health check endpoint."""
    return JSONResponse({
        "status": "ok",
        "target_host": TARGET_HOST,
    })


@app.api_route("/{full_path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
async def reverse_proxy_handler(request: Request):
    """
    A catch-all reverse proxy that forwards requests to the target host.
    It forwards the path and query parameters exactly as received.
    """
    start_time = time.monotonic()
    client: httpx.AsyncClient = request.app.state.http_client

    # --- THE CORE FIX: Forward the path as-is, without adding any prefix ---
    url = httpx.URL(path=request.url.path, query=request.url.query.encode("utf-8"))

    request_headers = dict(request.headers)
    request_headers.pop("host", None)
    
    random_ip = generate_random_ip()
    request_headers.update({
        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36",
        "x-forwarded-for": random_ip,
        "x-real-ip": random_ip,
    })

    request_body = await request.body()
    
    last_exception = None
    for attempt in range(MAX_RETRIES):
        try:
            req = client.build_request(
                method=request.method,
                url=url,
                headers=request_headers,
                content=request_body,
            )
            
            # Use debug level for this log as it can be very noisy
            logger.debug(f"Attempt {attempt + 1}/{MAX_RETRIES} -> {req.method} {req.url}")
            
            resp = await client.send(req, stream=True)

            if resp.status_code not in RETRY_STATUS_CODES or attempt == MAX_RETRIES - 1:
                duration_ms = (time.monotonic() - start_time) * 1000
                log_func = logger.info if resp.is_success else logger.warning
                log_func(
                    f"Request finished: {request.method} {request.url.path} -> {resp.status_code} "
                    f"[{resp.reason_phrase}] latency={duration_ms:.2f}ms"
                )
                
                return StreamingResponse(
                    resp.aiter_raw(),
                    status_code=resp.status_code,
                    headers=resp.headers,
                    background=BackgroundTask(resp.aclose),
                )

            logger.warning(
                f"Attempt {attempt + 1}/{MAX_RETRIES} for {url.path} failed with status {resp.status_code}. Retrying..."
            )
            await resp.aclose()
            time.sleep(1)

        except (httpx.ConnectError, httpx.ReadTimeout, httpx.ConnectTimeout) as e:
            last_exception = e
            logger.warning(f"Attempt {attempt + 1}/{MAX_RETRIES} for {url.path} failed with connection error: {e}")
            if attempt < MAX_RETRIES - 1:
                time.sleep(1)
            continue

    duration_ms = (time.monotonic() - start_time) * 1000
    logger.critical(
        f"Request failed after {MAX_RETRIES} attempts. Cannot connect to target. "
        f"path={request.url.path} latency={duration_ms:.2f}ms"
    )
    
    raise HTTPException(
        status_code=502,
        detail=f"Bad Gateway: Cannot connect to target service at {TARGET_HOST} after {MAX_RETRIES} attempts. Last error: {last_exception}"
    )


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)