Blame for server/main.py

import asyncio
import csv
import io
import json
import os
import re
import shutil
import sqlite3
import tempfile
from typing import Any
from difflib import get_close_matches

from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Request
from fastapi.responses import HTMLResponse, JSONResponse
from pydantic import ValidationError

try:
    from dotenv import load_dotenv
except Exception:  # pragma: no cover - optional dependency guard
    load_dotenv = None

from .environment import (
    SQLQualityEnvironment,
    analyze_uploaded_database,
    first_sql_keyword,
    is_readonly_statement,
    quote_identifier,
    rewrite_known_quoted_identifiers,
    strip_sql_comments,
    split_sql_statements,
)
from .models import (
    HealthResponse,
    ResetRequest,
    SQLAction,
    SQLObservation,
    SQLState,
    StepResponse,
)

if load_dotenv:
    load_dotenv()

from fastapi import APIRouter

app = FastAPI(
    title="SQL Quality Audit Agent Environment",
    description="OpenEnv-compatible environment for training AI agents on SQL data quality tasks",
    version="0.1.0",
)

# Create /v1/ prefixed router for OpenEnv standard compatibility
v1_router = APIRouter(prefix="/v1")

# Module-level singleton — ONE instance shared across all requests
_env = SQLQualityEnvironment()

# Module-level custom database state (separate from RL environment)
_custom_db_path: str | None = None
_custom_db_schema: list[dict] = []
_custom_db_filename: str = ""
_custom_db_analysis: dict[str, Any] = {}

# NEW: RL task support on uploaded database
_uploaded_db_path: str | None = None
_uploaded_db_schema: dict = {}
_uploaded_db_tasks: list[str] = []  # Which tasks are compatible

# HTTP-level request lock to serialize concurrent requests
_request_lock = asyncio.Lock()

TASK_CATALOG = [
    {
        "id": "null_audit",
        "task_id": "null_audit",
        "name": "null_audit",
        "difficulty": "easy",
        "description": "Identify the column with null rate exceeding 50 percent",
        "max_attempts": 1,
        "grader": True,
        "has_grader": True,
        "scored": True,
        "grader_fn": "server.graders:grade_null_audit",
    },
    {
        "id": "type_repair",
        "task_id": "type_repair",
        "name": "type_repair",
        "difficulty": "medium",
        "description": "Fix numeric type mismatch in transaction records",
        "max_attempts": 3,
        "grader": True,
        "has_grader": True,
        "scored": True,
        "grader_fn": "server.graders:grade_type_repair",
    },
    {
        "id": "dedup_bug",
        "task_id": "dedup_bug",
        "name": "dedup_bug",
        "difficulty": "hard",
        "description": "Find and repair duplication bug inflating join volumes",
        "max_attempts": 5,
        "grader": True,
        "has_grader": True,
        "scored": True,
        "grader_fn": "server.graders:grade_dedup_bug",
    },
    {
        "id": "fk_integrity",
        "task_id": "fk_integrity",
        "name": "fk_integrity",
        "difficulty": "hard",
        "description": "Remove orphan child records violating referential integrity",
        "max_attempts": 4,
        "grader": True,
        "has_grader": True,
        "scored": True,
        "grader_fn": "server.graders:grade_fk_integrity",
    },
    {
        "id": "schema_drift",
        "task_id": "schema_drift",
        "name": "schema_drift",
        "difficulty": "medium",
        "description": "Normalize drifted schemas into canonical analytics artifact",
        "max_attempts": 4,
        "grader": True,
        "has_grader": True,
        "scored": True,
        "grader_fn": "server.graders:grade_schema_drift",
    },
]
TASK_IDS = tuple(task["id"] for task in TASK_CATALOG)
TASK_CATALOG_BY_ID = {task["id"]: task for task in TASK_CATALOG}
DEFAULT_TASK_ID = TASK_IDS[0]
DOCUMENTATION_URL = "https://huggingface.co/spaces/Deadman922/sql-quality-env"

JSONRPC_PARSE_ERROR = -32700
JSONRPC_INVALID_REQUEST = -32600
JSONRPC_METHOD_NOT_FOUND = -32601
JSONRPC_INVALID_PARAMS = -32602
JSONRPC_INTERNAL_ERROR = -32603


def open_sqlite_db(db_path: str) -> sqlite3.Connection:
    """Open a SQLite database with row access tuned for API usage."""
    conn = sqlite3.connect(db_path, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def replace_temp_db(existing_path: str | None, suffix: str = ".sqlite") -> str:
    """Create a fresh temp database file and remove the prior temp file when possible."""
    if existing_path and os.path.exists(existing_path):
        try:
            os.unlink(existing_path)
        except OSError:
            pass
    handle = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    handle.close()
    return handle.name


def schema_table_map(schema: list[dict]) -> dict[str, dict]:
    """Return schema rows indexed by table name."""
    return {table["table_name"]: table for table in schema}


def _first_non_empty(*values: str) -> str:
    """Return the first non-empty stripped string."""
    for value in values:
        candidate = str(value or "").strip()
        if candidate:
            return candidate
    return ""


def _resolve_llm_runtime() -> dict[str, Any]:
    """Resolve OpenAI-compatible runtime config from environment variables."""
    groq_key = _first_non_empty(os.getenv("GROQ_API_KEY", ""))
    if groq_key:
        return {
            "configured": True,
            "provider": "groq",
            "api_key": groq_key,
            "base_url": _first_non_empty(
                os.getenv("GROQ_BASE_URL", ""),
                os.getenv("API_BASE_URL", ""),
                "https://api.groq.com/openai/v1",
            ),
            "model": _first_non_empty(
                os.getenv("GROQ_MODEL", ""),
                os.getenv("MODEL_NAME", ""),
                "llama-3.3-70b-versatile",
            ),
            "missing": [],
        }

    hf_token = _first_non_empty(os.getenv("HF_TOKEN", ""))
    openai_key = _first_non_empty(os.getenv("OPENAI_API_KEY", ""), os.getenv("API_KEY", ""))
    provider = "huggingface" if hf_token else ("openai" if openai_key else "none")
    api_key = hf_token or openai_key

    if provider == "openai":
        base_url = _first_non_empty(
            os.getenv("OPENAI_BASE_URL", ""),
            os.getenv("API_BASE_URL", ""),
            "https://api.openai.com/v1",
        )
    else:
        base_url = _first_non_empty(os.getenv("API_BASE_URL", ""), os.getenv("OPENAI_BASE_URL", ""))

    model_name = _first_non_empty(os.getenv("MODEL_NAME", ""), os.getenv("OPENAI_MODEL", ""))

    missing: list[str] = []
    if not api_key:
        missing.append("HF_TOKEN or OPENAI_API_KEY")
    if not base_url:
        missing.append("API_BASE_URL")
    if not model_name:
        missing.append("MODEL_NAME")

    return {
        "configured": bool(api_key and base_url and model_name),
        "provider": provider,
        "api_key": api_key,
        "base_url": base_url,
        "model": model_name,
        "missing": missing,
    }


def _format_llm_missing_message(missing: list[str]) -> str:
    """Build a user-friendly missing-config message."""
    if not missing:
        return "LLM configuration incomplete."
    return f"Missing LLM configuration: {', '.join(missing)}"


def load_static_html() -> str:
    """Load the static HTML interface for the App section."""
    # Look for static/index.html in various locations
    possible_paths = [
        os.path.join(os.path.dirname(__file__), "..", "static", "index.html"),
        os.path.join(os.path.dirname(__file__), "static", "index.html"),
        "/app/static/index.html",
        "static/index.html",
    ]
    for path in possible_paths:
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                return f.read()
    # Fallback simple HTML if file not found
    return """<!DOCTYPE html>
<html>
<head><title>SQL Quality Environment</title></head>
<body>
    <h1>SQL Quality Audit Agent Environment</h1>
    <p>Environment is running. Use API endpoints directly:</p>
    <ul>
        <li>GET /health - Health check</li>
        <li>POST /reset - Reset task</li>
        <li>POST /step - Execute action</li>
        <li>GET /state - Get current state</li>
    </ul>
</body>
</html>"""


def _task_catalog_records() -> list[dict[str, Any]]:
    """Return shallow copies of task records for response payloads."""
    return [dict(task) for task in TASK_CATALOG]


def _valid_task_ids_display() -> list[str]:
    """Return task IDs in canonical order for errors and schemas."""
    return list(TASK_IDS)


def _validate_task_id_or_raise(task_id: str | None) -> str:
    """Resolve and validate task IDs against the canonical catalog."""
    resolved_task_id = task_id or DEFAULT_TASK_ID
    if resolved_task_id not in TASK_CATALOG_BY_ID:
        raise HTTPException(
            status_code=422,
            detail=(
                f"Invalid task_id '{resolved_task_id}'. "
                f"Must be one of: {_valid_task_ids_display()}"
            ),
        )
    return resolved_task_id


def _patch_task_id_field_schema(field_schema: dict[str, Any]) -> None:
    """Inject machine-readable task enums into a single schema field."""
    if "anyOf" in field_schema:
        for variant in field_schema.get("anyOf", []):
            if isinstance(variant, dict) and variant.get("type") == "string":
                variant["enum"] = _valid_task_ids_display()
    elif field_schema.get("type") == "string":
        field_schema["enum"] = _valid_task_ids_display()


def _patch_task_id_enums(schema: Any) -> None:
    """Recursively patch schema objects so task_id is an explicit enum."""
    if isinstance(schema, dict):
        properties = schema.get("properties")
        if isinstance(properties, dict) and isinstance(properties.get("task_id"), dict):
            _patch_task_id_field_schema(properties["task_id"])
        for value in schema.values():
            _patch_task_id_enums(value)
    elif isinstance(schema, list):
        for item in schema:
            _patch_task_id_enums(item)


def _patched_model_schema(model_cls: Any) -> dict[str, Any]:
    """Return a model JSON schema with task enums patched in."""
    schema = model_cls.model_json_schema()
    _patch_task_id_enums(schema)
    return schema


def _build_schema_payload() -> dict[str, Any]:
    """Build OpenEnv-style action/observation/state schema payload."""
    return {
        "action": _patched_model_schema(SQLAction),
        "observation": _patched_model_schema(SQLObservation),
        "state": _patched_model_schema(SQLState),
    }


def _build_tasks_payload() -> dict[str, Any]:
    """Build the task listing payload from the canonical catalog."""
    tasks = _task_catalog_records()
    return {
        "tasks": tasks,
        "total": len(tasks),
        "graded": sum(1 for task in tasks if task.get("has_grader")),
    }


def _jsonrpc_success(result: Any, request_id: Any = None) -> dict[str, Any]:
    """Build a JSON-RPC 2.0 success response."""
    return {
        "jsonrpc": "2.0",
        "result": result,
        "id": request_id,
    }


def _jsonrpc_error(
    code: int,
    message: str,
    request_id: Any = None,
    data: Any = None,
) -> dict[str, Any]:
    """Build a JSON-RPC 2.0 error response."""
    error: dict[str, Any] = {
        "code": code,
        "message": message,
    }
    if data is not None:
        error["data"] = data
    return {
        "jsonrpc": "2.0",
        "error": error,
        "id": request_id,
    }


def _mcp_tool_specs() -> list[dict[str, Any]]:
    """Return MCP tool specs for task discovery and episode control."""
    reset_schema = ResetRequest.model_json_schema()
    _patch_task_id_enums(reset_schema)
    action_schema = _patched_model_schema(SQLAction)
    return [
        {
            "name": "list_tasks",
            "description": "Return the registered task catalog for this environment.",
            "inputSchema": {
                "type": "object",
                "properties": {},
                "additionalProperties": False,
            },
        },
        {
            "name": "episode_reset",
            "description": "Reset the environment for a specific task and return the initial observation.",
            "inputSchema": reset_schema,
        },
        {
            "name": "episode_step",
            "description": "Execute one environment step using the SQL action payload.",
            "inputSchema": action_schema,
        },
        {
            "name": "episode_state",
            "description": "Return the current state for the active environment episode.",
            "inputSchema": {
                "type": "object",
                "properties": {},
                "additionalProperties": False,
            },
        },
    ]


def _schema_rows_from_observation(obs: SQLObservation) -> list[dict[str, Any]]:
    """Convert observation schema info to frontend-friendly rows."""
    schema_rows: list[dict[str, Any]] = []
    if obs.schema_info:
        for table_info in obs.schema_info:
            schema_rows.append(
                {
                    "table_name": table_info.table_name,
                    "columns": table_info.columns,
                    "row_count": table_info.row_count,
                }
            )
    return schema_rows


async def _reset_observation(request: ResetRequest | None = None) -> SQLObservation:
    """Shared reset helper used by HTTP routes and MCP tools."""
    global _uploaded_db_path, _uploaded_db_tasks, _custom_db_analysis

    if request is None:
        request = ResetRequest()

    async with _request_lock:
        task_id = _validate_task_id_or_raise(request.task_id)
        conn = None
        try:
            if _uploaded_db_path is not None and task_id in _uploaded_db_tasks:
                conn = open_sqlite_db(_uploaded_db_path)
                return _env.reset_with_connection(
                    task_id,
                    conn,
                    _uploaded_db_schema,
                    _custom_db_analysis,
                )
            if _uploaded_db_path is not None and _custom_db_analysis.get("task_candidates", {}).get(task_id):
                reason = _custom_db_analysis["task_candidates"][task_id].get(
                    "reason",
                    f"Task '{task_id}' is not supported for the uploaded dataset.",
                )
                raise ValueError(reason)
            return _env.reset(task_id, seed=request.seed)
        except ValueError as e:
            raise HTTPException(status_code=422, detail=str(e)) from e
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e)) from e
        finally:
            if conn is not None:
                conn.close()


