src/app/api.py

"""
FastAPI backend -- expose pipeline (Schema Linker -> Model sinh SQL -> SQL
Validator -> Executor -> Auto Visualization) qua REST API, de frontend rieng
(Stitch hoac bat ky frontend nao khac) goi qua fetch/AJAX. Tai su dung dung
cac module da viet cho ban Streamlit (schema_builder/inference/sql_validator/
executor/visualizer) -- 2 frontend (Streamlit + Stitch) dung chung 1 backend
logic, khong code lap lai.

Ho tro 2 nguon du lieu, chuyen doi qua _active["mode"] (state toan cuc don
gian -- demo 1-nguoi-dung, KHONG can session/multi-user thuc su):
  - "fixed"    : database Spider co dinh (retail_store -- 7-table relational DB)
  - "uploaded" : file nguoi dung upload (CSV/SQLite/Excel -- nhieu bang duoc ho tro)

mode chi quyet dinh SCHEMA STRING duoc build nhu the nao (Spider-style co
PK/FK vs WikiSQL-style khong PK/FK) -- KHONG con ep checkpoint nao phai dung
theo mode. Nguoi dung TU CHON model_id (1 trong 12 checkpoint cua
inference.MODEL_CATALOG, span 4 stage x 3 architecture) qua /api/model, doc
lap voi mode -- da verify checkpoint Stage 2/3 van sinh dung SQL tren CA HAI
dang schema (xem inference.py).

Chay:  ./venv/Scripts/python.exe -m uvicorn src.app.api:app --reload --port 8000
Docs interactif (Swagger UI, tham khao luc thiet ke Stitch prompt): http://localhost:8000/docs
"""

import json
import sys
import time
from pathlib import Path
from typing import Any, List, Optional

from fastapi import FastAPI, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel

_STATIC_DIR = Path(__file__).resolve().parent / "giao_dien" / "text_to_sql_dashboard"

sys.path.insert(0, str(Path(__file__).resolve().parent))
import csv_loader  # noqa: E402
import executor  # noqa: E402
import inference  # noqa: E402
import schema_builder  # noqa: E402
import sql_validator  # noqa: E402
import visualizer  # noqa: E402

DEFAULT_DB_ID = "retail_store"
DEFAULT_DB_PATH = str(
    Path(__file__).resolve().parents[2]
    / "spider_data"
    / "database"
    / DEFAULT_DB_ID
    / f"{DEFAULT_DB_ID}.sqlite"
)

EXAMPLE_QUESTIONS = [
    # Simple — single metric
    "How many orders are there in total?",
    "How many customers are from hanoi?",
    "What is the average product price?",
    # Simple — bar chart (GROUP BY, single table)
    "How many orders are there for each status?",
    "How many customers live in each city?",
    # Medium — bar chart (no JOIN)
    "Show each region and the number of stores in that region.",
    # JOIN — bar chart
    "Show each category name and the number of products in it.",
    "For each store, show the store name and the number of orders.",
    "Show each category name and the total cost of all products in it.",
    "Show the 5 customers with the most orders. List customer name and order count.",
]

def _default_active():
    return {
        "mode": "fixed",
        "db_id": DEFAULT_DB_ID,
        "db_path": DEFAULT_DB_PATH,
        "schema_string": schema_builder.get_schema_string(DEFAULT_DB_ID),
        "model_id": inference.DEFAULT_MODEL_ID_FOR_MODE["fixed"],
    }


_active = _default_active()

app = FastAPI(
    title="Text-to-SQL Demo API",
    description="Question (English) -> SQL -> SQLite result, backed by fine-tuned T5/CodeT5/BART checkpoints.",
)

# allow_origins=["*"] -- demo noi bo/luan van, khong co auth/du lieu nhay cam.
# Sieta lai thanh domain cu the neu deploy public lau dai.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class QueryRequest(BaseModel):
    question: str


class QueryResponse(BaseModel):
    question: str
    sql: str
    gen_time: float
    exec_time: float = 0.0
    is_safe: bool
    validator_message: str
    exec_error: Optional[str] = None
    columns: List[str] = []
    rows: List[List[Any]] = []
    chart_type: str = "table"


class UploadResponse(BaseModel):
    db_id: str
    schema_text: str
    columns: List[str]
    row_count: int
    rows: List[List[Any]]
    truncated: bool
    model_id: str


class ModelRequest(BaseModel):
    model_id: str


@app.get("/")
def serve_frontend():
    return FileResponse(str(_STATIC_DIR / "code.html"), media_type="text/html")


@app.get("/usth_logo.png")
def serve_logo():
    return FileResponse(str(_STATIC_DIR / "usth_logo.png"), media_type="image/png")


@app.get("/api/health")
def health():
    return {"status": "ok"}


@app.get("/api/meta")
def meta():
    """Thong tin nguon du lieu dang active (schema, model, vi du cau hoi) --
    frontend goi lai sau moi /api/upload, /api/reset, hoac /api/model de cap
    nhat UI."""
    cfg = inference.get_checkpoint_config(_active["model_id"])
    return {
        "mode": _active["mode"],
        "db_id": _active["db_id"],
        "schema": _active["schema_string"],
        "model_id": _active["model_id"],
        "model": cfg["label"],
        # Vi du cau hoi chi hop ly cho database co dinh (retail_store) --
        # frontend nen an phan nay khi mode="uploaded".
        "examples": EXAMPLE_QUESTIONS if _active["mode"] == "fixed" else [],
    }


