main.py

import os, httpx
from fastapi import FastAPI, Request, HTTPException, UploadFile, File, Form
from fastapi.responses import JSONResponse, Response, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from slowapi import Limiter
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware

# --- Secrets set in the PUBLIC Space ---
UPSTREAM_URL = os.environ.get("hf_url", "").rstrip("/")             # url to access Private Space
HF_TOKEN     = os.environ.get("hf_token")                           # HF access token to access Private Space
DEMO_FORWARD_URL   = os.getenv("DEMO_FORWARD_URL", "").rstrip("/")  # url to acces demo space 
DATFID_DEMO_TOKEN  = os.getenv("DATFID_DEMO_TOKEN", "")             # token to access demo space

if not HF_TOKEN:
    raise RuntimeError("Missing secret 'hf_token' in public Space.")
if not UPSTREAM_URL.startswith("https://"):
    raise RuntimeError("Missing/invalid secret 'hf_url' in public Space.")

app = FastAPI(title="DATFID Public Proxy", docs_url="/docs", redoc_url=None)

# Prefer X-Forwarded-For (HF sits behind a proxy)
def client_ip(request: Request):
    xff = request.headers.get("x-forwarded-for")
    return xff.split(",")[0].strip() if xff else (request.client.host or "0.0.0.0")

limiter = Limiter(key_func=client_ip)  # or get_remote_address
app.state.limiter = limiter
app.add_middleware(SlowAPIMiddleware)

@app.exception_handler(RateLimitExceeded)
async def ratelimit_handler(request: Request, exc: RateLimitExceeded):
    return JSONResponse(status_code=429, content={"detail": "Too many requests, slow down."})

SDK_MAX_BODY_BYTES = int(os.getenv("SDK_MAX_BODY_BYTES", "25000000"))  # 25MB default for demo
SDK_MAX_BODY_BYTES_extended = int(os.getenv("SDK_MAX_BODY_BYTES_extended", "125000000")) # 125MB default for prod 

# Early global body-size guard (runs before routes)
@app.middleware("http")
async def limit_body_size(request: Request, call_next):
    cl = request.headers.get("content-length")

    # demo routes use the smaller cap, others use extended
    path = request.url.path or ""
    cap = SDK_MAX_BODY_BYTES if "-demo" in path else SDK_MAX_BODY_BYTES_extended

    if cl and int(cl) > cap:
        return JSONResponse(
            {"detail": f"Payload too large (> {cap} bytes)"},
            status_code=413,
        )
    return await call_next(request)

# CORS for browser
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://datfid.com",
        "https://www.datfid.com"
    ],
    # Optional: allow Vercel preview domains
    # allow_origin_regex=r"^https:\/\/.*\.vercel\.app$",
    allow_methods=["POST", "OPTIONS"],
    allow_headers=["Content-Type"],   # no Authorization header needed from browser
    allow_credentials=False,
    max_age=86400,
)

# to ensure we don’t leak hop-by-hop headers
def _filter_resp_headers(h):
    # pass through useful headers but strip hop-by-hop
    allowed = {"content-type", "content-disposition", "content-length"}
    return {k: v for k, v in h.items() if k.lower() in allowed}

def _extract_user_token(req: Request) -> str | None:
    """
    Read user's DATFID token from Authorization: Bearer <dt+...>.
    We do NOT verify it here; the private Space does that.
    """
    auth = req.headers.get("authorization", "")
    if not auth.lower().startswith("bearer "):
        return None
    return auth.split(" ", 1)[1].strip()

async def _forward(path: str, method: str = "GET", json_body=None, user_token: str | None = None):
    """
    Forward request to the PRIVATE Space:
      - 'Authorization: Bearer <HF_TOKEN>' to pass HF private gate
      - 'X-API-Key: <dt+...>' so your private app can validate the user token
    """
    url = f"{UPSTREAM_URL}{path}"
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Accept": "application/json",
    }
    if user_token:
        headers["X-API-Key"] = user_token

    timeout = httpx.Timeout(600.0)
    async with httpx.AsyncClient(timeout=timeout) as client:
        r = await client.request(method, url, headers=headers, json=json_body)

    ct = r.headers.get("content-type", "")

    if "application/json" in ct:
        try:
            return JSONResponse(status_code=r.status_code, content=r.json())
        except Exception:
            return JSONResponse(status_code=r.status_code, content={"error": r.text[:500]})
    # Fallback: return short text envelope if non-JSON
    return JSONResponse(status_code=r.status_code, content={"text": r.text[:1000]})