async def _execute_step(action: SQLAction) -> StepResponse:
    """Shared step helper used by HTTP routes and MCP tools."""
    async with _request_lock:
        try:
            return _env.step(action)
        except HTTPException:
            raise
        except ValueError as e:
            raise HTTPException(status_code=422, detail=str(e)) from e
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e)) from e


async def _current_state() -> SQLState:
    """Shared state helper used by HTTP routes and MCP tools."""
    try:
        return _env.state()
    except ValueError as e:
        raise HTTPException(status_code=422, detail=str(e)) from e


_original_openapi = app.openapi


def _custom_openapi() -> dict[str, Any]:
    """Patch OpenAPI component schemas so task_id becomes a real enum."""
    if app.openapi_schema:
        return app.openapi_schema

    openapi_schema = _original_openapi()
    if isinstance(openapi_schema, dict):
        components = openapi_schema.get("components", {}).get("schemas", {})
        for schema_name in ("ResetRequest", "SQLAction", "SQLObservation", "SQLState"):
            schema = components.get(schema_name)
            if isinstance(schema, dict):
                _patch_task_id_enums(schema)
        app.openapi_schema = openapi_schema
    return openapi_schema


app.openapi = _custom_openapi


# Validate LLM configuration at startup
runtime = _resolve_llm_runtime()
LLM_AVAILABLE = runtime["configured"]
if not LLM_AVAILABLE:
    missing_vars = ", ".join(runtime["missing"])
    print(f"[WARNING] LLM env vars not set: {missing_vars}. Server starting in mock mode.", flush=True)


@app.get("/", response_class=HTMLResponse)
async def root():
    """Root endpoint for HF Space App section - provides interactive testing interface."""
    return load_static_html()


@app.get("/health")
async def health() -> dict:
    """Health check endpoint for monitoring and load balancers.

    Returns:
        dict: Health status with server info
    """
    runtime = _resolve_llm_runtime()
    mock_mode = not bool(runtime["configured"])
    llm_available = bool(runtime["configured"])
    
    if mock_mode:
        message = _format_llm_missing_message(runtime["missing"])
    else:
        message = runtime["provider"]

    return {
        "status": "healthy",
        "llm": "mock_mode" if mock_mode else "connected",
        "llm_available": llm_available,
        "mock_mode": mock_mode,
        "provider": runtime["provider"],
        "message": message,
    }


@app.get("/metadata")
async def get_metadata() -> dict:
    """Get environment metadata for openenv-core.

    Returns:
        dict: Environment metadata including all task IDs with graders
    """
    tasks = _task_catalog_records()
    return {
        "name": "sql-quality-env",
        "version": "0.1.0",
        "description": "A multi-step SQL data quality environment for training and evaluating data engineering agents",
        "author": "Tharak",
        "documentation": DOCUMENTATION_URL,
        "documentation_url": DOCUMENTATION_URL,
        "tasks": tasks,
        "task_count": len(tasks),
        "graded_count": sum(1 for t in tasks if t.get("has_grader")),
    }


@app.get("/tasks")
async def list_tasks():
    return _build_tasks_payload()


@app.get("/schema")
async def get_schema() -> dict:
    """Return action, observation, and state schemas for OpenEnv clients."""
    return _build_schema_payload()


@app.post("/mcp")
async def mcp_endpoint(request_raw: Request) -> dict[str, Any]:
    """Minimal MCP JSON-RPC surface for task discovery and episode control."""
    try:
        body = await request_raw.body()
        request_dict = json.loads(body or b"{}")
    except json.JSONDecodeError:
        return _jsonrpc_error(JSONRPC_PARSE_ERROR, "Parse error")

    if not isinstance(request_dict, dict):
        return _jsonrpc_error(JSONRPC_INVALID_REQUEST, "Invalid Request")

    request_id = request_dict.get("id")
    if request_dict.get("jsonrpc") != "2.0" or not isinstance(request_dict.get("method"), str):
        return _jsonrpc_error(JSONRPC_INVALID_REQUEST, "Invalid Request", request_id=request_id)

    method = request_dict["method"]
    params = request_dict.get("params", {})
    if params is None:
        params = {}
    if not isinstance(params, dict):
        return _jsonrpc_error(
            JSONRPC_INVALID_PARAMS,
            "Invalid params: 'params' must be an object",
            request_id=request_id,
        )

    if method == "tools/list":
        return _jsonrpc_success({"tools": _mcp_tool_specs()}, request_id=request_id)

    if method != "tools/call":
        return _jsonrpc_error(
            JSONRPC_METHOD_NOT_FOUND,
            f"Method not found: {method}",
            request_id=request_id,
        )

    tool_name = params.get("name")
    arguments = params.get("arguments", {})
    if arguments is None:
        arguments = {}
    if not tool_name:
        return _jsonrpc_error(
            JSONRPC_INVALID_PARAMS,
            "Missing 'name' in params",
            request_id=request_id,
        )
    if not isinstance(arguments, dict):
        return _jsonrpc_error(
            JSONRPC_INVALID_PARAMS,
            "Invalid params: 'arguments' must be an object",
            request_id=request_id,
        )

    try:
        if tool_name == "list_tasks":
            result: Any = _build_tasks_payload()
        elif tool_name == "episode_reset":
            reset_request = ResetRequest.model_validate(arguments)
            result = (await _reset_observation(reset_request)).model_dump()
        elif tool_name == "episode_step":
            action = SQLAction.model_validate(arguments)
            result = (await _execute_step(action)).model_dump()
        elif tool_name == "episode_state":
            result = (await _current_state()).model_dump()
        else:
            return _jsonrpc_error(
                JSONRPC_INVALID_PARAMS,
                f"Tool not found: {tool_name}",
                request_id=request_id,
            )
    except ValidationError as e:
        return _jsonrpc_error(
            JSONRPC_INVALID_PARAMS,
            "Invalid params",
            request_id=request_id,
            data=e.errors(),
        )
    except HTTPException as e:
        code = JSONRPC_INVALID_PARAMS if e.status_code == 422 else JSONRPC_INTERNAL_ERROR
        return _jsonrpc_error(code, str(e.detail), request_id=request_id)
    except Exception as e:
        return _jsonrpc_error(JSONRPC_INTERNAL_ERROR, str(e), request_id=request_id)

    return _jsonrpc_success(result, request_id=request_id)


