backend/main.py

import os
import re
import time
from datetime import datetime
from collections import defaultdict
from fastapi import FastAPI, HTTPException, UploadFile, File, Response
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Dict, Any

from backend import config
from backend.services import csv_service, sql_service, llm_service, insight_service
from backend.services.schema_service import schema_service
from backend.schemas import (
    DatasetsResponse, DatasetInfo, DatasetSchemaResponse, ColumnDetail,
    UploadResponse, DeleteResponse, QueryRequest, QueryResponse, SQLResults,
    HealthResponse, SuggestionsResponse, QueryHistoryResponse, QueryHistoryItem,
    SchemaDatasetsResponse, SchemaUploadResponse, SchemaInfoResponse,
    SchemaPreviewsResponse, SchemaQueryResponse, SchemaQueryRequest
)


# Global in-memory caches
query_histories = defaultdict(list)  # db_name -> list of dicts
suggestions_cache = {}              # db_name -> list of strings

app = FastAPI(
    title="Natural Language Data Analyst API",
    description="Backend API for Natural Language Data Analyst application",
    version="0.1.0"
)

# Enforce CORS allowing all origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["X-Process-Time"]  # Make sure the frontend can read this header
)

def validate_db_name(db_name: str):
    """
    Validates that db_name only contains alphanumeric characters and underscores.
    Prevents path traversal and injection.
    """
    if not re.match(r"^[a-zA-Z0-9_]+$", db_name):
        raise HTTPException(
            status_code=400,
            detail="db_name must only contain alphanumeric characters and underscores."
        )

# 1. GET /datasets
@app.get("/datasets", response_model=DatasetsResponse)
def get_datasets():
    print("[API] GET /datasets called")
    try:
        # Load sample datasets first
        samples = csv_service.get_sample_datasets()
        sample_names = {ds["name"] for ds in samples}

        datasets_list = []
        # Append sample datasets mapped to output schema
        for ds in samples:
            datasets_list.append(
                DatasetInfo(
                    db_name=ds["name"],
                    display_name=ds["display_name"],
                    description=ds["description"],
                    columns=[col["name"] for col in ds["columns"]],
                    row_count=ds["row_count"],
                    is_sample=True
                )
            )

        # Scan databases/ directory for any user uploaded ones
        if os.path.exists(config.DB_DIR):
            for filename in os.listdir(config.DB_DIR):
                if filename.endswith(".db"):
                    db_name = os.path.splitext(filename)[0]
                    if db_name not in sample_names:
                        try:
                            schema = csv_service.get_table_schema(db_name)
                            datasets_list.append(
                                DatasetInfo(
                                    db_name=db_name,
                                    display_name=db_name.replace("_", " ").capitalize(),
                                    description=f"User-uploaded dataset from {db_name}.csv",
                                    columns=[col["name"] for col in schema["columns"]],
                                    row_count=schema["row_count"],
                                    is_sample=False
                                )
                            )
                        except Exception as e:
                            print(f"[API] Error reading user DB '{db_name}': {e}")

        return DatasetsResponse(datasets=datasets_list, total=len(datasets_list))
    except Exception as e:
        print(f"[API] Error in GET /datasets: {e}")
        raise HTTPException(status_code=500, detail=f"Failed to fetch datasets: {str(e)}")

