app.py

import os
import boto3
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse, HTMLResponse
from botocore.exceptions import ClientError

app = FastAPI(title="Energy Modelling Tools")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://starter-data-kit.github.io", "http://localhost", "*"],
    allow_methods=["GET"],
    allow_headers=["*"],
)

# --- S3 Configuration (use HF Spaces secrets / env vars) ---
AWS_ACCESS_KEY_ID = os.environ.get("AWS_ACCESS_KEY_ID", "AKIA4GFCYKIGZOAX2ELA")
AWS_SECRET_ACCESS_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY", "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")

s3 = boto3.client(
    "s3",
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
)

# --- Tool / dataset / bucket configuration ---
TOOL_CONFIG = {
    "maed": {
        "bucket": "maed",
        "label": "MAED",
        "full_name": "Model for Analysis of Energy Demand",
        "datasets": {
            "GDP by country": "GDP by country",
            "Jan2026_updated": "Jan2026_updated",
            "household_size": "household_size",
            "ilo_activeLabour": "ilo_activeLabour",
            "rural_population": "rural_population",
            "urban_population": "urban_population",
            "steel": "steel",
        },
    },
    "OnStove": {
        "bucket": "geospatialsdk",
        "label": "OnStove",
        "full_name": "Open-source Spatial Clean Cooking",
        "datasets": {
            "Socio-economic files": "Africa_OnStove/OnStove inputs and outputs/Inputs/Socio-economic files",
            "Techno-economic files": "Africa_OnStove/OnStove inputs and outputs/Inputs/Techno-economic files",
        },
    },
}

VALID_EXT = (".csv", ".xlsx", ".xls")


def list_s3_files(bucket: str, prefix: str):
    """List files in an S3 prefix, filtered by extension."""
    try:
        paginator = s3.get_paginator("list_objects_v2")
        files = []
        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            for obj in page.get("Contents", []):
                key = obj["Key"]
                if not key.endswith("/") and key.endswith(VALID_EXT):
                    # Return filename relative to prefix
                    rel = key[len(prefix) :].lstrip("/")
                    if rel:
                        files.append(rel)
        return sorted(files)
    except ClientError as e:
        raise HTTPException(status_code=500, detail=str(e))


# ---- API Endpoints ----

@app.get("/api/tools")
def get_tools():
    return {
        k: {"label": v["label"], "full_name": v["full_name"]}
        for k, v in TOOL_CONFIG.items()
    }


@app.get("/api/datasets/{tool}")
def get_datasets(tool: str):
    if tool not in TOOL_CONFIG:
        raise HTTPException(404, "Tool not found")
    return {
        "datasets": [
            {"label": label, "value": prefix}
            for label, prefix in TOOL_CONFIG[tool]["datasets"].items()
        ]
    }


@app.get("/api/files/{tool}")
def get_files(tool: str, dataset: str):
    if tool not in TOOL_CONFIG:
        raise HTTPException(404, "Tool not found")
    bucket = TOOL_CONFIG[tool]["bucket"]
    return {"files": list_s3_files(bucket, dataset)}


@app.get("/api/download/{tool}")
def download_file(tool: str, dataset: str, file: str):
    if tool not in TOOL_CONFIG:
        raise HTTPException(404, "Tool not found")
    bucket = TOOL_CONFIG[tool]["bucket"]
    key = f"{dataset}/{file}"
    try:
        obj = s3.get_object(Bucket=bucket, Key=key)
        content_type = "application/octet-stream"
        if file.endswith(".csv"):
            content_type = "text/csv"
        elif file.endswith((".xlsx", ".xls")):
            content_type = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        return StreamingResponse(
            obj["Body"],
            media_type=content_type,
            headers={"Content-Disposition": f'attachment; filename="{file}"'},
        )
    except ClientError as e:
        raise HTTPException(500, str(e))


# ---- Serve static files & pages ----

@app.get("/")
def index():
    return FileResponse("static/index.html")


@app.get("/tool/{tool_id}")
def tool_page(tool_id: str):
    return FileResponse("static/download.html")


app.mount("/static", StaticFiles(directory="static"), name="static")