@app.get("/status")
async def status() -> dict:
    """Status endpoint with LLM heartbeat check for Hugging Face Spaces.
    
    Performs an actual LLM connection test by checking environment variables
    and optionally making a lightweight API call.
    
    Returns:
        dict: Status response with llm_connected boolean
    """
    runtime = _resolve_llm_runtime()
    mock_mode = not bool(runtime["configured"])
    llm_connected = bool(runtime["configured"])
    llm_error = None

    if llm_connected:
        try:
            import httpx

            headers = {"Authorization": f"Bearer {runtime['api_key']}"}
            probe_url = f"{runtime['base_url'].rstrip('/')}/models"
            async with httpx.AsyncClient() as client:
                response = await client.get(probe_url, headers=headers, timeout=5)
            if response.status_code == 200:
                llm_connected = True
            elif response.status_code in {401, 403}:
                llm_connected = False
                llm_error = f"LLM auth failed (HTTP {response.status_code})"
            else:
                # Some providers do not expose /models; keep configured-as-connected.
                llm_connected = True
                llm_error = f"LLM heartbeat warning (HTTP {response.status_code})"
        except Exception as e:
            # Keep configured-as-connected to avoid false negatives from transient network errors.
            llm_connected = True
            llm_error = f"LLM heartbeat skipped: {e}"
    else:
        llm_error = _format_llm_missing_message(runtime["missing"])

    return {
        "status": "online",
        "llm": "mock_mode" if mock_mode else ("connected" if llm_connected else "disconnected"),
        "llm_connected": llm_connected,
        "llm_error": llm_error,
        "mock_mode": mock_mode,
        "provider": runtime["provider"],
    }


@app.post("/ping", response_model=HealthResponse)
async def ping() -> HealthResponse:
    """Ping endpoint for submission validator — accepts empty body.

    Returns:
        dict: Status response indicating the server is responsive
    """
    return HealthResponse(status="ok")


@app.post("/initialize")
async def initialize(request: ResetRequest) -> dict:
    """Initialize episode and return full state including schema for UI.
    
    This endpoint combines reset() with schema extraction for the frontend.
    
    Args:
        request: ResetRequest containing task_id and optional seed
        
    Returns:
        dict: Full episode state with schema_info for Table Schema panel
    """
    obs = await _reset_observation(request)
    return {
        "status": "ok",
        "episode_id": obs.episode_id,
        "task_id": obs.task_id,
        "task_description": obs.task_description,
        "schema": _schema_rows_from_observation(obs),
        "sample_rows": obs.sample_rows,
        "attempt_number": obs.attempt_number,
        "max_attempts": obs.max_attempts,
        "investigation_hint": obs.investigation_hint,
        "task_context": _env.task_context(),
    }


@app.post("/reset", response_model=SQLObservation)
async def reset(request: ResetRequest = None) -> SQLObservation:
    """Reset environment and return initial observation for specified task.

    Args:
        request: ResetRequest containing task_id and optional deterministic seed

    Returns:
        SQLObservation: Initial state of database for requested task

    Raises:
        HTTPException 422: If task_id is not a valid task identifier
        HTTPException 500: If environment initialization fails
    """
    return await _reset_observation(request)


@app.post("/step", response_model=StepResponse)
async def step(action: SQLAction) -> StepResponse:
    """Submit an action and receive observation, reward, done flag, and info.

    Args:
        action: SQLAction containing task_id and fix_query or answer

    Returns:
        StepResponse: Updated observation, reward, done status, and info

    Raises:
        HTTPException 422: If action is invalid or episode is already done
        HTTPException 500: If environment step execution fails
    """
    return await _execute_step(action)


@app.get("/state", response_model=SQLState)
async def state() -> SQLState:
    """Get current episode state.

    Returns:
        SQLState: Current episode information including task_id, episode_id, attempt_number, cumulative_reward, and done status

    Raises:
        HTTPException 422: If no active episode exists
    """
    return await _current_state()


# ─── V1 PREFIXED ENDPOINTS FOR OPENENV STANDARD ─────────────────────────

@v1_router.get("/health")
async def v1_health() -> dict:
    """Health check endpoint at /v1/health for OpenEnv compatibility."""
    return await health()

@v1_router.get("/tasks")
async def v1_list_tasks():
    """List all graded tasks at /v1/tasks for OpenEnv compatibility."""
    return _build_tasks_payload()

@v1_router.get("/metadata")
async def v1_get_metadata() -> dict:
    """Get environment metadata at /v1/metadata for OpenEnv compatibility."""
    return await get_metadata()

@v1_router.get("/schema")
async def v1_get_schema() -> dict:
    """Return action, observation, and state schemas at /v1/schema."""
    return _build_schema_payload()

@v1_router.post("/reset", response_model=SQLObservation)
async def v1_reset(request: ResetRequest) -> SQLObservation:
    """Reset environment at /v1/reset for OpenEnv compatibility."""
    return await reset(request)

@v1_router.post("/step", response_model=StepResponse)
async def v1_step(action: SQLAction) -> StepResponse:
    """Execute action at /v1/step for OpenEnv compatibility."""
    return await step(action)

@v1_router.get("/state", response_model=SQLState)
async def v1_state() -> SQLState:
    """Get current state at /v1/state for OpenEnv compatibility."""
    return await state()

# Include the v1 router in the app
app.include_router(v1_router)


# ─── CUSTOM DATABASE ENDPOINTS ─────────────────────────────────────────

def infer_sqlite_type(values: list[str]) -> str:
    """Infer SQLite column type from a list of string values with smart casting."""
    non_empty = [v for v in values if v and v.strip()]
    if not non_empty:
        return 'TEXT'

    strict_numeric_pattern = re.compile(r'^[+-]?\d+(?:\.\d+)?$')
    
    # Try INTEGER first (with cleaning)
    int_count = 0
    for v in non_empty:
        try:
            raw = str(v).strip()
            if strict_numeric_pattern.match(raw):
                int(raw)
                int_count += 1
        except (ValueError, TypeError):
            pass
    
    # If 80% can be integers, use INTEGER
    if int_count / len(non_empty) >= 0.8:
        return 'INTEGER'
    
    # Try REAL/float
    float_count = 0
    for v in non_empty:
        try:
            raw = str(v).strip()
            if strict_numeric_pattern.match(raw):
                float(raw)
                float_count += 1
        except (ValueError, TypeError):
            pass
    
    if float_count / len(non_empty) >= 0.8:
        return 'REAL'
    
    return 'TEXT'


def sanitize_identifier(name: str) -> str:
    """Sanitize table/column names to prevent SQL reserved word collisions."""
    # Replace invalid characters
    safe = re.sub(r'[^a-zA-Z0-9_]', '_', name)
    # Ensure starts with letter or underscore
    if safe and not safe[0].isalpha() and safe[0] != '_':
        safe = '_' + safe
    return safe


def quote_identifier(name: str) -> str:
    """Quote SQL identifier to handle reserved words and special characters."""
    safe = sanitize_identifier(name)
    return f'"{safe}"'


