app.py

import os
import io
import zipfile
import re
import difflib
import tempfile
import uuid
from typing import List, Optional, Dict, Any

from fastapi import FastAPI, UploadFile, File, HTTPException, Form, Header
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from langdetect import detect
from transformers import MarianMTModel, MarianTokenizer
from openai import OpenAI

# ---- Postgres (Neon) ----
import psycopg2
from psycopg2 import sql as pgsql

# ---- Supabase ----
from supabase import create_client, Client

SUPABASE_URL = "https://bnvmqgjawtaslczewqyd.supabase.co"
SUPABASE_ANON_KEY = (
    "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6ImJudm1x"
    "Z2phd3Rhc2xjemV3cXlkIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NjQ0NjM5NDAsImV4cCI6MjA4"
    "MDAzOTk0MH0.9zkyqrsm-QOSwMTUPZEWqyFeNpbbuar01rB7pmObkUI"
)

supabase: Client = create_client(SUPABASE_URL, SUPABASE_ANON_KEY)

# ======================================================
# 0) Configuración general de paths / modelo / OpenAI
# ======================================================

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

MODEL_DIR = os.getenv("MODEL_DIR", "stvnnnnnn/t5-large-nl2sql-spider")
DEVICE = torch.device("cpu")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
openai_client = OpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None

# DSN de Neon (Postgres) – EJEMPLO:
# postgres://user:pass@host/neondb?sslmode=require
POSTGRES_DSN = os.getenv("POSTGRES_DSN")

if not POSTGRES_DSN:
    raise RuntimeError(
        "⚠️ POSTGRES_DSN no está definido. "
        "Configúralo en los secrets del Space con la cadena de conexión de Neon."
    )

# ======================================================
# 1) Gestor de conexiones dinámicas: Postgres (Neon)
# ======================================================

