# app.py
from fastapi import FastAPI, Request, Form
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse
import wandb
import os
from typing import Optional

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

@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: Optional[str] = Form(None),  # Make token optional
    entity: str = Form(...),
    project: str = Form(...),
    run_id: str = Form(...)
):
    try:
        # Only set token and login if one is provided
        if token and token.strip():
            os.environ["WANDB_API_KEY"] = token
                    
        api = wandb.Api()
        run_path = f"{entity}/{project}/runs/{run_id}"
        run = api.run(run_path)
        iframe_html = run.to_html()
        iframe_html = iframe_html.replace('height:420px', 'height:100%')
        
        return templates.TemplateResponse(
            "index.html", 
            {
                "request": request,
                "token": token,
                "entity": entity,
                "project": project,
                "run_id": run_id,
                "iframe_html": iframe_html
            }
        )
    except Exception as e:
        return templates.TemplateResponse(
            "index.html", 
            {
                "request": request,
                "token": token,
                "entity": entity,
                "project": project,
                "run_id": run_id,
                "error": str(e)
            }
        )