import json
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.gzip import GZipMiddleware

from assembly_viz.config import ROOT, PROCESSED_DIR

_WEB_DIR  = ROOT / "src" / "web"
_DATA_DIR = ROOT / "data"
_NET_FILE = PROCESSED_DIR / "network_index.json"

app = FastAPI()
_network_index: dict = {}


class CacheMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        response = await call_next(request)
        path = request.url.path
        qs   = request.url.query
        if path == "/" or path.endswith(".html"):
            # HTML: always revalidate so users get the latest page
            response.headers["Cache-Control"] = "no-cache"
        elif path.endswith(".js") and qs:
            # Versioned JS (e.g. app.js?v=18): safe to cache for 1 day
            response.headers["Cache-Control"] = "public, max-age=86400, immutable"
        elif path.startswith("/data/processed/") and path.endswith(".json"):
            # Large data files: cache for 1 hour (data never changes at runtime)
            response.headers["Cache-Control"] = "public, max-age=3600"
        return response


app.add_middleware(GZipMiddleware, minimum_size=1000)
app.add_middleware(CacheMiddleware)


@app.on_event("startup")
async def _load_network_index():
    global _network_index
    if _NET_FILE.exists():
        with open(_NET_FILE) as f:
            _network_index = json.load(f)
        print(f"[server] network index: {len(_network_index)} molecules")


@app.get("/api/network")
async def get_network(request: Request):
    smiles = request.query_params.get("smiles", "")
    data   = _network_index.get(smiles)
    if data is None:
        raise HTTPException(status_code=404, detail="not found")
    return data


app.mount("/data", StaticFiles(directory=str(_DATA_DIR)), name="data")
app.mount("/", StaticFiles(directory=str(_WEB_DIR), html=True), name="web")