class PostgresManager:
    """
    Cada upload crea un *schema* aislado en Neon.
    connections[connection_id] = {
        "label": str,        # nombre de archivo original
        "engine": "postgres",
        "schema": str        # nombre del schema en Neon
    }
    """

    def __init__(self, dsn: str):
        self.dsn = dsn
        self.connections: Dict[str, Dict[str, Any]] = {}

    # ---------- utilidades internas ----------

    def _new_connection_id(self) -> str:
        return f"db_{uuid.uuid4().hex[:8]}"

    def _get_info(self, connection_id: str) -> Dict[str, Any]:
        if connection_id not in self.connections:
            raise KeyError(f"connection_id '{connection_id}' no registrado")
        return self.connections[connection_id]

    def _get_conn(self, autocommit: bool = True):
        conn = psycopg2.connect(self.dsn)
        conn.autocommit = autocommit
        return conn

    def _execute_pg_dump(self, cur, sql_text: str) -> None:
        """
        Ejecuta un script tipo pg_dump:

          - Ejecuta DDL/INSERTs normales con execute
          - Maneja bloques:

                COPY tabla (...) FROM stdin;
                datos...
                \.

            usando copy_expert.
        """
        lines = sql_text.splitlines()
        n = len(lines)
        i = 0
        buffer: List[str] = []

        def flush_buffer():
            stmt = "\n".join(buffer).strip()
            if not stmt:
                return
            # Partimos por ';' para ejecutar cada sentencia
            for piece in stmt.split(";"):
                piece = piece.strip()
                if piece:
                    cur.execute(piece)

        while i < n:
            line = lines[i]

            # ¿Inicio de bloque COPY ... FROM stdin; ?
            if re.match(r"^\s*copy\s+.+from\s+stdin;?\s*$", line, re.IGNORECASE):
                # Ejecutar lo acumulado antes del COPY
                flush_buffer()
                buffer = []

                copy_sql = line.strip()
                i += 1
                data_lines: List[str] = []

                # Acumular las filas hasta encontrar '\.'
                while i < n and lines[i].strip() != r"\.":
                    data_lines.append(lines[i])
                    i += 1

                # Saltar la línea '\.' si existe
                if i < n and lines[i].strip() == r"\.":
                    i += 1

                data_str = "\n".join(data_lines) + "\n"

                # Ejecutar el COPY con los datos
                cur.copy_expert(copy_sql, io.StringIO(data_str))
            else:
                buffer.append(line)
                i += 1

        # Ejecutar lo que quede al final
        flush_buffer()

    # ---------- creación de BD desde dump ----------

    def create_database_from_dump(self, label: str, sql_text: str) -> str:
        """
        Crea un schema aislado en Neon y restaura dentro de él
        un dump de Postgres tipo psql (con COPY FROM stdin, funciones, etc.).

        NOTA:
        - Ignoramos bloques CREATE FUNCTION ... $$ ... $$ (no los necesitamos
          para hacer SELECT sobre las tablas).
        - Ignoramos errores tipo “already exists” para que no reviente si
          el script crea dos veces la misma tabla/índice.
        - Ignoramos sentencias que cambian OWNER a 'postgres'.
        """
        connection_id = self._new_connection_id()
        schema_name = f"sess_{uuid.uuid4().hex[:8]}"

        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                # 1) Crear schema aislado y fijar search_path
                cur.execute(
                    pgsql.SQL("CREATE SCHEMA {}").format(
                        pgsql.Identifier(schema_name)
                    )
                )
                cur.execute(
                    pgsql.SQL("SET search_path TO {}").format(
                        pgsql.Identifier(schema_name)
                    )
                )

                in_copy = False
                copy_sql = ""
                copy_rows: list[str] = []

                # NUEVO: estado para funciones
                in_function = False
                function_delim: str | None = None
                function_delim_count: int = 0

                stmt_lines: list[str] = []

                def flush_statement():
                    """Ejecuta el statement acumulado si es útil."""
                    nonlocal stmt_lines
                    stmt = "\n".join(stmt_lines).strip()
                    stmt_lines.clear()
                    if not stmt or stmt == ";":
                        return

                    upper = stmt.upper()

                    # Saltar cosas globales/peligrosas
                    skip_prefixes = (
                        "SET ",
                        "SELECT PG_CATALOG.SET_CONFIG",
                        "COMMENT ON EXTENSION",
                        "DROP DATABASE",
                        "CREATE DATABASE",
                        "ALTER DATABASE",
                        "REVOKE ",
                        "GRANT ",
                        "BEGIN",
                        "COMMIT",
                        "ROLLBACK",
                    )
                    if upper.startswith(skip_prefixes):
                        return

                    # Saltar ALTER ... OWNER TO postgres (u otros OWNER)
                    if " OWNER TO " in upper:
                        return

                    # Quitar ';' final (psycopg2 no la necesita)
                    if stmt.endswith(";"):
                        stmt = stmt[:-1]

                    try:
                        cur.execute(stmt)
                    except Exception as e:
                        msg = str(e).lower()
                        # Ignorar cosas no fatales típicas de dumps
                        if "already exists" in msg or "duplicate key value" in msg:
                            print("[WARN] Ignorando error no crítico:", e)
                            return
                        raise

                for raw_line in sql_text.splitlines():
                    line = raw_line.rstrip("\n")
                    stripped = line.strip()

                    # Comentarios o líneas vacías (si no estamos en COPY/función)
                    if not in_copy and not in_function:
                        if not stripped or stripped.startswith("--"):
                            continue

                    # ====== BLOQUE COPY ... FROM stdin ======
                    if in_copy:
                        if stripped == r"\.":
                            data_str = "\n".join(copy_rows) + "\n"
                            copy_rows.clear()
                            in_copy = False
                            cur.copy_expert(copy_sql, io.StringIO(data_str))
                        else:
                            copy_rows.append(line)
                        continue

                    # ====== BLOQUE CREATE FUNCTION ... $$ ... $$ (IGNORAR) ======
                    if in_function:
                        # Contamos apariciones del delimitador ($$ o $body$)
                        if function_delim:
                            function_delim_count += line.count(function_delim)
                            # Cuando lo hemos visto al menos 2 veces
                            # (apertura y cierre) salimos del bloque función
                            if function_delim_count >= 2:
                                in_function = False
                                function_delim = None
                                function_delim_count = 0
                        continue

                    # Detectar inicio de COPY
                    if (
                        stripped.upper().startswith("COPY ")
                        and "FROM stdin" in stripped.upper()
                    ):
                        flush_statement()
                        in_copy = True
                        copy_sql = stripped
                        copy_rows = []
                        continue

                    # Detectar inicio de función: CREATE [OR REPLACE] FUNCTION ...
                    if stripped.upper().startswith("CREATE FUNCTION") or stripped.upper().startswith(
                        "CREATE OR REPLACE FUNCTION"
                    ):
                        # Entramos en bloque función → lo ignoramos entero
                        in_function = True

                        # Detectar delimitador tipo $$ o $algo$
                        m = re.search(r"\$(\w*)\$", stripped)
                        if m:
                            function_delim = f"${m.group(1)}$"
                        else:
                            function_delim = "$$"

                        # Contamos cuántas veces aparece ya en esta línea
                        function_delim_count = stripped.count(function_delim)

                        # Caso raro: función toda en una sola línea (apertura+cierre)
                        if function_delim_count >= 2:
                            in_function = False
                            function_delim = None
                            function_delim_count = 0

                        # No añadimos la línea al buffer, se ignora
                        continue

                    # ====== STATEMENTS NORMALES ======
                    stmt_lines.append(line)
                    if stripped.endswith(";"):
                        flush_statement()

                # Por si quedó algo pendiente sin ';'
                if stmt_lines:
                    flush_statement()

        except Exception as e:
            # Limpiar schema si algo salió mal
            try:
                with conn.cursor() as cur:
                    cur.execute(
                        pgsql.SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(
                            pgsql.Identifier(schema_name)
                        )
                    )
            except Exception:
                pass
            conn.close()
            raise RuntimeError(f"Error ejecutando dump SQL en Postgres: {e}")
        finally:
            conn.close()

        self.connections[connection_id] = {
            "label": label,
            "engine": "postgres",
            "schema": schema_name,
        }
        return connection_id

    # ---------- ejecución segura de SQL ----------

    def execute_sql(self, connection_id: str, sql_text: str) -> Dict[str, Any]:
        """
        Ejecuta un SELECT dentro del schema asociado al connection_id.
        Bloquea operaciones destructivas por seguridad.
        """
        info = self._get_info(connection_id)
        schema = info["schema"]

        forbidden = ["drop ", "delete ", "update ", "insert ", "alter ", "replace "]
        sql_low = sql_text.lower()
        if any(tok in sql_low for tok in forbidden):
            return {
                "ok": False,
                "error": "Query bloqueada por seguridad (operación destructiva).",
                "rows": None,
                "columns": [],
            }

        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                # usar el schema de la sesión
                cur.execute(
                    pgsql.SQL("SET search_path TO {}").format(
                        pgsql.Identifier(schema)
                    )
                )
                cur.execute(sql_text)

                if cur.description:
                    rows = cur.fetchall()
                    cols = [d[0] for d in cur.description]
                else:
                    rows, cols = [], []

            return {
                "ok": True,
                "error": None,
                "rows": [list(r) for r in rows],
                "columns": cols,
            }
        except Exception as e:
            return {"ok": False, "error": str(e), "rows": None, "columns": []}
        finally:
            conn.close()

    # ---------- introspección de esquema ----------

    def get_schema(self, connection_id: str) -> Dict[str, Any]:
        info = self._get_info(connection_id)
        schema = info["schema"]

        conn = self._get_conn()
        try:
            tables_info: Dict[str, Dict[str, Any]] = {}
            foreign_keys: List[Dict[str, Any]] = []

            with conn.cursor() as cur:
                # Tablas básicas
                cur.execute(
                    """
                    SELECT table_name
                    FROM information_schema.tables
                    WHERE table_schema = %s
                      AND table_type = 'BASE TABLE'
                    ORDER BY table_name;
                    """,
                    (schema,),
                )
                tables = [r[0] for r in cur.fetchall()]

                # Columnas por tabla
                for t in tables:
                    cur.execute(
                        """
                        SELECT column_name
                        FROM information_schema.columns
                        WHERE table_schema = %s
                          AND table_name = %s
                        ORDER BY ordinal_position;
                        """,
                        (schema, t),
                    )
                    cols = [r[0] for r in cur.fetchall()]
                    tables_info[t] = {"columns": cols}

                # Foreign keys
                cur.execute(
                    """
                    SELECT
                        tc.table_name      AS from_table,
                        kcu.column_name    AS from_column,
                        ccu.table_name     AS to_table,
                        ccu.column_name    AS to_column
                    FROM information_schema.table_constraints AS tc
                    JOIN information_schema.key_column_usage AS kcu
                      ON tc.constraint_name = kcu.constraint_name
                     AND tc.table_schema   = kcu.table_schema
                    JOIN information_schema.constraint_column_usage AS ccu
                      ON ccu.constraint_name = tc.constraint_name
                     AND ccu.table_schema   = tc.table_schema
                    WHERE tc.constraint_type = 'FOREIGN KEY'
                      AND tc.table_schema    = %s;
                    """,
                    (schema,),
                )
                for ft, fc, tt, tc2 in cur.fetchall():
                    foreign_keys.append(
                        {
                            "from_table": ft,
                            "from_column": fc,
                            "to_table": tt,
                            "to_column": tc2,
                        }
                    )

            return {
                "tables": tables_info,
                "foreign_keys": foreign_keys,
            }
        finally:
            conn.close()

    # ---------- preview de tabla ----------

    def get_preview(
        self, connection_id: str, table: str, limit: int = 20
    ) -> Dict[str, Any]:
        info = self._get_info(connection_id)
        schema = info["schema"]

        conn = self._get_conn()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    pgsql.SQL("SET search_path TO {}").format(
                        pgsql.Identifier(schema)
                    )
                )
                query = pgsql.SQL("SELECT * FROM {} LIMIT %s").format(
                    pgsql.Identifier(table)
                )
                cur.execute(query, (int(limit),))
                rows = cur.fetchall()
                cols = [d[0] for d in cur.description] if cur.description else []

            return {
                "columns": cols,
                "rows": [list(r) for r in rows],
            }
        finally:
            conn.close()


