app.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from typing import Any, Dict, Optional
import time
import misata
import pandas as pd

app = FastAPI(
    title="Misata API",
    description="Synthetic data generation from natural language. Powered by Misata.",
    version=misata.__version__,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["POST", "GET"],
    allow_headers=["*"],
)

# ── Request / response models ─────────────────────────────────────────────────

class GenerateRequest(BaseModel):
    story: str = Field(..., description="Plain-English description of the dataset")
    rows: int = Field(default=500, ge=1, le=5000,
                      description="Number of rows for the primary table (max 5 000)")
    seed: Optional[int] = Field(default=None, description="Random seed for reproducibility")

class ParseRequest(BaseModel):
    story: str
    rows: int = Field(default=500, ge=1, le=5000)

# ── Helpers ───────────────────────────────────────────────────────────────────

def _df_to_records(df: pd.DataFrame) -> list:
    """Convert a DataFrame to JSON-safe records."""
    return df.where(pd.notnull(df), None).to_dict(orient="records")

def _tables_to_json(tables: Dict[str, Any]) -> Dict[str, list]:
    return {name: _df_to_records(df) for name, df in tables.items()}

# ── Endpoints ─────────────────────────────────────────────────────────────────

@app.get("/", include_in_schema=False)
def root():
    return {"name": "Misata API", "version": misata.__version__,
            "docs": "/docs", "health": "/health"}

@app.get("/health")
def health():
    return {"status": "ok", "version": misata.__version__}

@app.post("/generate")
def generate(req: GenerateRequest):
    """
    Parse a plain-English story and return generated tables as JSON.

    Returns `{ tables: { table_name: [ {col: val, ...}, ... ] }, meta: {...} }`
    """
    t0 = time.perf_counter()
    try:
        schema = misata.parse(req.story, rows=req.rows)
        if req.seed is not None:
            schema.seed = req.seed
        tables = misata.generate_from_schema(schema)
    except Exception as exc:
        raise HTTPException(status_code=422, detail=str(exc))

    elapsed = round(time.perf_counter() - t0, 3)
    row_counts = {name: len(df) for name, df in tables.items()}

    return {
        "tables": _tables_to_json(tables),
        "meta": {
            "story": req.story,
            "seed": req.seed,
            "elapsed_seconds": elapsed,
            "row_counts": row_counts,
            "schema_summary": schema.summary(),
        },
    }

@app.post("/parse")
def parse(req: ParseRequest):
    """Return the inferred schema without generating data."""
    try:
        schema = misata.parse(req.story, rows=req.rows)
    except Exception as exc:
        raise HTTPException(status_code=422, detail=str(exc))
    return {
        "name": schema.name,
        "domain": schema.domain,
        "tables": [
            {
                "name": t.name,
                "row_count": t.row_count,
                "columns": [c.name for c in schema.get_columns(t.name)],
            }
            for t in schema.tables
        ],
        "relationships": [
            f"{r.parent_table}.{r.parent_key} → {r.child_table}.{r.child_key}"
            for r in schema.relationships
        ],
        "summary": schema.summary(),
    }