import re
from fastapi import FastAPI, Request, Form
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse
import wandb
import os

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: str = Form(...),
    entity: str = Form(...),
    project: str = Form(...),
    run_id: str = Form(...)
):
    try:
        # Set the token as an environment variable
        os.environ["WANDB_API_KEY"] = token
        # Login with the anonymous parameter set
        wandb.login(key=token)
        
        api = wandb.Api()
        run_path = f"{entity}/{project}/runs/{run_id}"
        run = api.run(run_path)
        iframe_html = run.to_html()

        # Modify the iframe height to be 100%
        iframe_html = re.sub(r'height:\d+px', 'height:100%', iframe_html)
        
        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)
            }
        )