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 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 app.config import OUTPUT_DIR, STATIC_DIR, TEMPLATES_DIR
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

app = FastAPI(title="PIMCO Chart Generator")
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
templates = Jinja2Templates(directory=str(TEMPLATES_DIR))

OUTPUT_DIR.mkdir(exist_ok=True)

# Simple in-memory session store: session_id -> {spec, data_path, sheets, summary, history}
_sessions: dict[str, dict] = {}


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


@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:
        # Save uploaded file
        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)

        # Load and prepare data
        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)

        # Interpret brief with Claude
        spec = interpret_brief(brief, summary)
        spec = validate_and_fix_spec(spec, prepared)

        # Store session
        _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."},
            ],
        }

        # Render
        svg, filename, spec_json = _render_and_save(spec, prepared, width, height)

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

    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"]

        # Add the edit to history
        history.append({"role": "user", "message": edit})

        # Ask Claude to refine the spec
        new_spec = refine_spec(old_spec, edit, summary, history)
        new_spec = validate_and_fix_spec(new_spec, session["prepared"])

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

        # Render
        svg, filename, spec_json = _render_and_save(
            new_spec, session["prepared"], session["width"], session["height"]
        )

        return templates.TemplateResponse("preview.html", {
            "request": request,
            "svg_content": svg,
            "filename": filename,
            "spec_json": spec_json,
            "session_id": session_id,
            "history": history,
        })

    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)
    return FileResponse(
        filepath,
        media_type="image/svg+xml",
        filename=filename,
    )


def _render_and_save(
    spec: ChartSpec,
    prepared: dict,
    width: int,
    height: int,
) -> tuple[str, str, str]:
    """Render a ChartSpec and save the SVG. Returns (svg_str, filename, spec_json)."""
    LAYOUT["single"]["width"] = width
    LAYOUT["single"]["height"] = height
    LAYOUT["dual_panel"]["width"] = width
    LAYOUT["dual_panel"]["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)

    filename = f"chart_{uuid.uuid4().hex[:8]}.svg"
    svg_path = OUTPUT_DIR / filename
    with open(svg_path, "w") as f:
        f.write(svg)

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