def perform_health_audit(conn: sqlite3.Connection, table_name: str) -> dict:
    """Perform automatic health audit on uploaded table.
    
    Returns health metrics including:
    - Total row count
    - Column-level null/empty statistics
    - Data type distribution
    - Dirty data indicators (NaN, -1, empty strings)
    """
    audit = {
        "table_name": table_name,
        "total_rows": 0,
        "columns": {},
        "health_score": 100,
        "issues": []
    }
    
    try:
        # Get row count
        cursor = conn.execute(f'SELECT COUNT(*) FROM {quote_identifier(table_name)}')
        audit["total_rows"] = cursor.fetchone()[0]
        
        if audit["total_rows"] == 0:
            audit["health_score"] = 0
            audit["issues"].append("Table is empty")
            return audit
        
        # Get column info
        pragma_cursor = conn.execute(f'PRAGMA table_info({quote_identifier(table_name)})')
        columns = [(row[1], row[2]) for row in pragma_cursor.fetchall()]  # (name, type)
        
        total_issues = 0
        
        for col_name, col_type in columns:
            col_audit = {
                "type": col_type,
                "null_count": 0,
                "empty_string_count": 0,
                "nan_count": 0,
                "negative_one_count": 0,
                "dirty_data_count": 0,
                "health_percentage": 100.0
            }
            
            # Count NULLs
            try:
                cursor = conn.execute(
                    f'SELECT COUNT(*) FROM {quote_identifier(table_name)} WHERE {quote_identifier(col_name)} IS NULL'
                )
                col_audit["null_count"] = cursor.fetchone()[0]
            except:
                pass
            
            # Count empty strings
            try:
                cursor = conn.execute(
                    f'''SELECT COUNT(*) FROM {quote_identifier(table_name)} 
                    WHERE {quote_identifier(col_name)} = '' OR {quote_identifier(col_name)} = ' ' 
                    OR TRIM(CAST({quote_identifier(col_name)} AS TEXT)) = '''
                )
                col_audit["empty_string_count"] = cursor.fetchone()[0]
            except:
                pass
            
            # Count 'NaN', 'N/A', 'null' strings
            try:
                cursor = conn.execute(
                    f'''SELECT COUNT(*) FROM {quote_identifier(table_name)} 
                    WHERE UPPER(CAST({quote_identifier(col_name)} AS TEXT)) IN ('NAN', 'N/A', 'NULL', 'NONE', '')'''
                )
                col_audit["nan_count"] = cursor.fetchone()[0]
            except:
                pass
            
            # Count -1 values (common placeholder for missing data)
            try:
                cursor = conn.execute(
                    f'''SELECT COUNT(*) FROM {quote_identifier(table_name)} 
                    WHERE CAST({quote_identifier(col_name)} AS TEXT) = '-1' 
                    OR {quote_identifier(col_name)} = -1'''
                )
                col_audit["negative_one_count"] = cursor.fetchone()[0]
            except:
                pass
            
            # Calculate total dirty data
            col_audit["dirty_data_count"] = (
                col_audit["null_count"] + 
                col_audit["empty_string_count"] + 
                col_audit["nan_count"] +
                col_audit["negative_one_count"]
            )
            
            # Calculate health percentage
            if audit["total_rows"] > 0:
                clean_rows = audit["total_rows"] - col_audit["dirty_data_count"]
                col_audit["health_percentage"] = round((clean_rows / audit["total_rows"]) * 100, 1)
            
            audit["columns"][col_name] = col_audit
            total_issues += col_audit["dirty_data_count"]
        
        # Calculate overall health score
        if audit["total_rows"] > 0 and len(columns) > 0:
            total_cells = audit["total_rows"] * len(columns)
            audit["health_score"] = round(((total_cells - total_issues) / total_cells) * 100, 1)
        
        # Identify worst columns
        sorted_cols = sorted(
            audit["columns"].items(), 
            key=lambda x: x[1]["health_percentage"]
        )
        
        if sorted_cols and sorted_cols[0][1]["health_percentage"] < 50:
            audit["issues"].append(
                f"Column '{sorted_cols[0][0]}' has {sorted_cols[0][1]['health_percentage']:.1f}% clean data"
            )
        
    except Exception as e:
        audit["issues"].append(f"Audit error: {str(e)}")
        audit["health_score"] = 0
    
    return audit


def build_safe_null_check(column: str) -> str:
    """Build a comprehensive NULL/dirty data check for a column.
    
    Checks for: NULL, '', ' ', 'NaN', 'N/A', 'NULL', 'NONE', -1
    """
    col = quote_identifier(column)
    checks = [
        f"{col} IS NULL",
        f"TRIM(CAST({col} AS TEXT)) = ''",
        f"UPPER(CAST({col} AS TEXT)) IN ('NAN', 'N/A', 'NULL', 'NONE')",
        f"CAST({col} AS TEXT) = '-1'"
    ]
    return " OR ".join(checks)


def suggest_query_alternatives(conn: sqlite3.Connection, sql: str, error_msg: str) -> list[str]:
    """Suggest alternative queries when a query fails or returns 0 rows."""
    suggestions = []

    sql_clean = strip_sql_comments(sql)
    sql_upper = sql_clean.upper()
    schema = get_db_schema(conn)
    table_names = [table["table_name"] for table in schema]
    column_names = []
    for table in schema:
        for column in table.get("columns", []):
            column_names.append(column.split(" (")[0])

    missing_table = re.search(r"no such table: ([^\s]+)", error_msg, re.IGNORECASE)
    if missing_table:
        bad_table = missing_table.group(1).strip('"')
        close = get_close_matches(bad_table, table_names, n=3, cutoff=0.4)
        if close:
            suggestions.extend([f'Try table "{name}" instead.' for name in close])

    missing_column = re.search(r"no such column: ([^\s]+)", error_msg, re.IGNORECASE)
    if missing_column:
        bad_column = missing_column.group(1).strip('"')
        close = get_close_matches(bad_column, column_names, n=3, cutoff=0.4)
        if close:
            suggestions.extend([f'Try column "{name}" instead.' for name in close])

    if "0 rows" in error_msg.lower() or "no rows" in error_msg.lower():
        match = re.search(r'FROM\s+("?[\w_]+"\b|[\w_]+)', sql_clean, re.IGNORECASE)
        if match:
            table = match.group(1).strip('"')
            suggestions.append(f'SELECT COUNT(*) FROM "{table}"')
            suggestions.append(f'SELECT * FROM "{table}" LIMIT 5')

    if "IS NULL" in sql_upper and ("0 rows" in error_msg.lower() or "no rows" in error_msg.lower()):
        col_match = re.search(r'WHERE\s+("?[\w_]+"\b|[\w_]+)\s+IS\s+NULL', sql_clean, re.IGNORECASE)
        from_match = re.search(r'FROM\s+("?[\w_]+"\b|[\w_]+)', sql_clean, re.IGNORECASE)
        if col_match and from_match:
            col = col_match.group(1).strip('"')
            table = from_match.group(1).strip('"')
            suggestions.append(
                f"SELECT COUNT(*) FROM {quote_identifier(table)} WHERE "
                f"{quote_identifier(col)} IS NULL OR TRIM(CAST({quote_identifier(col)} AS TEXT)) = '' "
                f"OR LOWER(TRIM(CAST({quote_identifier(col)} AS TEXT))) IN ('n/a','na','nan','null','none','missing','unknown','-1')"
            )

    if "syntax error" in error_msg.lower():
        suggestions.append("Try running EXPLAIN QUERY PLAN before the query to validate complex SQL.")

    return suggestions if suggestions else ["Try: SELECT * FROM table LIMIT 5 to explore data"]


def extract_sql_table_names(sql: str) -> list[str]:
    """Return table names referenced in FROM/JOIN clauses."""
    sql_clean = strip_sql_comments(sql)
    return [
        match.strip('"')
        for match in re.findall(r"\b(?:FROM|JOIN)\s+(\"?[\w_]+\"?)", sql_clean, re.IGNORECASE)
    ]


def build_query_quality_warnings(sql: str, analysis: dict[str, Any] | None) -> list[str]:
    """Generate source-aware warnings for analytically risky SQL."""
    if not analysis:
        return []

    sql_clean = strip_sql_comments(sql)
    sql_upper = sql_clean.upper()
    warnings: list[str] = []
    table_profiles = analysis.get("table_profiles", {})

    for table_name, table_profile in table_profiles.items():
        for column_name, profile in table_profile.get("column_profiles", {}).items():
            quoted_column = f'"{column_name}"'
            bare_column = re.escape(column_name)
            quoted_pattern = re.escape(quoted_column)
            aggregate_patterns = [
                rf"\bAVG\s*\(\s*(?:{quoted_pattern}|\b{bare_column}\b)\s*\)",
                rf"\bSUM\s*\(\s*(?:{quoted_pattern}|\b{bare_column}\b)\s*\)",
                rf"\bMIN\s*\(\s*(?:{quoted_pattern}|\b{bare_column}\b)\s*\)",
                rf"\bMAX\s*\(\s*(?:{quoted_pattern}|\b{bare_column}\b)\s*\)",
            ]
            aggregate_hit = any(re.search(pattern, sql_clean, re.IGNORECASE) for pattern in aggregate_patterns)

            if profile.get("header_like_count", 0) > 0 and (
                re.search(quoted_pattern, sql_clean, re.IGNORECASE)
                or re.search(rf"\b{bare_column}\b", sql_clean, re.IGNORECASE)
            ):
                warnings.append(
                    f'Column "{column_name}" contains {profile["header_like_count"]} header-like value(s); the CSV may include repeated header rows or malformed values.'
                )

            declared_type = str(profile.get("declared_type", "") or "").upper()
            suspicious_name = any(token in column_name.lower() for token in ("salary", "revenue", "price", "amount", "cost", "income"))
            numeric_ratio = float(profile.get("numeric_ratio", 0.0) or 0.0)
            if aggregate_hit and (numeric_ratio >= 0.7 or suspicious_name or declared_type in {"", "TEXT"}):
                if "CAST(" not in sql_upper and "REPLACE(" not in sql_upper:
                    warnings.append(
                        f'Column "{column_name}" looks numeric-like text. Direct aggregates may coerce dirty values to 0; clean/cast the values first.'
                    )

    return warnings[:4]


