src/app/csv_loader.py

"""
File upload loader -- supports CSV, SQLite/DB files, and Excel workbooks.

CSV:   single table "data_table" -> WikiSQL-format schema (no FK/PK)
SQLite/DB: copy to temp path, read multi-table schema via PRAGMA -> Spider-format schema
Excel: each sheet -> one table in temp SQLite -> Spider-format schema

The schema string format matches what models were trained on:
  CSV     -> WikiSQL format (process_wikisql.build_schema_string)
  SQLite/Excel -> Spider format  (table ( col*:type , ... ) | FK: ...)
"""

import shutil
import sqlite3
import sys
import tempfile
from pathlib import Path

import pandas as pd

_SRC_DIR = Path(__file__).resolve().parent.parent  # src/
if str(_SRC_DIR) not in sys.path:
    sys.path.insert(0, str(_SRC_DIR))

from process_wikisql import build_schema_string as _wikisql_schema  # noqa: E402

TABLE_NAME = "data_table"
_UPLOAD_DB_PATH = Path(tempfile.gettempdir()) / "_uploaded_data.sqlite"
MAX_PREVIEW_ROWS = 500


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _infer_wikisql_types(df: pd.DataFrame) -> list:
    return ["real" if pd.api.types.is_numeric_dtype(dt) else "text" for dt in df.dtypes]


def _lowercase_text_columns(df: pd.DataFrame, types: list) -> pd.DataFrame:
    """WikiSQL-trained models expect lowercase string literals in WHERE clauses."""
    df = df.copy()
    for col, t in zip(df.columns, types):
        if t == "text":
            df[col] = df[col].apply(lambda v: v.lower() if isinstance(v, str) else v)
    return df


def _sqlite_col_type(declared_type: str) -> str:
    """Map SQLite declared type -> 'number' or 'text' for schema string."""
    upper = declared_type.upper()
    if any(k in upper for k in ("INT", "REAL", "FLOAT", "NUMERIC", "DECIMAL", "DOUBLE")):
        return "number"
    return "text"


def _preview_from_db(db_path: str) -> dict:
    """Return preview rows from first table and total row count across all tables."""
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    tables = [r[0] for r in cur.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name"
    ).fetchall()]

    total_rows = 0
    first_table_df = None
    for tbl in tables:
        cnt = cur.execute(f'SELECT COUNT(*) FROM "{tbl}"').fetchone()[0]
        total_rows += cnt
        if first_table_df is None:
            first_table_df = pd.read_sql_query(
                f'SELECT * FROM "{tbl}" LIMIT {MAX_PREVIEW_ROWS}', conn
            )
    conn.close()

    if first_table_df is None or first_table_df.empty:
        return {"columns": [], "rows": [], "row_count": 0, "truncated": False}

    preview_df = first_table_df.astype(object).where(first_table_df.notna(), None)
    return {
        "columns": list(first_table_df.columns),
        "rows": preview_df.values.tolist(),
        "row_count": total_rows,
        "truncated": total_rows > MAX_PREVIEW_ROWS,
    }


def get_table_names(db_path: str) -> list:
    """Return all user table names in a SQLite database (alphabetical order)."""
    conn = sqlite3.connect(db_path)
    tables = [r[0] for r in conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name"
    ).fetchall()]
    conn.close()
    return tables


def get_preview_for_table(db_path: str, table: str) -> dict:
    """Return preview rows (up to MAX_PREVIEW_ROWS) from a specific table."""
    conn = sqlite3.connect(db_path)
    total = conn.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0]
    df = pd.read_sql_query(f'SELECT * FROM "{table}" LIMIT {MAX_PREVIEW_ROWS}', conn)
    conn.close()
    df = df.astype(object).where(df.notna(), None)
    return {
        "columns": list(df.columns),
        "rows": df.values.tolist(),
        "row_count": total,
        "truncated": total > MAX_PREVIEW_ROWS,
    }