async def _forward_stream(path: str, files=None, data=None, user_token: str | None = None, method: str = "POST"):
    url = f"{UPSTREAM_URL}{path}"
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Accept": "*/*",
    }
    if user_token:
        headers["X-API-Key"] = user_token

    timeout = httpx.Timeout(600.0)
    async with httpx.AsyncClient(timeout=timeout) as client:
        async with client.stream(method, url, headers=headers, files=files, data=data) as resp:
            if resp.status_code >= 400:
                text = await resp.aread()
                return Response(content=text, status_code=resp.status_code, media_type=resp.headers.get("content-type","text/plain"))
            return StreamingResponse(resp.aiter_raw(), status_code=resp.status_code, headers=_filter_resp_headers(resp.headers))

# for demo
async def _forward_demo_stream(path: str, *, files: dict | None, data: dict | None, method: str = "POST"):
    if not DEMO_FORWARD_URL or not HF_TOKEN or not DATFID_DEMO_TOKEN:
        raise HTTPException(status_code=500, detail="Demo not configured.")

    url = DEMO_FORWARD_URL + path
    headers = {
        # HF private-space gate:
        "Authorization": f"Bearer {HF_TOKEN}",
        # App-level demo token (checked by the private API):
        "X-DATFID-Token": DATFID_DEMO_TOKEN,
    }

    timeout = httpx.Timeout(120.0)
    async with httpx.AsyncClient(timeout=timeout) as client:
        # stream response back to the caller (so large CSVs don’t load fully into memory)
        async with client.stream(method, url, headers=headers, files=files, data=data) as resp:
            if resp.status_code >= 400:
                text = await resp.aread()
                raise HTTPException(status_code=resp.status_code, detail=text.decode(errors="ignore"))
            return StreamingResponse(resp.aiter_raw(), status_code=resp.status_code, headers=_filter_resp_headers(resp.headers))

@app.get("/")
async def root(req: Request):
    # Forward to private root (private gate still needs HF token)
    user_token = _extract_user_token(req)  # optional here
    return await _forward("/", "GET", user_token=user_token)

@app.get("/secure-ping/")
async def secure_ping(req: Request):
    # Require user's DATFID token in Authorization: Bearer <dt+...>
    user_token = _extract_user_token(req)
    if not user_token:
        raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
    return await _forward("/secure-ping/", "GET", user_token=user_token)

@app.post("/modelfit/")
async def modelfit(req: Request):
    user_token = _extract_user_token(req)
    if not user_token:
        raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
    body = await req.json()
    return await _forward("/modelfit/", "POST", json_body=body, user_token=user_token)

@app.post("/modelforecast/")
async def modelforecast(req: Request):
    user_token = _extract_user_token(req)
    if not user_token:
        raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
    body = await req.json()
    return await _forward("/modelforecast/", "POST", json_body=body, user_token=user_token)

@app.post("/modelfit-file/")
async def modelfit_file(
    req: Request,
    file: UploadFile = File(...),
    id_col: str = Form(...),
    time_col: str = Form(...),
    y: str = Form(...),
    # optional knobs
    lag_y: str = Form(""),
    lagged_features: str = Form(""),          # JSON string or empty
    current_features: str = Form(""),         # "all" | JSON string | ""
    filter_by_significance: str = Form("false"),
    meanvar_test: str = Form("false"),
):
    user_token = _extract_user_token(req)
    if not user_token:
        raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")

    raw = await file.read()
    if len(raw) > SDK_MAX_BODY_BYTES_extended:
        raise HTTPException(status_code=413, detail="Payload too large.")

    files = {
        "file": (file.filename, raw, file.content_type or "application/octet-stream"),
    }
    data = {
        "id_col": id_col,
        "time_col": time_col,
        "y": y,
        "lag_y": lag_y,
        "lagged_features": lagged_features,
        "current_features": current_features,
        "filter_by_significance": filter_by_significance,
        "meanvar_test": meanvar_test,
    }
    return await _forward_stream("/modelfit-file/", files=files, data=data, user_token=user_token, method="POST")

@app.post("/modelforecast-file/")
async def modelforecast_file(
    req: Request,
    df_forecast: UploadFile = File(...),
):
    user_token = _extract_user_token(req)
    if not user_token:
        raise HTTPException(status_code=401, detail="Missing Authorization Bearer token (dt+...)")
    
    raw = await df_forecast.read()
    if len(raw) > SDK_MAX_BODY_BYTES_extended:
        raise HTTPException(status_code=413, detail="Payload too large.")

    files = {
        "df_forecast": (df_forecast.filename, raw, df_forecast.content_type or "application/octet-stream"),
    }
    return await _forward_stream("/modelforecast-file/", files=files, data=None, user_token=user_token, method="POST")

@limiter.limit("10/10minute")  # 10 calls per 10 minutes per IP
@app.post("/modelfit-file-demo/")
async def modelfit_file_demo(
    request: Request,
    file: UploadFile = File(...),
    id_col: str = Form(...),
    time_col: str = Form(...),
    y: str = Form(...),
    # optional knobs
    lag_y: str = Form(""),
    lagged_features: str = Form(""),
    current_features: str = Form(""),
    filter_by_significance: str = Form("false"),
    meanvar_test: str = Form("false"),
):
    
    # read once, size-guard it, then forward
    raw = await file.read()
    if len(raw) > SDK_MAX_BODY_BYTES:
        raise HTTPException(status_code=413, detail="Payload too large.")
    
    files = {
        "file": (file.filename, raw, file.content_type or "application/octet-stream"),
    }
    data = {
        "id_col": id_col,
        "time_col": time_col,
        "y": y,
        "lag_y": lag_y,
        "lagged_features": lagged_features,
        "current_features": current_features,
        "filter_by_significance": filter_by_significance,
        "meanvar_test": meanvar_test,
    }
    # the path goes to the private demo route
    return await _forward_demo_stream("/modelfit-file-demo/", files=files, data=data, method="POST")

@limiter.limit("10/10minute")
@app.post("/modelforecast-file-demo/")
async def modelforecast_file_demo(
    request: Request,
    df_forecast: UploadFile = File(...),
):
    
    raw = await df_forecast.read()
    if len(raw) > SDK_MAX_BODY_BYTES:
        raise HTTPException(status_code=413, detail="Payload too large.")
    
    files = {
        "df_forecast": (df_forecast.filename, raw, df_forecast.content_type or "application/octet-stream"),
    }
    return await _forward_demo_stream("/modelforecast-file-demo/", files=files, data=None, method="POST")

@app.get("/health-demo-proxy")
async def health_demo_proxy():
    # convenience endpoint to test private-space + demo token hop
    try:
        return await _forward_demo_stream("/health-demo", files=None, data=None, method="GET")
    except HTTPException as e:
        # bubble up errors so you can diagnose missing tokens, wrong URL, etc.
        raise e