def build_dataset_task_suggestions(sql: str, analysis: dict[str, Any] | None) -> list[str]:
    """Prefer backend task candidates over result-grid heuristics."""
    if not analysis:
        return []

    sql_clean = strip_sql_comments(sql)
    referenced_tables = set(extract_sql_table_names(sql_clean))
    candidates = []
    for task_id, candidate in analysis.get("task_candidates", {}).items():
        if not candidate.get("supported"):
            continue
        related_tables = set(candidate.get("related_tables") or ([candidate.get("table")] if candidate.get("table") else []))
        relevance = 1.0 if not referenced_tables else (1.0 if related_tables & referenced_tables else 0.4)
        candidates.append(
            (
                relevance * float(candidate.get("confidence", 0.0) or 0.0),
                task_id,
                candidate,
            )
        )

    suggestions = []
    for _, task_id, candidate in sorted(candidates, key=lambda item: item[0], reverse=True)[:3]:
        hint = candidate.get("context_hint") or candidate.get("description") or task_id
        suggestions.append(f'{task_id}: {hint}')
    return suggestions


def build_source_data_insights(sql: str, analysis: dict[str, Any] | None) -> list[str]:
    """Summarize the most relevant source-data quality signals."""
    if not analysis:
        return []

    insights: list[str] = []
    warnings = build_query_quality_warnings(sql, analysis)
    insights.extend(warnings)

    task_candidates = analysis.get("task_candidates", {})
    for task_id in ("type_repair", "null_audit", "dedup_bug", "fk_integrity", "schema_drift"):
        candidate = task_candidates.get(task_id, {})
        if not candidate.get("supported"):
            continue
        insights.append(candidate.get("context_hint") or candidate.get("description") or task_id)
        if len(insights) >= 4:
            break

    return insights[:6]


def is_aggregate_like_column_name(column_name: str) -> bool:
    """Heuristic for result columns that are derived aggregates or ranks."""
    tokens = [token for token in re.split(r"[^a-z0-9]+", (column_name or "").lower()) if token]
    aggregate_tokens = {"avg", "sum", "count", "total", "min", "max", "rank", "highest", "lowest"}
    return bool(tokens and any(token in aggregate_tokens for token in tokens))


def is_key_like_column_name(column_name: str) -> bool:
    """Heuristic for key-like result columns where duplicate values are suspicious."""
    tokens = [token for token in re.split(r"[^a-z0-9]+", (column_name or "").lower()) if token]
    key_tokens = {"id", "email", "phone", "ssn", "uuid", "employee", "patient"}
    return bool(tokens and any(token in key_tokens for token in tokens))


def get_db_schema(conn: sqlite3.Connection) -> list[dict]:
    """Extract schema information from a SQLite connection."""
    schema = []
    cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = cursor.fetchall()
    
    for (table_name,) in tables:
        # Get row count
        count_cursor = conn.execute(f'SELECT COUNT(*) FROM "{table_name}"')
        row_count = count_cursor.fetchone()[0]
        
        # Get columns
        pragma_cursor = conn.execute(f'PRAGMA table_info("{table_name}")')
        columns = []
        for row in pragma_cursor.fetchall():
            col_name = row[1]
            col_type = row[2]
            columns.append(f"{col_name} ({col_type})")
        
        schema.append({
            "table_name": table_name,
            "row_count": row_count,
            "columns": columns
        })
    
    return schema


import logging
import json
from datetime import datetime
from pathlib import Path

# Configure confidence logging
confidence_logger = logging.getLogger('confidence_scoring')
confidence_logger.setLevel(logging.INFO)

# Create logs directory if it doesn't exist
log_dir = Path('logs')
log_dir.mkdir(exist_ok=True)

# File handler for confidence scoring
confidence_handler = logging.FileHandler(log_dir / 'confidence_scoring.log')
confidence_handler.setLevel(logging.INFO)

# Formatter for detailed logging
confidence_formatter = logging.Formatter('%(asctime)s - %(message)s')
confidence_handler.setFormatter(confidence_formatter)
confidence_logger.addHandler(confidence_handler)

import uuid
from typing import Dict

# Session-based table alias mapping for manual task override
table_aliases: Dict[str, Dict[str, str]] = {}  # {session_id: {actual_table: alias_as}}

# Session management
active_sessions: Dict[str, dict] = {}  # {session_id: {created_at, last_used}}


def log_confidence_calculation(table_name: str, task_id: str, confidence: float, 
                              reasoning: str, fingerprint: dict, session_id: str = None) -> dict:
    """Log confidence score calculations for judge review and UI display."""
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "session_id": session_id or "default",
        "table_name": table_name,
        "task_id": task_id,
        "confidence_score": round(confidence, 3),
        "reasoning": reasoning,
        "fingerprint_analysis": {
            "row_count": fingerprint.get('row_count', 0),
            "column_count": fingerprint.get('column_count', 0),
            "null_counts": fingerprint.get('null_counts', {}),
            "semantic_hints": fingerprint.get('semantic_hints', []),
            "data_types": fingerprint.get('data_types', {}),
            "total_nulls": sum(fingerprint.get('null_counts', {}).values()),
            "null_percentage": round(
                sum(fingerprint.get('null_counts', {}).values()) / max(1, fingerprint.get('row_count', 1)) * 100, 2
            )
        },
        "calculation_breakdown": {
            "base_confidence": confidence,
            "adjustments": [],
            "final_score": round(confidence, 3)
        }
    }
    
    # Log to file for judge review
    confidence_logger.info(f"CONFIDENCE_CALCULATION: {json.dumps(log_entry, indent=2)}")
    
    return log_entry


def sanitize_data_value(value: str) -> any:
    """Sanitize dirty data values during upload."""
    if value is None:
        return None
    
    value_str = str(value).strip()
    
    # Convert common dirty data patterns to NULL
    if (value_str == '' or 
        value_str == ' ' or 
        value_str.upper() in ['NAN', 'N/A', 'NULL', 'NONE', 'NA', '#N/A', '-1']):
        return None

    return value_str


def analyze_table_fingerprint(conn: sqlite3.Connection, table_name: str) -> dict:
    """Analyze table structure and content for semantic task mapping."""
    fingerprint = {
        "table_name": table_name,
        "column_count": 0,
        "row_count": 0,
        "columns": [],
        "null_counts": {},
        "data_types": {},
        "semantic_hints": []
    }
    
    try:
        # Get basic info
        cursor = conn.execute(f'SELECT COUNT(*) FROM {quote_identifier(table_name)}')
        fingerprint["row_count"] = cursor.fetchone()[0]
        
        # Get column info
        pragma_cursor = conn.execute(f'PRAGMA table_info({quote_identifier(table_name)})')
        columns = [(row[1], row[2]) for row in pragma_cursor.fetchall()]
        fingerprint["column_count"] = len(columns)
        
        for col_name, col_type in columns:
            fingerprint["columns"].append({"name": col_name, "type": col_type})
            
            # Count nulls
            try:
                null_cursor = conn.execute(
                    f'SELECT COUNT(*) FROM {quote_identifier(table_name)} WHERE {quote_identifier(col_name)} IS NULL'
                )
                null_count = null_cursor.fetchone()[0]
                fingerprint["null_counts"][col_name] = null_count
            except:
                fingerprint["null_counts"][col_name] = 0
            
            # Sample data for type analysis
            try:
                sample_cursor = conn.execute(
                    f'SELECT {quote_identifier(col_name)} FROM {quote_identifier(table_name)} WHERE {quote_identifier(col_name)} IS NOT NULL LIMIT 10'
                )
                samples = [row[0] for row in sample_cursor.fetchall()]
                fingerprint["data_types"][col_name] = infer_sqlite_type([str(s) for s in samples])
            except:
                fingerprint["data_types"][col_name] = 'TEXT'
        
        # Generate semantic hints
        col_names_lower = [col["name"].lower() for col in fingerprint["columns"]]
        
        # Check for user-like patterns
        user_patterns = ['id', 'name', 'email', 'age', 'city', 'country', 'address', 'phone']
        if any(pattern in col_names_lower for pattern in user_patterns):
            fingerprint["semantic_hints"].append("user_data")
        
        # Check for transaction-like patterns
        trans_patterns = ['amount', 'price', 'total', 'cost', 'quantity', 'order', 'transaction', 'date']
        if any(pattern in col_names_lower for pattern in trans_patterns):
            fingerprint["semantic_hints"].append("transaction_data")
        
        # Check for job/posting patterns
        job_patterns = ['job', 'title', 'company', 'salary', 'rating', 'description', 'skills']
        if any(pattern in col_names_lower for pattern in job_patterns):
            fingerprint["semantic_hints"].append("job_market_data")
        
        # Check for event patterns
        event_patterns = ['event', 'timestamp', 'type', 'category', 'action']
        if any(pattern in col_names_lower for pattern in event_patterns):
            fingerprint["semantic_hints"].append("event_data")
            
    except Exception as e:
        fingerprint["error"] = str(e)
    
    return fingerprint