# Instancia global de PostgresManager
sql_manager = PostgresManager(POSTGRES_DSN)

# ======================================================
# 2) Inicialización de FastAPI
# ======================================================

app = FastAPI(
    title="NL2SQL Backend (Supabase + Postgres/Neon)",
    version="3.0.0",
)

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

# ======================================================
# 3) Modelo NL→SQL y traductor ES→EN
# ======================================================

t5_tokenizer = None
t5_model = None
mt_tokenizer = None
mt_model = None


def load_nl2sql_model():
    """Carga el modelo NL→SQL (T5-large fine-tuned en Spider) desde HF Hub."""
    global t5_tokenizer, t5_model
    if t5_model is not None:
        return
    print(f"🔁 Cargando modelo NL→SQL desde: {MODEL_DIR}")
    t5_tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=True)
    t5_model = AutoModelForSeq2SeqLM.from_pretrained(
        MODEL_DIR, torch_dtype=torch.float32
    )
    t5_model.to(DEVICE)
    t5_model.eval()
    print("✅ Modelo NL→SQL listo en memoria.")


def load_es_en_translator():
    """Carga el modelo Helsinki-NLP para traducción ES→EN (solo una vez)."""
    global mt_tokenizer, mt_model
    if mt_model is not None:
        return
    model_name = "Helsinki-NLP/opus-mt-es-en"
    print(f"🔁 Cargando traductor ES→EN: {model_name}")
    mt_tokenizer = MarianTokenizer.from_pretrained(model_name)
    mt_model = MarianMTModel.from_pretrained(model_name)
    mt_model.to(DEVICE)
    mt_model.eval()
    print("✅ Traductor ES→EN listo.")