# 2. GET /datasets/{db_name}/schema
@app.get("/datasets/{db_name}/schema", response_model=DatasetSchemaResponse)
def get_dataset_schema(db_name: str):
    print(f"[API] GET /datasets/{db_name}/schema called")
    validate_db_name(db_name)

    db_path = os.path.join(config.DB_DIR, f"{db_name}.db")
    if not os.path.exists(db_path):
        raise HTTPException(status_code=404, detail="Dataset not found.")

    try:
        schema = csv_service.get_table_schema(db_name)
        return DatasetSchemaResponse(
            db_name=db_name,
            table_name=schema["table_name"],
            columns=[
                ColumnDetail(
                    name=col["name"],
                    type=col["type"],
                    sample_values=col["sample_values"]
                ) for col in schema["columns"]
            ],
            row_count=schema["row_count"]
        )
    except Exception as e:
        print(f"[API] Error fetching schema for '{db_name}': {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 3. POST /upload
@app.post("/upload", response_model=UploadResponse)
def upload_csv(file: UploadFile = File(...)):
    print(f"[API] POST /upload called with file: {file.filename}")
    if not file.filename.endswith(".csv"):
        raise HTTPException(status_code=400, detail="Only .csv files are allowed.")

    # Sanitize file name to produce db_name
    base_name = os.path.splitext(file.filename)[0].lower().replace(" ", "_")
    db_name = re.sub(r"[^a-zA-Z0-9_]", "", base_name)
    if not db_name:
        raise HTTPException(status_code=400, detail="Invalid filename characters.")

    db_path = os.path.join(config.DB_DIR, f"{db_name}.db")
    if os.path.exists(db_path):
        raise HTTPException(status_code=409, detail="Dataset already uploaded. Delete it first.")

    # Save to temp location in DB_DIR
    temp_path = os.path.join(config.DB_DIR, f"temp_{db_name}.csv")
    try:
        with open(temp_path, "wb") as buffer:
            buffer.write(file.file.read())
        
        # Load CSV to SQLite database
        schema_info = csv_service.load_csv_to_sqlite(temp_path, db_name)
        
        display_name = db_name.replace("_", " ").capitalize()
        columns = [col["name"] for col in schema_info["columns"]]
        
        return UploadResponse(
            success=True,
            db_name=db_name,
            display_name=display_name,
            columns=columns,
            row_count=schema_info["row_count"],
            message=f"Dataset '{display_name}' uploaded and processed successfully."
        )
    except Exception as e:
        print(f"[API] Error uploading CSV: {e}")
        # If database was created but failed load, clean it up
        if os.path.exists(db_path):
            try:
                os.remove(db_path)
            except Exception:
                pass
        raise HTTPException(status_code=500, detail=f"Failed to process CSV file: {str(e)}")
    finally:
        # Cleanup temporary CSV file
        if os.path.exists(temp_path):
            try:
                os.remove(temp_path)
            except Exception as e:
                print(f"[API] Failed to delete temp file '{temp_path}': {e}")

# 4. DELETE /datasets/{db_name}
@app.delete("/datasets/{db_name}", response_model=DeleteResponse)
def delete_dataset(db_name: str):
    print(f"[API] DELETE /datasets/{db_name} called")
    validate_db_name(db_name)

    if db_name in ["sales", "employees", "ecommerce"]:
        raise HTTPException(status_code=403, detail="Cannot delete sample datasets.")

    db_path = os.path.join(config.DB_DIR, f"{db_name}.db")
    if not os.path.exists(db_path):
        raise HTTPException(status_code=404, detail="Dataset not found.")

    try:
        os.remove(db_path)
        return DeleteResponse(
            success=True,
            message=f"Dataset '{db_name}' deleted successfully."
        )
    except Exception as e:
        print(f"[API] Error deleting dataset '{db_name}': {e}")
        raise HTTPException(status_code=500, detail=f"Failed to delete dataset: {str(e)}")

# 5. POST /query
@app.post("/query", response_model=QueryResponse)
def run_query(request: QueryRequest, response: Response):
    print(f"[API] POST /query called for db: {request.db_name}, question: '{request.question}'")
    start_time = time.perf_counter()
    
    db_path = os.path.join(config.DB_DIR, f"{request.db_name}.db")
    if not os.path.exists(db_path):
        raise HTTPException(status_code=404, detail=f"Database '{request.db_name}' not found.")

    try:
        # a. Get schema
        schema = csv_service.get_table_schema(request.db_name)

        # b. Generate SQL
        generation = llm_service.generate_sql(request.question, schema)
        sql = generation["sql"]

        # c. Validate SQL
        if not sql_service.validate_sql(sql):
            raise HTTPException(
                status_code=400,
                detail="Dangerous or invalid SQL query generated. Only SELECT operations are allowed."
            )

        # d. Execute SQL with auto-retry
        retries = 0
        results = None
        current_sql = sql
        last_error = ""

        while True:
            try:
                results = sql_service.execute_query(request.db_name, current_sql)
                # Success, break loop
                break
            except Exception as e:
                last_error = str(e)
                print(f"[API] Query execution failed: {last_error}")
                
                if retries < config.MAX_SQL_RETRIES:
                    retries += 1
                    print(f"[API] Retrying query fix (Attempt {retries}/{config.MAX_SQL_RETRIES})...")
                    try:
                        fix_result = llm_service.fix_sql(current_sql, last_error, schema)
                        fixed_sql = fix_result["sql"]
                        
                        # Validate fixed SQL
                        if not sql_service.validate_sql(fixed_sql):
                            raise ValueError("Fixed SQL failed validation check.")
                        
                        current_sql = fixed_sql
                    except Exception as fix_err:
                        print(f"[API] Failed to generate/validate fix: {fix_err}")
                        # Keep the loop going, next attempts might fix it or we exit with retries exceeded
                else:
                    # Retries exceeded
                    print("[API] Maximum SQL retries exceeded. Aborting.")
                    raise HTTPException(
                        status_code=422,
                        detail=f"SQL execution failed after {config.MAX_SQL_RETRIES} retries. Last error: {last_error}"
                    )

        # e. Generate insight
        insight = insight_service.generate_insight(request.question, current_sql, results)

        # f. Calculate processing time and inject header
        process_time_ms = (time.perf_counter() - start_time) * 1000.0
        response.headers["X-Process-Time"] = f"{process_time_ms:.2f}ms"
        print(f"[API] Request processed in {process_time_ms:.2f}ms. Header added.")

        # Save to query history
        query_histories[request.db_name].append({
            "question": request.question,
            "sql": current_sql,
            "insight": insight,
            "timestamp": datetime.now().isoformat(),
            "row_count": results["row_count"]
        })
        if len(query_histories[request.db_name]) > 20:
            query_histories[request.db_name] = query_histories[request.db_name][-20:]

        return QueryResponse(
            question=request.question,
            sql=current_sql,
            results=SQLResults(
                columns=results["columns"],
                rows=results["rows"],
                row_count=results["row_count"],
                execution_time_ms=results["execution_time_ms"]
            ),
            insight=insight,
            retries=retries,
            db_name=request.db_name
        )

    except HTTPException as he:
        raise he
    except Exception as e:
        print(f"[API] Internal error in /query pipeline: {e}")
        raise HTTPException(status_code=500, detail=f"Internal pipeline error: {str(e)}")

# 6. GET /datasets/{db_name}/preview
@app.get("/datasets/{db_name}/preview", response_model=SQLResults)
def preview_dataset(db_name: str):
    print(f"[API] GET /datasets/{db_name}/preview called")
    validate_db_name(db_name)

    db_path = os.path.join(config.DB_DIR, f"{db_name}.db")
    if not os.path.exists(db_path):
        raise HTTPException(status_code=404, detail="Dataset not found.")

    try:
        preview_sql = "SELECT * FROM data LIMIT 10"
        results = sql_service.execute_query(db_name, preview_sql)
        return SQLResults(
            columns=results["columns"],
            rows=results["rows"],
            row_count=results["row_count"],
            execution_time_ms=results["execution_time_ms"]
        )
    except Exception as e:
        print(f"[API] Error in preview for '{db_name}': {e}")
        raise HTTPException(status_code=500, detail=f"Failed to generate preview: {str(e)}")

# 7. GET /health
@app.get("/health", response_model=HealthResponse)
def health_check():
    print("[API] GET /health called")
    groq_configured = bool(config.GROQ_API_KEY)
    
    db_dir_accessible = False
    datasets_loaded = 0
    if os.path.exists(config.DB_DIR):
        db_dir_accessible = os.access(config.DB_DIR, os.R_OK | os.W_OK)
        try:
            datasets_loaded = len([f for f in os.listdir(config.DB_DIR) if f.endswith(".db")])
        except Exception:
            pass

    status = "healthy" if (groq_configured and db_dir_accessible) else "unhealthy"

    return HealthResponse(
        status=status,
        groq_configured=groq_configured,
        datasets_loaded=datasets_loaded,
        sample_datasets=["sales", "employees", "ecommerce"]
    )

# 8. GET /datasets/{db_name}/suggestions
@app.get("/datasets/{db_name}/suggestions", response_model=SuggestionsResponse)
def get_dataset_suggestions(db_name: str):
    print(f"[API] GET /datasets/{db_name}/suggestions called")
    validate_db_name(db_name)

    db_path = os.path.join(config.DB_DIR, f"{db_name}.db")
    if not os.path.exists(db_path):
        raise HTTPException(status_code=404, detail="Dataset not found.")

    # Check cache first
    if db_name in suggestions_cache:
        print(f"[API] Suggestions cache hit for '{db_name}'")
        return SuggestionsResponse(db_name=db_name, suggestions=suggestions_cache[db_name])

    try:
        # Retrieve table schema
        schema = csv_service.get_table_schema(db_name)
        # Generate question suggestions
        suggestions = llm_service.generate_question_suggestions(schema)
        # Store in cache
        suggestions_cache[db_name] = suggestions
        
        return SuggestionsResponse(db_name=db_name, suggestions=suggestions)
    except Exception as e:
        print(f"[API] Error generating suggestions for '{db_name}': {e}")
        raise HTTPException(status_code=500, detail=f"Failed to generate suggestions: {str(e)}")

# 9. GET /query/history
@app.get("/query/history", response_model=QueryHistoryResponse)
def get_query_history(db_name: str):
    print(f"[API] GET /query/history called for '{db_name}'")
    validate_db_name(db_name)
    
    db_path = os.path.join(config.DB_DIR, f"{db_name}.db")
    is_schema = False
    if not os.path.exists(db_path):
        schema_db_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.db")
        if os.path.exists(schema_db_path):
            is_schema = True
            db_path = schema_db_path

    if not os.path.exists(db_path):
        raise HTTPException(status_code=404, detail="Dataset not found.")

    history_key = f"schema_{db_name}" if is_schema else db_name
    history_list = query_histories.get(history_key, [])
    # Map raw history dict items to QueryHistoryItem
    items = [
        QueryHistoryItem(
            question=h["question"],
            sql=h["sql"],
            insight=h["insight"],
            timestamp=h["timestamp"],
            row_count=h["row_count"]
        ) for h in history_list
    ]
    # Return in reverse chronological order (latest queries first)
    items.reverse()

    return QueryHistoryResponse(
        db_name=db_name,
        history=items,
        total=len(items)
    )

# 10. GET /schema-datasets
@app.get("/schema-datasets", response_model=SchemaDatasetsResponse)
def get_schema_datasets():
    print("[API] GET /schema-datasets called")
    try:
        datasets_list = schema_service.get_all_schema_datasets()
        return SchemaDatasetsResponse(datasets=datasets_list, total=len(datasets_list))
    except Exception as e:
        print(f"[API] Error in GET /schema-datasets: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 11. POST /upload/schema
@app.post("/upload/schema", response_model=SchemaUploadResponse)
def upload_schema_db(
    schema_file: UploadFile = File(None),
    db_file: UploadFile = File(None),
    erd_image: UploadFile = File(None)
):
    schema_fn = schema_file.filename if schema_file else "None"
    db_fn = db_file.filename if db_file else "None"
    erd_fn = erd_image.filename if erd_image else "None"
    print(f"[API] POST /upload/schema called with schema: {schema_fn}, db: {db_fn}, erd_image: {erd_fn}")
    
    if not schema_file and not erd_image:
        raise HTTPException(
            status_code=400,
            detail="Please provide either a .sql schema file or an ERD diagram image"
        )
        
    if schema_file and not schema_file.filename.endswith(".sql"):
        raise HTTPException(status_code=400, detail="Schema file must have a .sql extension.")
        
    if erd_image:
        ext = os.path.splitext(erd_image.filename)[1].lower()
        if ext not in (".png", ".jpg", ".jpeg", ".pdf"):
            raise HTTPException(
                status_code=400,
                detail="ERD image must be a .png, .jpg, .jpeg, or .pdf file."
            )
            
    if db_file and not db_file.filename.endswith((".db", ".sqlite")):
        raise HTTPException(status_code=400, detail="Database file must have a .db or .sqlite extension.")
        
    try:
        # Determine db_name from the uploaded db file or fall back to the schema file name
        if db_file:
            base_name = os.path.splitext(db_file.filename)[0].lower().replace(" ", "_")
        elif schema_file:
            base_name = os.path.splitext(schema_file.filename)[0].lower().replace(" ", "_")
            # Strip common suffixes like _schema, _ddl
            for suffix in ["_schema", "_ddl", "_create", "_tables"]:
                if base_name.endswith(suffix):
                    base_name = base_name[:-len(suffix)]
                    break
        else:
            base_name = "uploaded_db"
        db_name = re.sub(r"[^a-zA-Z0-9_]", "", base_name)
        if not db_name:
            raise HTTPException(status_code=400, detail="Invalid filename characters.")
            
        db_dest_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.db")
        if os.path.exists(db_dest_path):
            raise HTTPException(status_code=409, detail="Database dataset already uploaded. Delete it first.")
            
        # Save uploaded db to a temporary file (if provided)
        temp_db_path = None
        if db_file:
            temp_db_path = os.path.join(config.SCHEMA_DB_DIR, f"temp_{db_name}.db")
            with open(temp_db_path, "wb") as buffer:
                buffer.write(db_file.file.read())
            
        # Optional sql content
        schema_content = None
        if schema_file:
            schema_content = schema_file.file.read().decode("utf-8")
            
        # Optional image path
        temp_img_path = None
        if erd_image:
            ext = os.path.splitext(erd_image.filename)[1].lower()
            temp_img_path = os.path.join(config.SCHEMA_DB_DIR, f"temp_{db_name}{ext}")
            with open(temp_img_path, "wb") as buffer:
                buffer.write(erd_image.file.read())
                
        # Register schema db (uploaded_db_path may be None if building from SQL)
        info = schema_service.register_schema_db(
            db_name=db_name,
            schema_sql_content=schema_content,
            uploaded_db_path=temp_db_path,
            erd_image_path=temp_img_path
        )
        
        # Validation check: at least 2 tables detected after cleaning
        if info["total_tables"] < 2:
            try:
                db_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.db")
                sql_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.sql")
                json_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.json")
                for p in [db_path, sql_path, json_path]:
                    if os.path.exists(p):
                        os.remove(p)
            except Exception:
                pass
            raise HTTPException(
                status_code=400,
                detail="Validation failed: At least 2 tables must be detected in the schema."
            )
        
        # Clean up temp db and image
        if temp_db_path and os.path.exists(temp_db_path):
            try:
                os.remove(temp_db_path)
            except Exception:
                pass
        if temp_img_path and os.path.exists(temp_img_path):
            try:
                os.remove(temp_img_path)
            except Exception:
                pass
                
        return SchemaUploadResponse(
            success=True,
            db_name=db_name,
            total_tables=info["total_tables"],
            total_rows=info["total_rows"],
            tables=[t["name"] for t in info["tables"]],
            relationships=info["relationships"],
            message="Schema database registered successfully"
        )
    except HTTPException as he:
        raise he
    except Exception as e:
        print(f"[API] Error in POST /upload/schema: {e}")
        # Clean up temp files on error
        if 'temp_db_path' in locals() and temp_db_path and os.path.exists(temp_db_path):
            try:
                os.remove(temp_db_path)
            except Exception:
                pass
        if 'temp_img_path' in locals() and temp_img_path and os.path.exists(temp_img_path):
            try:
                os.remove(temp_img_path)
            except Exception:
                pass
        raise HTTPException(status_code=500, detail=str(e))

# 12. GET /schema-datasets/{db_name}/schema
@app.get("/schema-datasets/{db_name}/schema", response_model=SchemaInfoResponse)
def get_schema_details(db_name: str):
    print(f"[API] GET /schema-datasets/{db_name}/schema called")
    validate_db_name(db_name)
    
    try:
        info = schema_service.get_schema_db_info(db_name)
        return SchemaInfoResponse(
            db_name=db_name,
            tables=info["tables"],
            relationships=info["relationships"],
            total_tables=info["total_tables"],
            total_rows=info["total_rows"]
        )
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="Schema dataset not found.")
    except Exception as e:
        print(f"[API] Error in GET schema details for '{db_name}': {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 13. GET /schema-datasets/{db_name}/preview
@app.get("/schema-datasets/{db_name}/preview", response_model=SchemaPreviewsResponse)
def get_schema_preview(db_name: str):
    print(f"[API] GET /schema-datasets/{db_name}/preview called")
    validate_db_name(db_name)
    
    try:
        info = schema_service.get_schema_db_info(db_name)
        previews = []
        
        for tbl in info["tables"]:
            tbl_name = tbl["name"]
            preview_sql = f"SELECT * FROM {tbl_name} LIMIT 5"
            results = sql_service.execute_query(db_name, preview_sql, mode="schema")
            previews.append({
                "table_name": tbl_name,
                "columns": results["columns"],
                "rows": results["rows"],
                "total_rows": tbl["row_count"]
            })
            
        return SchemaPreviewsResponse(db_name=db_name, previews=previews)
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="Schema dataset not found.")
    except Exception as e:
        print(f"[API] Error in GET schema preview for '{db_name}': {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 14. POST /schema-query
@app.post("/schema-query", response_model=SchemaQueryResponse)
def run_schema_query(request: SchemaQueryRequest, response: Response):
    print(f"[API] POST /schema-query called for db: {request.db_name}, question: '{request.question}'")
    start_time = time.perf_counter()
    
    db_path = os.path.join(config.SCHEMA_DB_DIR, f"{request.db_name}.db")
    if not os.path.exists(db_path):
        raise HTTPException(status_code=404, detail=f"Database '{request.db_name}' not found.")
        
    try:
        # a. Get schema info
        schema_info = schema_service.get_schema_db_info(request.db_name)
        known_tables = [t["name"] for t in schema_info["tables"]]
        
        # b. Generate SQL
        generation = llm_service.generate_sql_schema(request.question, schema_info)
        sql = generation["sql"]
        
        # c. Validate SQL
        if not sql_service.validate_sql(sql):
            raise HTTPException(
                status_code=400,
                detail="Dangerous or invalid SQL query generated. Only SELECT operations are allowed."
            )
            
        # d. Execute query with auto-retry
        retries = 0
        results = None
        current_sql = sql
        last_error = ""
        
        while True:
            try:
                results = sql_service.execute_query(request.db_name, current_sql, mode="schema")
                break
            except Exception as e:
                last_error = str(e)
                print(f"[API] Schema Query execution failed: {last_error}")
                
                if retries < config.MAX_SQL_RETRIES:
                    retries += 1
                    print(f"[API] Retrying query fix (Attempt {retries}/{config.MAX_SQL_RETRIES})...")
                    try:
                        fix_result = llm_service.fix_sql_schema(current_sql, last_error, schema_info)
                        fixed_sql = fix_result["sql"]
                        
                        if not sql_service.validate_sql(fixed_sql):
                            raise ValueError("Fixed SQL failed validation check.")
                        current_sql = fixed_sql
                    except Exception as fix_err:
                        print(f"[API] Failed to generate/validate fix: {fix_err}")
                else:
                    print("[API] Maximum SQL retries exceeded in schema mode. Aborting.")
                    raise HTTPException(
                        status_code=422,
                        detail=f"SQL execution failed after {config.MAX_SQL_RETRIES} retries. Last error: {last_error}"
                    )
                    
        # e. Generate insight
        insight = insight_service.generate_insight(request.question, current_sql, results)
        
        # f. Process time header
        process_time_ms = (time.perf_counter() - start_time) * 1000.0
        response.headers["X-Process-Time"] = f"{process_time_ms:.2f}ms"
        
        # g. Parse tables_used
        matches = re.findall(r"\b(?:FROM|JOIN)\s+([a-zA-Z0-9_]+)", current_sql, re.IGNORECASE)
        tables_used = []
        for m in matches:
            tbl = m.lower().strip()
            if tbl in known_tables and tbl not in tables_used:
                tables_used.append(tbl)
                
        # Save to query history
        query_histories[f"schema_{request.db_name}"].append({
            "question": request.question,
            "sql": current_sql,
            "insight": insight,
            "timestamp": datetime.now().isoformat(),
            "row_count": results["row_count"]
        })
        if len(query_histories[f"schema_{request.db_name}"]) > 20:
            query_histories[f"schema_{request.db_name}"] = query_histories[f"schema_{request.db_name}"][-20:]
            
        return SchemaQueryResponse(
            question=request.question,
            sql=current_sql,
            results=SQLResults(
                columns=results["columns"],
                rows=results["rows"],
                row_count=results["row_count"],
                execution_time_ms=results["execution_time_ms"]
            ),
            insight=insight,
            tables_used=tables_used,
            retries=retries,
            db_name=request.db_name
        )
    except HTTPException as he:
        raise he
    except Exception as e:
        print(f"[API] Error in run_schema_query: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# 15. DELETE /schema-datasets/{db_name}
@app.delete("/schema-datasets/{db_name}", response_model=DeleteResponse)
def delete_schema_dataset(db_name: str):
    print(f"[API] DELETE /schema-datasets/{db_name} called")
    validate_db_name(db_name)
    
    if db_name == "ecommerce":
        raise HTTPException(status_code=403, detail="Cannot delete sample datasets.")
        
    db_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.db")
    sql_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.sql")
    
    if not os.path.exists(db_path) or not os.path.exists(sql_path):
        raise HTTPException(status_code=404, detail="Schema dataset not found.")
        
    try:
        os.remove(db_path)
        os.remove(sql_path)
        json_path = os.path.join(config.SCHEMA_DB_DIR, f"{db_name}.json")
        if os.path.exists(json_path):
            os.remove(json_path)
        return DeleteResponse(
            success=True,
            message=f"Schema dataset '{db_name}' deleted successfully."
        )
    except Exception as e:
        print(f"[API] Error deleting schema dataset '{db_name}': {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)