def detect_tasks_semantically(schema: list[dict], fingerprints: dict[str, dict], session_id: str = None) -> list[dict]:
    """Detect compatible tasks using semantic analysis instead of hard-coded names."""
    task_suggestions = []
    
    for table_info in schema:
        table_name = table_info['table_name']
        fingerprint = fingerprints.get(table_name, {})
        
        semantic_hints = fingerprint.get('semantic_hints', [])
        null_counts = fingerprint.get('null_counts', {})
        total_nulls = sum(null_counts.values())
        
        # Suggest null audit if there are significant nulls
        if total_nulls > 0:
            confidence = min(0.9, total_nulls / fingerprint.get('row_count', 1))
            reasoning = f"Found {total_nulls} null values across {len(null_counts)} columns"
            
            # Log the calculation for judge visibility
            log_entry = log_confidence_calculation(
                table_name, 'null_audit', confidence, reasoning, fingerprint, session_id
            )
            
            task_suggestions.append({
                "task_id": "null_audit",
                "table_name": table_name,
                "confidence": confidence,
                "reason": reasoning,
                "requires_alias": table_name.lower() != 'users',
                "debug_info": log_entry  # AI Reasoning for judges
            })
        
        # Suggest type repair if we see suspicious type patterns
        data_types = fingerprint.get('data_types', {})
        suspicious_numeric_cols = [col for col in ['amount', 'price', 'rating', 'age', 'quantity'] 
                                 if col in data_types and data_types.get(col, '') == 'TEXT']
        
        if suspicious_numeric_cols:
            confidence = 0.7 + (len(suspicious_numeric_cols) * 0.05)  # Higher confidence for more issues
            reasoning = f"Detected {len(suspicious_numeric_cols)} numeric columns stored as text: {', '.join(suspicious_numeric_cols)}"
            
            # Log the calculation
            log_entry = log_confidence_calculation(
                table_name, 'type_repair', confidence, reasoning, fingerprint, session_id
            )
            
            task_suggestions.append({
                "task_id": "type_repair",
                "table_name": table_name,
                "confidence": min(0.9, confidence),
                "reason": reasoning,
                "requires_alias": table_name.lower() != 'transactions',
                "debug_info": log_entry
            })
        
        # Suggest dedup if we have potential duplicate indicators
        columns = [col.split(' (')[0].lower() for col in table_info.get('columns', [])]
        if 'id' not in columns:
            confidence = 0.6
            reasoning = "No unique ID column detected - potential duplicates"
            
            # Log the calculation
            log_entry = log_confidence_calculation(
                table_name, 'dedup_bug', confidence, reasoning, fingerprint, session_id
            )
            
            task_suggestions.append({
                "task_id": "dedup_bug",
                "table_name": table_name,
                "confidence": confidence,
                "reason": reasoning,
                "requires_alias": table_name.lower() != 'order_items',
                "debug_info": log_entry
            })
        
        # Job market specific task
        if 'job_market_data' in semantic_hints:
            confidence = 0.8
            reasoning = "Detected job market data patterns (job titles, companies, ratings)"
            
            # Log the calculation
            log_entry = log_confidence_calculation(
                table_name, 'job_market_audit', confidence, reasoning, fingerprint, session_id
            )
            
            task_suggestions.append({
                "task_id": "job_market_audit",
                "table_name": table_name,
                "confidence": confidence,
                "reason": reasoning,
                "requires_alias": True,
                "suggested_alias": "users",
                "debug_info": log_entry
            })
    
    # Sort by confidence
    task_suggestions.sort(key=lambda x: x['confidence'], reverse=True)
    return task_suggestions


def detect_tasks_from_schema(schema: list[dict]) -> list[str]:
    """Legacy hard-coded mapping - kept for compatibility."""
    compatible_tasks = []
    table_names = {t['table_name'].lower() for t in schema}
    
    # null_audit: requires 'users' table
    if 'users' in table_names:
        compatible_tasks.append('null_audit')
    
    # type_repair: requires 'transactions' table
    if 'transactions' in table_names:
        compatible_tasks.append('type_repair')
    
    # dedup_bug: requires 'order_items' table
    if 'order_items' in table_names:
        compatible_tasks.append('dedup_bug')
    
    # fk_integrity: requires 'invoices' and 'customers' tables
    if 'invoices' in table_names and 'customers' in table_names:
        compatible_tasks.append('fk_integrity')
    
    # schema_drift: requires 'events_v1' and 'events_v2' tables
    if 'events_v1' in table_names and 'events_v2' in table_names:
        compatible_tasks.append('schema_drift')
    
    return compatible_tasks


# Add base64 upload endpoint for agent automation
@app.post("/upload-db-base64")
async def upload_db_base64(
    filename: str = Form(...),
    content_base64: str = Form(...),
    use_for_tasks: bool = Form(True),
    session_id: str = Form(None)
) -> dict:
    """Upload a database file using base64 encoding for agent automation.
    
    This endpoint allows programmatic file upload without using the OS file picker,
    enabling test automation agents to upload files directly.
    
    Args:
        filename: Original filename (e.g., "jobs_raw.csv")
        content_base64: Base64-encoded file content
        use_for_tasks: If True, enables RL tasks on this database
        session_id: Optional session identifier for alias isolation
        
    Returns:
        dict: Status, schema info, and compatible RL tasks
    """
    import base64 as b64
    
    try:
        # Decode base64 content
        content = b64.b64decode(content_base64)
        
        # Create a mock UploadFile
        class MockUploadFile:
            def __init__(self, filename, content):
                self.filename = filename
                self._content = content
                
        mock_file = MockUploadFile(filename, content)
        
        # Reuse the existing upload logic by calling the internal handler
        return await _process_upload(mock_file, use_for_tasks, session_id)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to decode base64 content: {str(e)}")


# Internal upload processor (refactored from upload_db)
async def _process_upload(
    file,
    use_for_tasks: bool = True,
    session_id: str = None
) -> dict:
    """Internal upload processing logic shared between upload endpoints."""
    global _custom_db_path, _custom_db_schema, _custom_db_filename, _custom_db_analysis
    global _uploaded_db_path, _uploaded_db_schema, _uploaded_db_tasks
    global table_aliases, active_sessions
    
    # Auto-clear aliases for this session on new upload
    if session_id and session_id in table_aliases:
        table_aliases[session_id].clear()
        active_sessions[session_id] = {
            "created_at": active_sessions.get(session_id, {}).get("created_at", datetime.now().isoformat()),
            "last_used": datetime.now().isoformat()
        }
    
    # Validate file extension
    filename = file.filename or ""
    ext = os.path.splitext(filename)[1].lower()
    valid_extensions = {'.sqlite', '.db', '.csv'}
    
    if ext not in valid_extensions:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid file extension: {ext}. Must be one of: {valid_extensions}"
        )
    
    # Read file content
    if hasattr(file, 'read'):
        content = await file.read()
    else:
        content = file._content
    
    # Validate file size (15MB limit)
    max_size = 15 * 1024 * 1024
    if len(content) > max_size:
        raise HTTPException(
            status_code=413,
            detail=f"File too large: {len(content)} bytes. Maximum is {max_size} bytes (15MB)"
        )
    
    try:
        # Hard reset any active episode so a new upload cannot inherit stale task state.
        old_conn = getattr(_env, "_original_conn", None)
        if old_conn is not None:
            try:
                old_conn.close()
            except Exception:
                pass
        _env._original_conn = None
        _env._active = False
        _env._done = False
        _env._episode_id = None
        _env._current_task_id = None
        _env._trajectory = []
        _env._task_context = {}

        _uploaded_db_path = None
        _uploaded_db_schema = {}
        _uploaded_db_tasks = []

        db_path = replace_temp_db(_custom_db_path)
        conn = None

        if ext == '.csv':
            conn = open_sqlite_db(db_path)
            csv_text = content.decode('utf-8')
            csv_file = io.StringIO(csv_text)
            reader = csv.DictReader(csv_file)
            
            if not reader.fieldnames:
                raise HTTPException(status_code=400, detail="CSV file has no headers")
            
            # Infer types from first 100 rows
            rows = list(reader)
            if not rows:
                raise HTTPException(status_code=400, detail="CSV file is empty")
            
            # Create table
            table_name = os.path.splitext(os.path.basename(filename))[0]
            # Sanitize table name
            table_name = re.sub(r'[^a-zA-Z0-9_]', '_', table_name)
            if not table_name[0].isalpha() and table_name[0] != '_':
                table_name = '_' + table_name
            
            columns = []
            for col_name in reader.fieldnames:
                # Sanitize column name
                safe_col = re.sub(r'[^a-zA-Z0-9_]', '_', col_name)
                if not safe_col[0].isalpha() and safe_col[0] != '_':
                    safe_col = '_' + safe_col
                
                values = [row.get(col_name, '') for row in rows[:100]]
                col_type = infer_sqlite_type(values)
                columns.append(f'"{safe_col}" {col_type}')
            
            cols_joined = ", ".join(columns)
            create_sql = f'CREATE TABLE "{table_name}" ({cols_joined})'
            conn.execute(create_sql)
            
            # Insert rows with data sanitization
            placeholders = ', '.join(['?' for _ in reader.fieldnames])
            insert_sql = f'INSERT INTO "{table_name}" VALUES ({placeholders})'
            
            for row in rows:
                # Sanitize each value during insertion
                values = [sanitize_data_value(row.get(col_name, None)) for col_name in reader.fieldnames]
                conn.execute(insert_sql, values)
            
            conn.commit()
        else:
            with open(db_path, "wb") as handle:
                handle.write(content)
            conn = open_sqlite_db(db_path)

        _custom_db_path = db_path
        _custom_db_schema = get_db_schema(conn)
        _custom_db_filename = filename
        _custom_db_analysis = analyze_uploaded_database(conn, schema_table_map(_custom_db_schema))

        health_audits = {}
        fingerprints = {}
        for table_info in _custom_db_schema:
            table_name = table_info['table_name']
            health_audits[table_name] = perform_health_audit(conn, table_name)
            fingerprints[table_name] = analyze_table_fingerprint(conn, table_name)

        semantic_tasks = []
        for task_id, candidate in _custom_db_analysis.get("task_candidates", {}).items():
            semantic_tasks.append({
                "task_id": task_id,
                "table_name": candidate.get("table"),
                "confidence": candidate.get("confidence", 0.0),
                "reason": candidate.get("context_hint") or candidate.get("reason", ""),
                "supported": candidate.get("supported", False),
                "verification_sql": candidate.get("verification_sql"),
            })

        _uploaded_db_tasks = [
            task_id
            for task_id, candidate in _custom_db_analysis.get("task_candidates", {}).items()
            if candidate.get("supported")
        ]

        if use_for_tasks and _uploaded_db_tasks:
            _uploaded_db_path = _custom_db_path
            _uploaded_db_schema = schema_table_map(_custom_db_schema)
        else:
            _uploaded_db_path = None
            _uploaded_db_schema = {}

        conn.close()
        conn = None

        return {
            "status": "ok",
            "filename": filename,
            "tables": _custom_db_schema,
            "compatible_tasks": _uploaded_db_tasks,
            "semantic_tasks": semantic_tasks,
            "use_for_tasks": use_for_tasks and len(_uploaded_db_tasks) > 0,
            "health_audits": health_audits,
            "fingerprints": fingerprints,
            "task_contexts": _custom_db_analysis.get("task_candidates", {}),
            "session_id": session_id,
            "detected_table": _custom_db_schema[0]['table_name'] if _custom_db_schema else None
        }
        
    except HTTPException:
        raise
    except Exception as e:
        import traceback
        error_detail = f"Failed to process file: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        raise HTTPException(status_code=500, detail=error_detail)
    finally:
        if conn is not None:
            try:
                conn.close()
            except Exception:
                pass


