server/main.py

import os
import httpx
from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional
from starlette.background import BackgroundTask
from starlette.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

app = FastAPI()

# --- Configuration ---
# The URL of your PRIVATE calculation server.
PRIVATE_SERVER_URL = os.environ.get("PRIVATE_SERVER_URL")
# The broker will read your HF_TOKEN from the Space's secrets.
HF_TOKEN = os.environ.get("HF_TOKEN")
GOOGLE_MAPS_API_KEY = os.environ.get("VITE_GOOGLE_MAPS_API_KEY")

if not PRIVATE_SERVER_URL:
    print("WARNING: The `PRIVATE_SERVER_URL` environment variable is not set. The broker will not be able to forward requests.")

if not HF_TOKEN:
    print("WARNING: The `HF_TOKEN` environment variable is not set. Requests to the private server will not be authenticated.")

if not GOOGLE_MAPS_API_KEY:
    print("WARNING: The `VITE_GOOGLE_MAPS_API_KEY` environment variable is not set. The map on the visit log page will not load.")

# Read timeout from env var, default to 5 minutes (300 seconds) to accommodate slow model conversions.
BROKER_TIMEOUT = float(os.environ.get("BROKER_TIMEOUT", 300.0))
# Use a persistent client for performance (connection pooling).
client = httpx.AsyncClient(base_url=PRIVATE_SERVER_URL, timeout=BROKER_TIMEOUT)


async def _reverse_proxy(request: Request):
    """
    This function acts as a reverse proxy. It captures incoming requests,
    forwards them to the private backend server, and streams the response back.
    """
    print(f"Broker: Received request: {request.method} {request.url.path}")
    if HF_TOKEN:
        token_preview = HF_TOKEN[:4] + "..." + HF_TOKEN[-4:]
        print(f"Broker: Using HF_TOKEN: {token_preview}")
    else:
        print("Broker: WARNING - HF_TOKEN is not set. Request will be unauthenticated.")
    print(f"Broker: Target private server URL: {PRIVATE_SERVER_URL}")

    if not PRIVATE_SERVER_URL:
        return Response(status_code=503, content="Service Unavailable: Broker is not configured. `PRIVATE_SERVER_URL` is missing.")

    forward_headers = httpx.Headers(request.headers)
    forward_headers.pop("host", None)

    if HF_TOKEN:
        forward_headers["Authorization"] = f"Bearer {HF_TOKEN}"

    url = httpx.URL(path=request.url.path, query=request.url.query.encode("utf-8"))
    
    backend_request = client.build_request(method=request.method, url=url, headers=forward_headers, content=request.stream())
    print(f"Broker: Forwarding request to: {backend_request.method} {backend_request.url}")

    try:
        backend_response = await client.send(backend_request, stream=True)
        print(f"Broker: Received response from backend with status: {backend_response.status_code}")
        return StreamingResponse(backend_response.aiter_raw(), status_code=backend_response.status_code, headers=backend_response.headers, background=BackgroundTask(backend_response.aclose))
    except httpx.ConnectError as e:
        error_message = f"Connection to private server failed: {e}"
        print(f"Broker: CRITICAL - {error_message}")
        return Response(status_code=502, content=f"Bad Gateway: The broker could not connect to the private server. Details: {error_message}")

# --- Client Configuration Endpoint ---
# This endpoint must be defined *before* the reverse proxy routes.
class ClientConfig(BaseModel):
    googleMapsApiKey: Optional[str]

@app.get("/api/config", response_model=ClientConfig)
async def get_client_config():
    """Provides public-safe configuration variables to the client at runtime."""
    return ClientConfig(googleMapsApiKey=GOOGLE_MAPS_API_KEY)


# Add routes to capture all API and data requests and forward them.
app.add_route("/api/{path:path}", _reverse_proxy, ["GET", "POST", "PUT", "DELETE"])
app.add_route("/data/{path:path}", _reverse_proxy, ["GET"])

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- Static File Serving for Production ---
# This must come *after* the API routes to avoid conflicts.

# This mounts the 'static' directory (containing the built React app) at the root.
# The `html=True` argument enables SPA mode, serving `index.html` for any path
# that doesn't match a file. This is the recommended and most robust way to
# serve a Single-Page Application with client-side routing.
app.mount("/", StaticFiles(directory="static", html=True), name="static_root")