def _build_spider_schema_from_sqlite(db_path: str) -> str:
    """Read schema from a SQLite file and return a Spider-format schema string.

    Format:  table ( col*:type , col2:type ) | table2 ( ... ) | foreign_keys: t.c = t2.c , ...
    """
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()

    tables = [r[0] for r in cur.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name"
    ).fetchall()]

    table_parts = []
    fk_parts = []

    for tbl in tables:
        tbl_lower = tbl.lower()
        cur.execute(f'PRAGMA table_info("{tbl}")')
        cols = cur.fetchall()  # (cid, name, type, notnull, default, pk)

        col_strs = []
        for col in cols:
            col_name  = col[1].lower()
            col_type  = _sqlite_col_type(col[2])
            pk_marker = "*" if col[5] > 0 else ""
            col_strs.append(f"{col_name}{pk_marker}:{col_type}")

        table_parts.append(f"{tbl_lower} ( {' , '.join(col_strs)} )")

        cur.execute(f'PRAGMA foreign_key_list("{tbl}")')
        for fk in cur.fetchall():  # (id, seq, ref_table, from_col, to_col, ...)
            ref_table = fk[2].lower()
            from_col  = fk[3].lower()
            to_col    = fk[4].lower()
            fk_parts.append(f"{tbl_lower}.{from_col} = {ref_table}.{to_col}")

    conn.close()

    schema = " | ".join(table_parts)
    if fk_parts:
        schema += " | foreign_keys: " + " , ".join(fk_parts)
    return schema


# ---------------------------------------------------------------------------
# CSV loader (unchanged from original, single-table WikiSQL format)
# ---------------------------------------------------------------------------

def load_csv_to_sqlite(file_obj_or_path) -> dict:
    """Load a CSV file into a single-table SQLite DB and return metadata."""
    df = pd.read_csv(file_obj_or_path)

    header = list(df.columns)
    types  = _infer_wikisql_types(df)
    df     = _lowercase_text_columns(df, types)

    conn = sqlite3.connect(str(_UPLOAD_DB_PATH))
    try:
        df.to_sql(TABLE_NAME, conn, index=False, if_exists="replace")
    finally:
        conn.close()

    schema_string = _wikisql_schema(header, types=types)

    preview_df = df.head(MAX_PREVIEW_ROWS).astype(object).where(df.head(MAX_PREVIEW_ROWS).notna(), None)

    return {
        "db_path":       str(_UPLOAD_DB_PATH),
        "schema_string": schema_string,
        "columns":       header,
        "row_count":     len(df),
        "rows":          preview_df.values.tolist(),
        "truncated":     len(df) > MAX_PREVIEW_ROWS,
    }


# ---------------------------------------------------------------------------
# SQLite / .db file loader (multi-table, Spider format)
# ---------------------------------------------------------------------------

def load_sqlite_file(file_obj) -> dict:
    """Copy an uploaded .sqlite/.db file to temp location and build schema."""
    # Write uploaded bytes to temp file
    with open(str(_UPLOAD_DB_PATH), "wb") as out:
        shutil.copyfileobj(file_obj, out)

    schema_string = _build_spider_schema_from_sqlite(str(_UPLOAD_DB_PATH))
    preview       = _preview_from_db(str(_UPLOAD_DB_PATH))

    return {
        "db_path":       str(_UPLOAD_DB_PATH),
        "schema_string": schema_string,
        **preview,
    }


# ---------------------------------------------------------------------------
# Excel loader (each sheet = one table, Spider format)
# ---------------------------------------------------------------------------

def load_excel_to_sqlite(file_obj) -> dict:
    """Load an Excel workbook: each sheet becomes a table in the temp SQLite DB."""
    xl = pd.ExcelFile(file_obj)
    sheet_names = xl.sheet_names

    conn = sqlite3.connect(str(_UPLOAD_DB_PATH))
    try:
        for sheet in sheet_names:
            df = xl.parse(sheet)
            # Sanitize table name: lowercase, spaces to underscores
            tbl_name = sheet.lower().replace(" ", "_")
            df.to_sql(tbl_name, conn, index=False, if_exists="replace")
        conn.commit()
    finally:
        conn.close()

    schema_string = _build_spider_schema_from_sqlite(str(_UPLOAD_DB_PATH))
    preview       = _preview_from_db(str(_UPLOAD_DB_PATH))

    return {
        "db_path":       str(_UPLOAD_DB_PATH),
        "schema_string": schema_string,
        **preview,
    }


# ---------------------------------------------------------------------------
# Model input builder
# ---------------------------------------------------------------------------

def build_input(question: str, schema_string: str) -> str:
    """Build model input string from a pre-built schema string (any format)."""
    return f"question: {question.strip()} | schema: {schema_string}"