@app.post("/upload-db")
async def upload_db(
    file: UploadFile = File(...),
    use_for_tasks: bool = True,
    session_id: str = None
) -> dict:
    """Upload a .sqlite, .db, or .csv file for custom exploration and RL tasks.
    
    Args:
        file: The uploaded file (.sqlite, .db, or .csv)
        use_for_tasks: If True, enables RL tasks on this database
        session_id: Optional session identifier for alias isolation
        
    Returns:
        dict: Status, schema info, and compatible RL tasks
        
    Raises:
        HTTPException 400: If file extension is invalid
        HTTPException 413: If file is too large (>15MB)
        HTTPException 500: If processing fails
    """
    # Use the shared processing logic
    return await _process_upload(file, use_for_tasks, session_id)


@app.post("/set-table-alias")
async def set_table_alias(body: dict, session_id: str = None) -> dict:
    """Set a table alias for manual task override.
    
    Args:
        body: JSON body with {"actual_table": "cleaned_ds_jobs2", "alias_as": "users"}
        session_id: Optional session identifier for multi-user support
        
    Returns:
        dict: Status and current aliases
    """
    global table_aliases, active_sessions
    
    # Generate or use provided session ID
    if not session_id:
        session_id = str(uuid.uuid4())
    
    # Initialize session if not exists
    if session_id not in table_aliases:
        table_aliases[session_id] = {}
    
    actual_table = body.get("actual_table", "").strip()
    alias_as = body.get("alias_as", "").strip()
    
    if not actual_table or not alias_as:
        raise HTTPException(status_code=422, detail="Both 'actual_table' and 'alias_as' required")
    
    table_aliases[session_id][actual_table] = alias_as
    
    # Update session activity
    active_sessions[session_id] = {
        "created_at": datetime.now().isoformat(),
        "last_used": datetime.now().isoformat()
    }
    
    return {
        "status": "ok",
        "session_id": session_id,
        "message": f"Table '{actual_table}' will be treated as '{alias_as}' for RL tasks",
        "aliases": table_aliases[session_id].copy()
    }


@app.get("/table-aliases")
async def get_table_aliases(session_id: str = None) -> dict:
    """Get current table aliases for session."""
    if not session_id or session_id not in table_aliases:
        return {"aliases": {}, "session_id": session_id}
    
    return {
        "session_id": session_id,
        "aliases": table_aliases[session_id].copy()
    }


@app.delete("/table-aliases")
async def clear_table_aliases(session_id: str = None) -> dict:
    """Clear all table aliases for session or all sessions."""
    global table_aliases, active_sessions
    
    if session_id and session_id in table_aliases:
        table_aliases[session_id].clear()
        if session_id in active_sessions:
            del active_sessions[session_id]
        return {
            "status": "ok",
            "session_id": session_id,
            "message": f"Aliases cleared for session {session_id}"
        }
    else:
        # Clear all sessions
        table_aliases.clear()
        active_sessions.clear()
        return {
            "status": "ok",
            "message": "All table aliases cleared for all sessions"
        }


@app.post("/new-session")
async def create_session() -> dict:
    """Create a new session for isolated alias management."""
    session_id = str(uuid.uuid4())
    table_aliases[session_id] = {}
    active_sessions[session_id] = {
        "created_at": datetime.now().isoformat(),
        "last_used": datetime.now().isoformat()
    }
    
    return {
        "session_id": session_id,
        "status": "ok",
        "message": "New session created"
    }


def sanitize_sql_identifiers(sql: str) -> str:
    """Pass-through function kept for compatibility instead of aggressive quoting."""
    return sql


def apply_table_aliases(sql: str, session_id: str = None) -> str:
    """Apply table aliases to SQL queries intelligently without mangling complex queries."""
    global table_aliases
    import re
    
    # Get aliases for session or default to empty
    session_aliases = {}
    if session_id and session_id in table_aliases:
        session_aliases = table_aliases[session_id]
    
    # First sanitize all identifiers (now a pass-through)
    sanitized_sql = sanitize_sql_identifiers(sql)
    
    # Then apply aliases
    for actual, alias in session_aliases.items():
        # Replace unquoted FROM table_name
        sanitized_sql = re.sub(
            r'\bFROM\s+' + re.escape(actual) + r'\b', 
            f'FROM "{alias}"', 
            sanitized_sql, 
            flags=re.IGNORECASE
        )
        # Replace quoted FROM "table_name"
        sanitized_sql = re.sub(
            r'\bFROM\s+"' + re.escape(actual) + r'"\b', 
            f'FROM "{alias}"', 
            sanitized_sql, 
            flags=re.IGNORECASE
        )
        # Replace back-ticked FROM `table_name`
        sanitized_sql = re.sub(
            r'\bFROM\s+`' + re.escape(actual) + r'`\b', 
            f'FROM "{alias}"', 
            sanitized_sql, 
            flags=re.IGNORECASE
        )
        # Replace bracketed FROM [table_name]
        sanitized_sql = re.sub(
            r'\bFROM\s+\[' + re.escape(actual) + r'\]\b', 
            f'FROM "{alias}"', 
            sanitized_sql, 
            flags=re.IGNORECASE
        )

        # Replace JOIN table_name
        sanitized_sql = re.sub(
            r'\bJOIN\s+' + re.escape(actual) + r'\b', 
            f'JOIN "{alias}"', 
            sanitized_sql, 
            flags=re.IGNORECASE
        )
        sanitized_sql = re.sub(
            r'\bJOIN\s+"' + re.escape(actual) + r'"\b', 
            f'JOIN "{alias}"', 
            sanitized_sql, 
            flags=re.IGNORECASE
        )
        
        # Replace table.column with alias.column (unquoted)
        sanitized_sql = re.sub(
            r'\b' + re.escape(actual) + r'\.', 
            f'"{alias}".', 
            sanitized_sql,
            flags=re.IGNORECASE
        )
        # Replace "table".column with alias.column (quoted)
        sanitized_sql = re.sub(
            r'"' + re.escape(actual) + r'"\.', 
            f'"{alias}".', 
            sanitized_sql,
            flags=re.IGNORECASE
        )
    
    return sanitized_sql


