# app.py
from fastapi import FastAPI, Request, Form
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
import wandb

app = FastAPI()
templates = Jinja2Templates(directory="templates")

@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/", response_class=HTMLResponse)
async def process_form(
    request: Request,
    token: str = Form(...),
    entity: str = Form(...),
    project: str = Form(...),
    run_id: str = Form(...)
):
    try:
        wandb.login(token)
        api = wandb.Api()
        run_path = f"{entity}/{project}/runs/{run_id}"
        run = api.run(run_path)
        iframe_html = run.to_html()
        return templates.TemplateResponse(
            "index.html", 
            {
                "request": request, 
                "iframe_html": iframe_html
            }
        )
    except Exception as e:
        return templates.TemplateResponse(
            "index.html", 
            {
                "request": request, 
                "error": str(e)
            }
        )