def detect_language(text: str) -> str:
    try:
        return detect(text)
    except Exception:
        return "unknown"


def translate_es_to_en(text: str) -> str:
    """
    Usa Marian ES→EN solo si el texto se detecta como español ('es').
    Si no, devuelve el texto tal cual.
    """
    lang = detect_language(text)
    if lang != "es":
        return text
    if mt_model is None:
        load_es_en_translator()
    inputs = mt_tokenizer(text, return_tensors="pt", truncation=True).to(DEVICE)
    with torch.no_grad():
        out = mt_model.generate(**inputs, max_length=256)
    return mt_tokenizer.decode(out[0], skip_special_tokens=True)


# ======================================================
# 4) Capa de reparación de SQL (usa el schema real)
# ======================================================

def _normalize_name_for_match(name: str) -> str:
    s = name.lower()
    s = s.replace('"', "").replace("`", "")
    s = s.replace("_", "")
    if s.endswith("s") and len(s) > 3:
        s = s[:-1]
    return s


def _build_schema_indexes(
    tables_info: Dict[str, Dict[str, List[str]]]
) -> Dict[str, Dict[str, List[str]]]:
    table_index: Dict[str, List[str]] = {}
    column_index: Dict[str, List[str]] = {}

    for t, info in tables_info.items():
        tn = _normalize_name_for_match(t)
        table_index.setdefault(tn, [])
        if t not in table_index[tn]:
            table_index[tn].append(t)

        for c in info.get("columns", []):
            cn = _normalize_name_for_match(c)
            column_index.setdefault(cn, [])
            if c not in column_index[cn]:
                column_index[cn].append(c)

    return {"table_index": table_index, "column_index": column_index}


def _best_match_name(missing: str, index: Dict[str, List[str]]) -> Optional[str]:
    if not index:
        return None

    key = _normalize_name_for_match(missing)
    if key in index and index[key]:
        return index[key][0]

    candidates = difflib.get_close_matches(key, list(index.keys()), n=1, cutoff=0.7)
    if not candidates:
        return None
    best_key = candidates[0]
    if index[best_key]:
        return index[best_key][0]
    return None


DOMAIN_SYNONYMS_TABLE = {
    "song": "track",
    "songs": "track",
    "tracks": "track",
    "artist": "artist",
    "artists": "artist",
    "album": "album",
    "albums": "album",
    "order": "invoice",
    "orders": "invoice",
}

DOMAIN_SYNONYMS_COLUMN = {
    "song": "name",
    "songs": "name",
    "track": "name",
    "title": "name",
    "length": "milliseconds",
    "duration": "milliseconds",
}