@app.post("/query-custom")
async def query_custom(body: dict, session_id: str = None) -> dict:
    """Execute a SELECT query against the uploaded custom database.
    
    Args:
        body: JSON body with {"sql": "SELECT ..."}
        session_id: Optional session identifier for alias isolation
        
    Returns:
        dict: Query results with rows, columns, and row_count, plus suggestions if 0 rows
        
    Raises:
        HTTPException 422: If no database uploaded or SQL is invalid
    """
    global _custom_db_path

    if _custom_db_path is None:
        raise HTTPException(
            status_code=422,
            detail="No database uploaded yet. Use /upload-db first."
        )
    
    sql = body.get("sql", "").strip()
    if not sql:
        raise HTTPException(status_code=422, detail="No SQL provided")

    rewritten_sql, rewrites, unresolved = rewrite_known_quoted_identifiers(sql, _custom_db_schema)
    if unresolved:
        return {
            "error": "Unknown quoted identifier detected.",
            "message": (
                "SQL Error — a double-quoted name is not in the uploaded schema. "
                "SQLite may treat it as a string literal and produce misleading results."
            ),
            "suggestions": [
                f'Use "{target}" instead of "{source}".'
                for source, target in unresolved[:3]
            ] + ["Use single quotes for string literals and schema-panel names for identifiers."],
            "active_tables": [table["table_name"] for table in _custom_db_schema],
        }
    sql = rewritten_sql
    
    if not is_readonly_statement(sql):
        return {
            "error": "QUERY action is read-only. This action only accepts SELECT, WITH, EXPLAIN, or PRAGMA.",
            "message": "SQL Error — check your query. Only read-only SELECT/CTE/EXPLAIN/PRAGMA queries are allowed in custom mode.",
            "suggestions": ["Try: SELECT * FROM table LIMIT 5", "Check schema panel for table names"]
        }

    statements = split_sql_statements(sql)
    if any(not is_readonly_statement(stmt) for stmt in statements):
        return {
            "error": "Only read-only statements are allowed in Zone 1.",
            "message": "SQL Error — mixed read/write statements are not allowed.",
            "suggestions": ["Run a single SELECT/CTE/EXPLAIN/PRAGMA statement at a time."]
        }

    conn = open_sqlite_db(_custom_db_path)
    try:
        warnings = build_query_quality_warnings(sql, _custom_db_analysis)
        if rewrites:
            warnings = [
                "Auto-mapped quoted identifiers: "
                + ", ".join(f'"{src}" -> "{dst}"' for src, dst in rewrites[:3])
            ] + warnings
        task_suggestions = build_dataset_task_suggestions(sql, _custom_db_analysis)
        cursor = None
        for stmt in statements:
            cursor = conn.execute(stmt)

        if cursor is None:
            return {
                "rows": [],
                "columns": [],
                "row_count": 0,
                "message": "No executable SQL statement was found.",
                "suggestions": ["Try: SELECT * FROM table LIMIT 5 to explore data"]
            }

        if cursor.description:
            columns = [d[0] for d in cursor.description]
            rows = cursor.fetchall()
            row_dicts = [dict(row) for row in rows]
            suggestions = None
            if len(row_dicts) == 0:
                suggestions = suggest_query_alternatives(conn, sql, "0 rows")
            return {
                "rows": row_dicts,
                "columns": columns,
                "row_count": len(row_dicts),
                "suggestions": suggestions,
                "warnings": warnings,
                "task_suggestions": task_suggestions,
                "source_insights": build_source_data_insights(sql, _custom_db_analysis),
                "active_tables": [table["table_name"] for table in _custom_db_schema],
            }

        return {
            "rows": [],
            "columns": [],
            "row_count": 0,
            "message": "Query executed successfully but returned no rows",
            "suggestions": ["Try: SELECT * FROM table LIMIT 5 to explore data"],
            "warnings": warnings,
            "task_suggestions": task_suggestions,
            "source_insights": build_source_data_insights(sql, _custom_db_analysis),
            "active_tables": [table["table_name"] for table in _custom_db_schema],
        }
    except sqlite3.Error as e:
        error_msg = str(e)
        suggestions = suggest_query_alternatives(conn, sql, error_msg)
        return {
            "error": error_msg,
            "message": "SQL Error — check your query. Check table and column names in the schema panel.",
            "suggestions": suggestions,
            "warnings": build_query_quality_warnings(sql, _custom_db_analysis),
            "task_suggestions": build_dataset_task_suggestions(sql, _custom_db_analysis),
            "active_tables": [table["table_name"] for table in _custom_db_schema],
        }
    finally:
        conn.close()


@app.get("/custom-db-status")
async def custom_db_status() -> dict:
    """Check if a custom database is loaded and return its schema.
    
    Returns:
        dict: {loaded: true/false, tables: [...]} 
    """
    if _custom_db_path is None:
        return {"loaded": False}
    
    return {
        "loaded": True,
        "filename": _custom_db_filename,
        "tables": _custom_db_schema,
        "compatible_tasks": _uploaded_db_tasks,
        "task_contexts": _custom_db_analysis.get("task_candidates", {}),
    }


# ─── AI INSIGHTS ENDPOINT (Groq / OpenAI-compatible) ──────────────────────

def _get_ai_client():
    """Get an OpenAI-compatible client for AI insights.
    """
    runtime = _resolve_llm_runtime()
    if runtime["configured"]:
        try:
            from openai import OpenAI

            return OpenAI(
                api_key=runtime["api_key"],
                base_url=runtime["base_url"],
            ), runtime["model"]
        except Exception:
            pass

    return None, None


def _generate_local_insights(sql: str, rows: list, columns: list, row_count: int) -> list[str]:
    """Generate source-aware analysis without LLM (fast fallback)."""
    insights = []
    
    if not rows:
        insights.append(f"Query returned 0 rows. The SQL may be too restrictive or target wrong table names.")
        insights.extend(build_source_data_insights(sql, _custom_db_analysis))
        return insights[:6]
    
    insights.append(f"Query returned {row_count} rows across {len(columns)} columns")
    insights.extend(build_query_quality_warnings(sql, _custom_db_analysis))
    
    # Null analysis
    null_stats = {}
    for col in columns:
        null_count = sum(1 for r in rows if r.get(col) is None or r.get(col) == '')
        if null_count > 0:
            pct = round(null_count / len(rows) * 100, 1)
            null_stats[col] = {'count': null_count, 'pct': pct}
    
    if null_stats:
        worst = max(null_stats.items(), key=lambda x: x[1]['pct'])
        insights.append(f"Highest null rate: column '{worst[0]}' at {worst[1]['pct']}% ({worst[1]['count']} nulls)")
        if len(null_stats) > 1:
            insights.append(f"{len(null_stats)} columns have null/empty values — consider a null_audit task")
    else:
        insights.append("No NULL values detected in the returned rows, but source-table issues may still exist outside this result set")
    
    # Numeric analysis
    for col in columns:
        values = [r.get(col) for r in rows if r.get(col) is not None]
        numeric = []
        for v in values:
            try:
                numeric.append(float(str(v).replace('$','').replace(',','').replace('%','')))
            except (ValueError, TypeError):
                pass
        
        if len(numeric) > len(values) * 0.7 and len(numeric) > 3:
            avg = sum(numeric) / len(numeric)
            std = (sum((x - avg)**2 for x in numeric) / len(numeric)) ** 0.5
            outliers = sum(1 for x in numeric if abs(x - avg) > 2 * std)
            insights.append(f"Column '{col}': avg={avg:.2f}, std={std:.2f}, {outliers} outlier(s) (>2σ)")
    
    # Duplicate detection
    for col in columns:
        if is_aggregate_like_column_name(col) or not is_key_like_column_name(col):
            continue
        values = [str(r.get(col, '')) for r in rows if r.get(col) is not None]
        unique = set(values)
        if len(unique) < len(values) * 0.7 and len(values) > 5:
            dup_count = len(values) - len(unique)
            insights.append(f"Column '{col}' has {dup_count} duplicate values ({len(unique)}/{len(values)} unique) — check for dedup issues")
    
    for extra in build_dataset_task_suggestions(sql, _custom_db_analysis):
        if len(insights) >= 8:
            break
        insights.append(extra)

    return insights[:8]


@app.post("/ai-insights")
async def ai_insights(body: dict) -> dict:
    """Generate AI-powered insights from query results.
    
    Uses Groq API (or HF Inference) for intelligent analysis.
    Falls back to local heuristic analysis if no API key configured.
    
    Args:
        body: JSON with {sql, rows, row_count, columns}
        
    Returns:
        dict: {insights: [str, ...]}
    """
    sql = body.get("sql", "")
    rows = body.get("rows", [])
    row_count = body.get("row_count", len(rows))
    columns = body.get("columns", [])
    
    # Try AI-powered insights first
    client, model = _get_ai_client()
    
    if client and model:
        try:
            # Build a compact data summary for the LLM
            data_summary = f"SQL: {sql[:300]}\n"
            data_summary += f"Result: {row_count} rows, {len(columns)} columns\n"
            data_summary += f"Columns: {', '.join(columns[:15])}\n"
            
            if rows:
                # Include a few sample rows
                sample = rows[:5]
                data_summary += f"Sample data (first {len(sample)} rows):\n"
                for r in sample:
                    row_str = ', '.join(f"{k}={v}" for k, v in list(r.items())[:8])
                    data_summary += f"  {row_str}\n"
                
                # Calculate null stats
                null_info = []
                for col in columns:
                    null_count = sum(1 for r in rows if r.get(col) is None or r.get(col) == '')
                    if null_count > 0:
                        null_info.append(f"{col}: {null_count}/{len(rows)} nulls ({round(null_count/len(rows)*100)}%)")
                if null_info:
                    data_summary += f"Null columns: {'; '.join(null_info)}\n"
            
            import asyncio
            response = await asyncio.to_thread(
                lambda: client.chat.completions.create(
                    model=model,
                    messages=[
                        {"role": "system", "content": (
                            "You are a SQL data quality expert. Analyze the query results and provide "
                            "3-5 actionable insights about data quality issues. Be specific and concise. "
                            "Focus on: null rates, type mismatches, duplicates, outliers, referential issues. "
                            "Each insight should be 1 sentence. Use emoji prefixes: 📊 ⚠️ ✅ 🔍 📈 🔄 🔧"
                        )},
                        {"role": "user", "content": data_summary}
                    ],
                    temperature=0.3,
                    max_tokens=400,
                    stream=False
                )
            )
            
            text = response.choices[0].message.content or ""
            # Split into individual insights
            insights = [line.strip() for line in text.strip().split('\n') if line.strip() and len(line.strip()) > 5]
            
            if insights:
                return {"insights": insights[:6], "source": "ai"}
        
        except Exception as e:
            # Log but don't fail — fall back to local
            print(f"AI insights error: {e}")
    
    # Fallback: local heuristic insights
    insights = _generate_local_insights(sql, rows, columns, row_count)
    return {"insights": insights, "source": "local"}


def run_server() -> None:
    """
    Entry point for 'uv run server' command.

    Uses PORT environment variable if set, otherwise defaults to 7860
    (matching the Dockerfile EXPOSE and HF Spaces convention).
    Override locally: PORT=8000 uv run server
    """
    import os

    import uvicorn

    port = int(os.getenv("PORT", "7860"))
    uvicorn.run("server.main:app", host="0.0.0.0", port=port, reload=False)


def main() -> None:
    """Primary CLI entrypoint expected by OpenEnv validators."""
    run_server()