pimco-charts/app/main.py

"""FastAPI application: upload data, interpret brief, render SVG, iterate."""

from __future__ import annotations
import json
import uuid
import traceback
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, UploadFile, File, Form, Request
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from starlette.middleware.sessions import SessionMiddleware
from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware

from app.config import (
    OUTPUT_DIR, STATIC_DIR, TEMPLATES_DIR, SESSION_SECRET_KEY,
    AZURE_CLIENT_ID, AZURE_TENANT_ID, AZURE_REDIRECT_URI,
)
from app.auth.middleware import AuthMiddleware
from app.auth.routes import router as auth_router
from app.data.loader import load_file
from app.data.analyzer import summarize_data
from app.data.transformer import prepare_dataframe
from app.ai.brief_interpreter import interpret_brief, refine_spec
from app.ai.spec_validator import validate_and_fix_spec
from app.models.chart_spec import ChartSpec
from app.models.style import LAYOUT
from app.renderer.engine import render_chart
from app.renderer import export as svg_export


@asynccontextmanager
async def lifespan(app: FastAPI):
    await svg_export.init_browser()
    yield
    await svg_export.close_browser()


app = FastAPI(title="PIMCO Chart Generator", root_path="/Pimco-charts", lifespan=lifespan)
app.add_middleware(AuthMiddleware)
app.add_middleware(SessionMiddleware, secret_key=SESSION_SECRET_KEY, https_only=True, same_site="lax")
app.add_middleware(ProxyHeadersMiddleware, trusted_hosts="127.0.0.1")
app.include_router(auth_router)
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
templates = Jinja2Templates(directory=str(TEMPLATES_DIR))

OUTPUT_DIR.mkdir(exist_ok=True)

_sessions: dict[str, dict] = {}


@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    user = request.session.get("user")
    if not user:
        return templates.TemplateResponse("msal_shell.html", {
            "request": request,
            "client_id": AZURE_CLIENT_ID,
            "tenant_id": AZURE_TENANT_ID,
            "redirect_uri": AZURE_REDIRECT_URI,
            "root_path": request.scope.get("root_path", ""),
        })
    return templates.TemplateResponse("upload.html", {
        "request": request,
        "user": user,
    })


@app.post("/generate", response_class=HTMLResponse)
async def generate(
    request: Request,
    file: UploadFile = File(...),
    brief: str = Form(...),
    sheet: str = Form(""),
    width: int = Form(2560),
    height: int = Form(1440),
):
    try:
        session_id = uuid.uuid4().hex[:12]
        data_path = OUTPUT_DIR / f"data_{session_id}_{file.filename}"
        with open(data_path, "wb") as f:
            content = await file.read()
            f.write(content)

        sheets = load_file(data_path)
        if sheet and sheet in sheets:
            sheets = {sheet: sheets[sheet]}

        prepared = {}
        for name, df in sheets.items():
            prepared[name] = prepare_dataframe(df)

        summary = summarize_data(prepared)

        spec = interpret_brief(brief, summary)
        spec = validate_and_fix_spec(spec, prepared)

        _sessions[session_id] = {
            "spec": spec,
            "data_path": str(data_path),
            "prepared": prepared,
            "summary": summary,
            "width": width,
            "height": height,
            "history": [
                {"role": "user", "message": brief},
                {"role": "assistant", "message": "Chart generated."},
            ],
        }

        svg, base_name, spec_json = await _render_and_save(spec, prepared, width, height)

        return templates.TemplateResponse("preview.html", {
            "request": request,
            "svg_content": svg,
            "base_name": base_name,
            "spec_json": spec_json,
            "session_id": session_id,
            "history": _sessions[session_id]["history"],
            "user": request.session.get("user"),
        })

    except Exception as e:
        traceback.print_exc()
        return HTMLResponse(
            f'<div class="error"><h3>Error</h3><p>{str(e)}</p></div>',
            status_code=200,
        )


@app.post("/refine", response_class=HTMLResponse)
async def refine(
    request: Request,
    session_id: str = Form(...),
    edit: str = Form(...),
):
    try:
        session = _sessions.get(session_id)
        if not session:
            return HTMLResponse(
                '<div class="error"><h3>Session expired</h3>'
                '<p>Please re-upload your data and generate a new chart.</p></div>',
                status_code=200,
            )

        old_spec = session["spec"]
        summary = session["summary"]
        history = session["history"]

        history.append({"role": "user", "message": edit})

        new_spec = refine_spec(old_spec, edit, summary, history)
        new_spec = validate_and_fix_spec(new_spec, session["prepared"])

        session["spec"] = new_spec
        history.append({"role": "assistant", "message": "Chart updated."})

        svg, base_name, spec_json = await _render_and_save(
            new_spec, session["prepared"], session["width"], session["height"]
        )

        return templates.TemplateResponse("preview.html", {
            "request": request,
            "svg_content": svg,
            "base_name": base_name,
            "spec_json": spec_json,
            "session_id": session_id,
            "history": history,
            "user": request.session.get("user"),
        })

    except Exception as e:
        traceback.print_exc()
        return HTMLResponse(
            f'<div class="error"><h3>Error</h3><p>{str(e)}</p></div>',
            status_code=200,
        )


@app.get("/download/{filename}")
async def download(filename: str):
    filepath = OUTPUT_DIR / filename
    if not filepath.exists():
        return HTMLResponse("<p>File not found</p>", status_code=404)
    ext = Path(filename).suffix.lower()
    media_types = {
        ".svg": "image/svg+xml",
        ".png": "image/png",
        ".pdf": "application/pdf",
    }
    media_type = media_types.get(ext, "application/octet-stream")
    return FileResponse(filepath, media_type=media_type, filename=filename)


async def _render_and_save(
    spec: ChartSpec,
    prepared: dict,
    width: int,
    height: int,
) -> tuple[str, str, str]:
    """Render a ChartSpec and save SVG, PNG, and PDF. Returns (svg_str, base_name, spec_json)."""
    for layout_key in ("single", "dual_panel", "dual_y_axis"):
        LAYOUT[layout_key]["width"] = width
        LAYOUT[layout_key]["height"] = height

    data_dict = {}
    if len(prepared) == 1:
        data_dict["_default"] = next(iter(prepared.values()))
    else:
        data_dict = dict(prepared)
        data_dict["_default"] = next(iter(prepared.values()))

    svg = render_chart(spec, data_dict)

    base_name = f"chart_{uuid.uuid4().hex[:8]}"
    svg_bytes = svg.encode("utf-8")

    with open(OUTPUT_DIR / f"{base_name}.svg", "w") as f:
        f.write(svg)

    png_bytes = await svg_export.svg_to_png(svg, width, height)
    with open(OUTPUT_DIR / f"{base_name}.png", "wb") as f:
        f.write(png_bytes)

    pdf_bytes = await svg_export.svg_to_pdf(svg, width, height)
    with open(OUTPUT_DIR / f"{base_name}.pdf", "wb") as f:
        f.write(pdf_bytes)

    spec_json = json.dumps(spec.model_dump(), indent=2, default=str)
    return svg, base_name, spec_json