@app.get("/api/models")
def models():
    """Toan bo catalog 12 checkpoint (4 stage x 3 architecture) de frontend
    render dropdown -- nguoi dung tu chon model_id bat ky, khong gioi han theo
    mode dang active."""
    options = [
        {
            "model_id": model_id,
            "label": cfg["label"],
            "architecture": cfg["architecture"],
            "stage": cfg["stage"],
        }
        for model_id, cfg in inference.MODEL_CATALOG.items()
    ]
    options.sort(key=lambda o: (o["stage"] if isinstance(o["stage"], int) else 99, o["architecture"]))
    return {"current_model_id": _active["model_id"], "options": options}


@app.post("/api/model")
def set_model(req: ModelRequest):
    """Doi model_id (1 trong 12 checkpoint cua MODEL_CATALOG). Mode (fixed/
    uploaded) KHONG doi qua endpoint nay."""
    if req.model_id not in inference.MODEL_CATALOG:
        return {"status": "error", "message": f"Unknown model_id '{req.model_id}'"}
    _active["model_id"] = req.model_id
    return {"status": "ok", "model_id": _active["model_id"]}


@app.post("/api/upload", response_model=UploadResponse)
def upload(file: UploadFile):
    filename = file.filename or "uploaded"
    ext = Path(filename).suffix.lower()

    if ext in (".sqlite", ".db"):
        result = csv_loader.load_sqlite_file(file.file)
    elif ext in (".xlsx", ".xls"):
        result = csv_loader.load_excel_to_sqlite(file.file)
    else:
        result = csv_loader.load_csv_to_sqlite(file.file)

    _active.update(
        mode="uploaded",
        db_id=filename,
        db_path=result["db_path"],
        schema_string=result["schema_string"],
    )

    return UploadResponse(
        db_id=_active["db_id"],
        schema_text=result["schema_string"],
        columns=result["columns"],
        row_count=result["row_count"],
        rows=result["rows"],
        truncated=result["truncated"],
        model_id=_active["model_id"],
    )


@app.get("/api/preview")
def preview_table(table: str = ""):
    """Return rows from a specific table (or first available) for the active database."""
    db_path = _active["db_path"]
    table_names = csv_loader.get_table_names(db_path)
    if not table_names:
        return {"table_names": [], "active_table": "", "columns": [], "rows": [], "row_count": 0, "truncated": False}
    if table not in table_names:
        table = table_names[0]
    preview = csv_loader.get_preview_for_table(db_path, table)
    return {"table_names": table_names, "active_table": table, **preview}


@app.post("/api/reset")
def reset():
    """Quay lai database mac dinh (retail_store). Giu nguyen model_id nguoi
    dung dang chon (vd van la 1 checkpoint BIRD Stage3 sau reset)."""
    model_id = _active["model_id"]
    _active.clear()
    _active.update(_default_active())
    _active["model_id"] = model_id
    return {"status": "ok", "db_id": _active["db_id"]}


@app.post("/api/query", response_model=QueryResponse)
def query(req: QueryRequest):
    question = req.question.strip()
    model_id = _active["model_id"]

    if _active["mode"] == "fixed":
        model_input = schema_builder.build_input(question, _active["db_id"])
    else:
        model_input = csv_loader.build_input(question, _active["schema_string"])

    t0 = time.time()
    sql = inference.generate_sql(model_input, model_id=model_id)
    gen_time = time.time() - t0

    is_safe, msg = sql_validator.validate_sql(sql)
    resp = QueryResponse(
        question=question, sql=sql, gen_time=gen_time,
        is_safe=is_safe, validator_message=msg,
    )
    if not is_safe:
        return resp

    t_exec = time.time()
    try:
        df = executor.run_query(_active["db_path"], sql)
    except Exception as e:
        resp.exec_time = time.time() - t_exec
        resp.exec_error = str(e)
        return resp
    resp.exec_time = time.time() - t_exec

    chart_type, _chart_df = visualizer.make_chart_data(df)
    # Khi chart la bar/line, frontend dung r[0] lam label truc X.
    # Neu model sinh aggregate truoc (vd SELECT COUNT(*), city), reorder de
    # cot text (label) luon o vi tri 0 truoc khi serialize.
    if chart_type in ("bar", "line"):
        text_cols = df.select_dtypes(exclude="number").columns.tolist()
        if text_cols and df.columns.tolist()[0] != text_cols[0]:
            label = text_cols[0]
            rest = [c for c in df.columns if c != label]
            df = df[[label] + rest]
    # df.to_json() xu ly dung numpy int64/float64/NaN -> JSON-safe, khac voi
    # .values.tolist() truc tiep (co the tra ve kieu numpy khong serialize duoc).
    resp.columns = [str(c) for c in df.columns]
    resp.rows = json.loads(df.to_json(orient="values"))
    resp.chart_type = chart_type
    return resp


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)