def try_repair_sql(sql: str, error: str, schema_meta: Dict[str, Any]) -> Optional[str]:
    """
    Intenta reparar nombres de tablas/columnas basándose en el esquema real.
    Compatible con mensajes de Postgres y también con los de SQLite
    (por si algún día reusamos la lógica).
    """
    tables_info = schema_meta["tables"]
    idx = _build_schema_indexes(tables_info)
    table_index = idx["table_index"]
    column_index = idx["column_index"]

    repaired_sql = sql
    changed = False

    missing_table = None
    missing_column = None

    m_t = re.search(r'relation "([\w\.]+)" does not exist', error, re.IGNORECASE)
    if not m_t:
        m_t = re.search(r"no such table: ([\w\.]+)", error)
    if m_t:
        missing_table = m_t.group(1)

    m_c = re.search(r'column "([\w\.]+)" does not exist', error, re.IGNORECASE)
    if not m_c:
        m_c = re.search(r"no such column: ([\w\.]+)", error)
    if m_c:
        missing_column = m_c.group(1)

    if missing_table:
        short = missing_table.split(".")[-1]
        syn = DOMAIN_SYNONYMS_TABLE.get(short.lower())
        target = None
        if syn:
            target = _best_match_name(syn, table_index) or syn
        if not target:
            target = _best_match_name(short, table_index)

        if target:
            pattern = r"\b" + re.escape(short) + r"\b"
            new_sql = re.sub(pattern, target, repaired_sql)
            if new_sql != repaired_sql:
                repaired_sql = new_sql
                changed = True

    if missing_column:
        short = missing_column.split(".")[-1]
        syn = DOMAIN_SYNONYMS_COLUMN.get(short.lower())
        target = None
        if syn:
            target = _best_match_name(syn, column_index) or syn
        if not target:
            target = _best_match_name(short, column_index)

        if target:
            pattern = r"\b" + re.escape(short) + r"\b"
            new_sql = re.sub(pattern, target, repaired_sql)
            if new_sql != repaired_sql:
                repaired_sql = new_sql
                changed = True

    if not changed:
        return None
    return repaired_sql


# ======================================================
# 5) Prompt NL→SQL + re-ranking
# ======================================================

def build_prompt(question_en: str, db_id: str, schema_str: str) -> str:
    return (
        f"translate to SQL: {question_en} | "
        f"db: {db_id} | schema: {schema_str} | "
        f"note: use JOIN when foreign keys link tables"
    )


def normalize_score(raw: float) -> float:
    """Normaliza el score logit del modelo a un porcentaje 0-100."""
    norm = (raw + 20) / 25
    norm = max(0, min(1, norm))
    return round(norm * 100, 2)


def nl2sql_with_rerank(question: str, conn_id: str) -> Dict[str, Any]:
    if conn_id not in sql_manager.connections:
        raise HTTPException(
            status_code=404, detail=f"connection_id '{conn_id}' no registrado"
        )

    meta = sql_manager.get_schema(conn_id)
    tables_info = meta["tables"]

    parts = []
    for t, info in tables_info.items():
        cols = info.get("columns", [])
        parts.append(f"{t}(" + ", ".join(cols) + ")")
    schema_str = " ; ".join(parts) if parts else "(empty_schema)"

    detected = detect_language(question)
    question_en = translate_es_to_en(question) if detected == "es" else question

    prompt = build_prompt(question_en, db_id=conn_id, schema_str=schema_str)

    if t5_model is None:
        load_nl2sql_model()

    inputs = t5_tokenizer(
        [prompt], return_tensors="pt", truncation=True, max_length=768
    ).to(DEVICE)
    num_beams = 6
    num_return = 6

    with torch.no_grad():
        out = t5_model.generate(
            **inputs,
            max_length=220,
            num_beams=num_beams,
            num_return_sequences=num_return,
            return_dict_in_generate=True,
            output_scores=True,
        )

    sequences = out.sequences
    scores = out.sequences_scores
    if scores is not None:
        scores = scores.cpu().tolist()
    else:
        scores = [0.0] * sequences.size(0)

    candidates: List[Dict[str, Any]] = []
    best = None
    best_exec = False
    best_score = -1e9

    for i in range(sequences.size(0)):
        raw_sql = t5_tokenizer.decode(
            sequences[i], skip_special_tokens=True
        ).strip()
        cand: Dict[str, Any] = {
            "sql": raw_sql,
            "score": float(scores[i]),
            "repaired_from": None,
            "repair_note": None,
            "raw_sql_model": raw_sql,
        }

        exec_info = sql_manager.execute_sql(conn_id, raw_sql)

        err_lower = (exec_info["error"] or "").lower()
        if (not exec_info["ok"]) and (
            "no such table" in err_lower
            or "no such column" in err_lower
            or "does not exist" in err_lower
        ):
            current_sql = raw_sql
            last_error = exec_info["error"] or ""
            for step in range(1, 4):
                repaired_sql = try_repair_sql(current_sql, last_error, meta)
                if not repaired_sql or repaired_sql == current_sql:
                    break
                exec_info2 = sql_manager.execute_sql(conn_id, repaired_sql)
                cand["repaired_from"] = (
                    current_sql if cand["repaired_from"] is None else cand["repaired_from"]
                )
                cand["repair_note"] = f"auto-repair (table/column name, step {step})"
                cand["sql"] = repaired_sql
                exec_info = exec_info2
                current_sql = repaired_sql
                if exec_info2["ok"]:
                    break
                last_error = exec_info2["error"] or ""

        cand["exec_ok"] = exec_info["ok"]
        cand["exec_error"] = exec_info["error"]
        cand["rows_preview"] = (
            exec_info["rows"][:5] if exec_info["ok"] and exec_info["rows"] else None
        )
        cand["columns"] = exec_info["columns"]

        candidates.append(cand)

        if exec_info["ok"]:
            if (not best_exec) or cand["score"] > best_score:
                best_exec = True
                best_score = cand["score"]
                best = cand
        elif not best_exec and cand["score"] > best_score:
            best_score = cand["score"]
            best = cand

    if best is None and candidates:
        best = candidates[0]

    return {
        "question_original": question,
        "detected_language": detected,
        "question_en": question_en,
        "connection_id": conn_id,
        "schema_summary": schema_str,
        "best_sql": best["sql"],
        "best_exec_ok": best.get("exec_ok", False),
        "best_exec_error": best.get("exec_error"),
        "best_rows_preview": best.get("rows_preview"),
        "best_columns": best.get("columns", []),
        "candidates": candidates,
        "score_percent": normalize_score(best["score"]),
    }


