from fastapi import FastAPI, HTTPException
from fastapi.responses import Response, FileResponse
from fastapi.middleware.gzip import GZipMiddleware
import geopandas as gpd
import cbor2
from pathlib import Path

app = FastAPI()

# ✅ Add Gzip compression middleware
app.add_middleware(GZipMiddleware, minimum_size=500)

# Directory containing Parquet files
GEOHASH_FOLDER = Path("geohash")

@app.get("/get_geodata/")
async def get_geodata(geohash: str, format: str = "cbor"):
    """Fetch geospatial data by Geohash and return in the requested format with HTTP compression."""
    
    if len(geohash) < 4:
        raise HTTPException(status_code=400, detail="Geohash must be at least 4 characters long.")

    base_geohash = geohash[:4]
    parquet_file = GEOHASH_FOLDER / f"{base_geohash}.parquet"

    if not parquet_file.exists():
        raise HTTPException(status_code=404, detail=f"No data found for geohash '{base_geohash}'.")

    if format == "parquet":
        return FileResponse(parquet_file, media_type="application/octet-stream",
                            filename=f"{base_geohash}.parquet")

    gdf = gpd.read_parquet(parquet_file)

    if format == "geojson":
        geojson_data = gdf.to_json()
        return Response(content=geojson_data, media_type="application/geo+json",
                        headers={"Content-Disposition": f"attachment; filename={base_geohash}.geojson"})

    elif format == "cbor":
        geojson_dict = gdf.to_json()
        cbor_data = cbor2.dumps(geojson_dict)

        return Response(content=cbor_data, media_type="application/cbor",
                        headers={"Content-Disposition": f"attachment; filename={base_geohash}.cbor"})

    raise HTTPException(status_code=400, detail="Invalid format. Use 'parquet', 'geojson', or 'cbor'.")