"""FastAPI Backend for OpenNL2SQL with Groq AI Integration
Author: Amal SP  
Created: December 2025
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
import os
import logging
from groq import Groq
import json

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="OpenNL2SQL API",
    description="AI-powered Natural Language to SQL Analytics System",
    version="1.0.0"
)

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

# Initialize Groq client
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
# Initialize Groq client
try:
    groq_client = Groq(api_key=GROQ_API_KEY) if GROQ_API_KEY else None
except Exception as e:
    logger.warning(f"Error initializing Groq: {str(e)}")
    groq_client = None

if groq_client:
    logger.info("Groq client initialized successfully")
else:
    logger.warning("GROQ_API_KEY not found - running in demo mode")
# Request/Response Models
class QueryRequest(BaseModel):
    question: str
    session_id: Optional[str] = None

class QueryResponse(BaseModel):
    success: bool
    sql: Optional[str] = None
    results: Optional[List[Dict[str, Any]]] = None
    sql_explanation: Optional[str] = None
    results_explanation: Optional[str] = None
    error: Optional[str] = None
    session_id: str

def generate_sql_with_groq(question: str) -> tuple:
    """Generate SQL using Groq AI"""
    try:
        # Return demo SQL if Groq client is not available
        if not groq_client:
            demo_sql = "SELECT c.name, SUM(o.total) as order_total FROM customers c JOIN orders o ON c.id = o.customer_id GROUP BY c.name ORDER BY order_total DESC"
            return demo_sql, None
        # Sample database schema
        schema = """
        Database Schema:
        - customers (id, name, email, created_at)
        - orders (id, customer_id, total, status, created_at)
        - products (id, name, price, category)
        - order_items (id, order_id, product_id, quantity, price)
        """
        
        prompt = f"""{schema}

Convert this natural language question to a SQL query:
Question: {question}

Generate ONLY a valid SELECT SQL query. No explanations.
SQL Query:"""

        response = groq_client.chat.completions.create(
            model="mixtral-8x7b-32768",
            messages=[
                {"role": "system", "content": "You are a SQL expert. Generate only valid SQL SELECT queries without any explanations or markdown formatting."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.2,
            max_tokens=500
        )
        
        sql = response.choices[0].message.content.strip()
        # Clean up the SQL
        sql = sql.replace("```sql", "").replace("```", "").strip()
        
        return sql, None
    except Exception as e:
        logger.error(f"Error generating SQL: {str(e)}")
        return None, str(e)

def explain_sql_with_groq(sql: str, question: str) -> str:
    """Generate explanation for SQL query"""
    try:
        # Return demo explanation if Groq client is not available
        if not groq_client:
            return "This query retrieves data from the database. Full AI explanation unavailable in demo mode."
        prompt = f"""Explain this SQL query in simple terms:
        
Original Question: {question}
SQL Query: {sql}

Provide a brief, clear explanation:"""

        response = groq_client.chat.completions.create(
            model="mixtral-8x7b-32768",
            messages=[
                {"role": "system", "content": "You are a helpful assistant that explains SQL queries in simple terms."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
            max_tokens=300
        )
        
        return response.choices[0].message.content.strip()
    except Exception as e:
        logger.error(f"Error explaining SQL: {str(e)}")
        return "SQL query generated successfully."

@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service": "OpenNL2SQL API",
        "version": "1.0.0",
        "message": "FastAPI backend with Groq AI integration running on Hugging Face Spaces!",
        "groq_enabled": groq_client is not None
    }

@app.get("/health")
async def health_check():
    """Detailed health check"""
    return {
        "status": "healthy",
        "groq_api_configured": groq_client is not None,
        "service": "OpenNL2SQL API"
    }

@app.post("/query", response_model=QueryResponse)
async def process_query(request: QueryRequest):
    """Process natural language query with Groq AI"""
    session_id = request.session_id or "demo-session"
    
    # Check if Groq is available
    # Check if Groq API key is configured
    if not GROQ_API_KEY:
        return QueryResponse(
            success=False,
            error="GROQ_API_KEY not configured. Please add it in HF Spaces Settings > Variables.",
            session_id=session_id
        )
    
        sql, error = generate_sql_with_groq(request.question)
        
        if error:
            return QueryResponse(
                success=False,
                error=f"Failed to generate SQL: {error}",
                session_id=session_id
            )
        
        # Generate explanation
        explanation = explain_sql_with_groq(sql, request.question)
        
        # For demo: return mock results
        # In production, you'd execute the SQL against a real database
        results = [
            {"info": "SQL generated successfully! In production, this would execute against your database."},
            {"note": "Connect your database to see real query results."}
        ]
        
        return QueryResponse(
            success=True,
            sql=sql,
            results=results,
            sql_explanation=explanation,
            results_explanation=f"Generated SQL query for: '{request.question}'. Ready to execute against your database.",
            session_id=session_id
        )
        

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