# ======================================================
# 6) Schemas Pydantic
# ======================================================

class UploadResponse(BaseModel):
    connection_id: str
    label: str
    db_path: str
    note: Optional[str] = None


class ConnectionInfo(BaseModel):
    connection_id: str
    label: str
    engine: Optional[str] = None
    db_name: Optional[str] = None  # ya no usamos archivo, pero mantenemos campo


class SchemaResponse(BaseModel):
    connection_id: str
    schema_summary: str
    tables: Dict[str, Dict[str, List[str]]]


class PreviewResponse(BaseModel):
    connection_id: str
    table: str
    columns: List[str]
    rows: List[List[Any]]


class InferRequest(BaseModel):
    connection_id: str
    question: str


class InferResponse(BaseModel):
    question_original: str
    detected_language: str
    question_en: str
    connection_id: str
    schema_summary: str
    best_sql: str
    best_exec_ok: bool
    best_exec_error: Optional[str]
    best_rows_preview: Optional[List[List[Any]]]
    best_columns: List[str]
    candidates: List[Dict[str, Any]]


class SpeechInferResponse(BaseModel):
    transcript: str
    result: InferResponse


# ======================================================
# 7) Helpers para /upload (.sql y .zip)
# ======================================================

def _combine_sql_files_from_zip(zip_bytes: bytes) -> str:
    """
    Lee un ZIP, se queda solo con los .sql y los concatena.
    Orden:
      1) archivos con 'schema' o 'structure' en el nombre
      2) el resto (data, etc.)
    """
    try:
        with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
            names = [info.filename for info in zf.infolist() if not info.is_dir()]
            sql_names = [n for n in names if n.lower().endswith(".sql")]

            if not sql_names:
                raise ValueError("El ZIP no contiene archivos .sql utilizables.")

            def sort_key(name: str) -> int:
                nl = name.lower()
                if "schema" in nl or "structure" in nl:
                    return 0
                return 1

            sql_names_sorted = sorted(sql_names, key=sort_key)

            parts: List[str] = []
            for name in sql_names_sorted:
                with zf.open(name) as f:
                    text = f.read().decode("utf-8", errors="ignore")
                parts.append(f"-- FILE: {name}\n{text}\n")

            return "\n\n".join(parts)
    except zipfile.BadZipFile:
        raise ValueError("Archivo ZIP inválido o corrupto.")


# ======================================================
# 8) Endpoints FastAPI
# ======================================================

@app.on_event("startup")
async def startup_event():
    load_nl2sql_model()
    print("✅ Backend NL2SQL inicializado (engine Postgres/Neon).")
    print(f"MODEL_DIR={MODEL_DIR}, DEVICE={DEVICE}")
    print(f"Conexiones activas al inicio: {len(sql_manager.connections)}")


