app.py

import sys
import os
import uvicorn
import re
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from typing import List, Optional

# 🚨 FORCE PYTHON TO FIND THE 'src' FOLDER
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Import custom modules
from src.db_connector import Database
from src.rag_manager import RAGSystem 
from src.sql_generator import SQLGenerator

app = FastAPI()

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

class ChatRequest(BaseModel):
    question: str
    history: Optional[List[dict]] = [] 

# --- HELPER: CLEAN AI OUTPUT ---
def clean_sql(sql_text: str) -> str:
    if not sql_text: return ""
    cleaned = re.sub(r"```sql|```", "", sql_text, flags=re.IGNORECASE).strip()
    return cleaned.rstrip(';')

# --- 🚀 SAFE INITIALIZATION ---
# Initialize as None so the app doesn't crash if DB fails
db = None
rag = None
generator = None
startup_error = None

print("--- 🚀 SYSTEM STARTUP SEQUENCE ---")
try:
    print("   ...Connecting to Database")
    # ⚠️ IF THIS FAILS, THE APP WILL NOW CATCH IT GRACEFULLY
    db = Database() 
    print("   ✅ Database Connection: SUCCESS")
    
    print("   ...Initializing RAG System")
    rag = RAGSystem(db)
    print("   ✅ RAG System: ONLINE")
    
    print("   ...Loading AI Model")
    generator = SQLGenerator()
    print("   ✅ AI Model: LOADED")
    
except Exception as e:
    startup_error = str(e)
    print(f"   ❌ CRITICAL STARTUP ERROR: {e}")

# --- ROUTES ---
@app.get("/", response_class=HTMLResponse)
async def serve_ui():
    try:
        with open("index.html", "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        return "Error: index.html not found."

@app.post("/chat")
def chat_endpoint(request: ChatRequest):
    # 🚨 CHECK FOR STARTUP ERRORS FIRST
    if rag is None or db is None:
        return {
            "answer": [f"System Error: {startup_error}"],
            "sql": "-- Database Connection Failed",
            "message": f"I cannot connect to the database. Reason: {startup_error}. Please check the Logs tab.",
            "follow_ups": []
        }

    try:
        context = rag.get_relevant_schema(request.question)
        raw_sql, explanation, friendly_msg = generator.generate_sql(request.question, context, request.history)
        
        # Clean SQL
        cleaned_sql = clean_sql(raw_sql)
        
        # Security Check
        if not cleaned_sql.upper().startswith("SELECT"):
             return {"answer": [], "sql": cleaned_sql, "message": "Security Alert: Read-Only Mode.", "follow_ups": []}

        # Run Query
        try:
            results = db.run_query(cleaned_sql)
        except Exception as db_err:
            return {"answer": [f"SQL Error: {str(db_err)}"], "sql": cleaned_sql, "message": "Syntax Error in SQL.", "follow_ups": []}
        
        return {
            "answer": results,
            "sql": cleaned_sql,
            "message": friendly_msg, 
            "follow_ups": []
        }

    except Exception as e:
        return {"answer": [], "sql": "-- Error", "message": f"Processing Error: {str(e)}", "follow_ups": []}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)