@app.post("/upload", response_model=UploadResponse)
async def upload_database(
    mode: str = Form("full"),                 # "full" | "schema_data" | "zip"
    db_files: List[UploadFile] = File(...),  # uno o varios archivos
    authorization: Optional[str] = Header(None),
):
    """
    Sube uno o varios archivos SQL/ZIP según el modo:

    - mode = "full":
        * Espera EXACTAMENTE 1 archivo .sql
        * El .sql trae esquema + datos juntos (dump de PostgreSQL)

    - mode = "schema_data":
        * Espera EXACTAMENTE 2 archivos .sql
        * Uno de esquema y otro de datos (el orden lo resolvemos nosotros)

    - mode = "zip":
        * Espera EXACTAMENTE 1 archivo .zip
        * Dentro del zip buscamos SOLO archivos .sql (ignoramos el resto)
    """
    if authorization is None:
        raise HTTPException(401, "Missing Authorization header")

    jwt = authorization.replace("Bearer ", "")
    user = supabase.auth.get_user(jwt)
    if not user or not user.user:
        raise HTTPException(401, "Invalid Supabase token")

    if not db_files:
        raise HTTPException(400, "No se recibió ningún archivo.")

    mode = mode.lower().strip()

    # =======================
    # MODO 1: FULL (.sql único)
    # =======================
    if mode == "full":
        if len(db_files) != 1:
            raise HTTPException(
                400, "Modo FULL requiere exactamente 1 archivo .sql."
            )

        file = db_files[0]
        filename = file.filename or ""
        if not filename.lower().endswith(".sql"):
            raise HTTPException(400, "Modo FULL solo acepta archivos .sql.")

        contents = await file.read()
        sql_text = contents.decode("utf-8", errors="ignore")

    # ====================================
    # MODO 2: ESQUEMA + DATOS (2 archivos)
    # ====================================
    elif mode == "schema_data":
        if len(db_files) != 2:
            raise HTTPException(
                400,
                "Modo esquema+datos requiere exactamente 2 archivos .sql.",
            )

        files_info: List[tuple[str, str]] = []
        for f in db_files:
            fname = f.filename or ""
            if not fname.lower().endswith(".sql"):
                raise HTTPException(400, "Todos los archivos deben ser .sql.")
            contents = await f.read()
            files_info.append((fname, contents.decode("utf-8", errors="ignore")))

        # Intentamos poner primero el esquema y luego los datos
        def weight(name: str) -> int:
            nl = name.lower()
            if "schema" in nl or "structure" in nl:
                return 0
            if "data" in nl or "insert" in nl:
                return 1
            return 2

        files_info_sorted = sorted(files_info, key=lambda x: weight(x[0]))

        sql_parts: List[str] = []
        for fname, text in files_info_sorted:
            sql_parts.append(f"-- FILE: {fname}\n{text}\n")

        sql_text = "\n\n".join(sql_parts)
        # usamos el nombre del primer archivo como label "principal"
        filename = files_info_sorted[0][0]

    # ==================
    # MODO 3: ZIP (.zip)
    # ==================
    elif mode == "zip":
        if len(db_files) != 1:
            raise HTTPException(
                400, "Modo ZIP requiere exactamente 1 archivo .zip."
            )

        file = db_files[0]
        filename = file.filename or ""
        if not filename.lower().endswith(".zip"):
            raise HTTPException(400, "Modo ZIP solo acepta archivos .zip.")

        contents = await file.read()
        # tu helper ya ignora carpetas y solo concatena .sql
        sql_text = _combine_sql_files_from_zip(contents)

    else:
        raise HTTPException(400, f"Modo no soportado: {mode}")

    # --- crear schema dinámico en Postgres (Neon) ---
    try:
        conn_id = sql_manager.create_database_from_dump(
            label=filename, sql_text=sql_text
        )
    except Exception as e:
        raise HTTPException(400, f"Error creando BD: {e}")

    meta = sql_manager.connections[conn_id]

    # --- guardar metadatos en Supabase (sin romper el upload si falla) ---
    try:
        supabase.table("databases").insert(
            {
                "user_id": user.user.id,
                "filename": filename,
                "engine": meta["engine"],
                "connection_id": conn_id,
            }
        ).execute()
    except Exception as e:
        # Solo logeamos, pero NO rompemos el endpoint
        print("[WARN] No se pudieron guardar metadatos en Supabase:", repr(e))

    return UploadResponse(
        connection_id=conn_id,
        label=filename,
        db_path=f"{meta['engine']}://schema/{meta['schema']}",
        note="Database schema created in Neon and indexed in Supabase.",
    )


@app.get("/connections", response_model=List[ConnectionInfo])
async def list_connections():
    return [
        ConnectionInfo(
            connection_id=cid,
            label=meta.get("label", ""),
            engine=meta.get("engine"),
            db_name=meta.get("schema"),  # usamos schema como "nombre"
        )
        for cid, meta in sql_manager.connections.items()
    ]


@app.get("/schema/{connection_id}", response_model=SchemaResponse)
async def get_schema(connection_id: str):
    if connection_id not in sql_manager.connections:
        raise HTTPException(status_code=404, detail="connection_id no encontrado")

    meta = sql_manager.get_schema(connection_id)
    tables = meta["tables"]

    parts = []
    for t, info in tables.items():
        cols = info.get("columns", [])
        parts.append(f"{t}(" + ", ".join(cols) + ")")
    schema_str = " ; ".join(parts) if parts else "(empty_schema)"

    return SchemaResponse(
        connection_id=connection_id,
        schema_summary=schema_str,
        tables=tables,
    )


@app.get("/preview/{connection_id}/{table}", response_model=PreviewResponse)
async def preview_table(connection_id: str, table: str, limit: int = 20):
    if connection_id not in sql_manager.connections:
        raise HTTPException(status_code=404, detail="connection_id no encontrado")

    try:
        preview = sql_manager.get_preview(connection_id, table, limit)
    except Exception as e:
        raise HTTPException(
            status_code=400, detail=f"Error al leer tabla '{table}': {e}"
        )

    return PreviewResponse(
        connection_id=connection_id,
        table=table,
        columns=preview["columns"],
        rows=preview["rows"],
    )


@app.post("/infer", response_model=InferResponse)
async def infer_sql(
    req: InferRequest,
    authorization: Optional[str] = Header(None),
):
    if authorization is None:
        raise HTTPException(401, "Missing Authorization header")

    jwt = authorization.replace("Bearer ", "")
    user = supabase.auth.get_user(jwt)
    if not user or not user.user:
        raise HTTPException(401, "Invalid Supabase token")

    result = nl2sql_with_rerank(req.question, req.connection_id)
    score = normalize_score(result["candidates"][0]["score"])

    db_row = (
        supabase.table("databases")
        .select("id")
        .eq("connection_id", req.connection_id)
        .eq("user_id", user.user.id)
        .execute()
    )
    db_id = db_row.data[0]["id"] if db_row.data else None

    supabase.table("queries").insert(
        {
            "user_id": user.user.id,
            "db_id": db_id,
            "nl": result["question_original"],
            "sql_generated": result["best_sql"],
            "sql_repaired": result["candidates"][0]["sql"],
            "execution_ok": result["best_exec_ok"],
            "error": result["best_exec_error"],
            "rows_preview": result["best_rows_preview"],
            "score": score,
        }
    ).execute()

    result["score_percent"] = score
    return InferResponse(**result)


@app.post("/speech-infer", response_model=SpeechInferResponse)
async def speech_infer(
    connection_id: str = Form(...),
    audio: UploadFile = File(...),
):
    if openai_client is None:
        raise HTTPException(
            status_code=500,
            detail="OPENAI_API_KEY no está configurado en el backend.",
        )

    if audio.content_type is None:
        raise HTTPException(status_code=400, detail="Archivo de audio inválido.")

    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".webm") as tmp:
            tmp.write(await audio.read())
            tmp_path = tmp.name
    except Exception:
        raise HTTPException(
            status_code=500, detail="No se pudo procesar el audio recibido."
        )

    try:
        with open(tmp_path, "rb") as f:
            transcription = openai_client.audio.transcriptions.create(
                model="gpt-4o-transcribe",
                file=f,
            )
        transcript_text: str = transcription.text
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error al transcribir audio: {e}")

    result_dict = nl2sql_with_rerank(transcript_text, connection_id)
    infer_result = InferResponse(**result_dict)

    return SpeechInferResponse(
        transcript=transcript_text,
        result=infer_result,
    )


@app.get("/health")
async def health():
    return {
        "status": "ok",
        "model_loaded": t5_model is not None,
        "connections": len(sql_manager.connections),
        "device": str(DEVICE),
        "engine": "postgres",
    }


@app.get("/history")
def get_history(authorization: Optional[str] = Header(None)):
    if authorization is None:
        raise HTTPException(401, "Missing Authorization")

    jwt = authorization.replace("Bearer ", "")
    user = supabase.auth.get_user(jwt)

    rows = (
        supabase.table("queries")
        .select("*")
        .eq("user_id", user.user.id)
        .order("created_at", desc=True)
        .execute()
    )

    return rows.data


@app.get("/my-databases")
def get_my_databases(authorization: Optional[str] = Header(None)):
    if authorization is None:
        raise HTTPException(401, "Missing Authorization")

    jwt = authorization.replace("Bearer ", "")
    user = supabase.auth.get_user(jwt)

    rows = (
        supabase.table("databases")
        .select("*")
        .eq("user_id", user.user.id)
        .execute()
    )

    return rows.data


@app.get("/")
async def root():
    return {
        "message": "NL2SQL T5-large backend running on Postgres/Neon (no SQLite).",
        "endpoints": [
            "POST /upload            (subir .sql o .zip con .sql → crea schema en Neon)",
            "GET  /connections       (listar BDs subidas en esta instancia)",
            "GET  /schema/{id}       (esquema resumido)",
            "GET  /preview/{id}/{t}  (preview de tabla)",
            "POST /infer             (NL→SQL + ejecución en BD)",
            "POST /speech-infer      (voz → NL→SQL + ejecución)",
            "GET  /history           (historial de consultas en Supabase)",
            "GET  /my-databases      (BDs del usuario en Supabase)",
            "GET  /health            (estado del backend)",
            "GET  /docs              (OpenAPI UI)",
        ],
    }