import gradio as gr
import time
import json

# Import RAG system components
from rag_system.vector_store import VectorStore
from rag_system.retriever import SQLRetriever
from rag_system.prompt_engine import PromptEngine
from rag_system.sql_generator import SQLGenerator

# Initialize RAG system components
print("Initializing RAG system...")
try:
    vector_store = VectorStore()
    retriever = SQLRetriever(vector_store)
    prompt_engine = PromptEngine()
    sql_generator = SQLGenerator(retriever, prompt_engine)
    print("RAG system initialized successfully!")
except Exception as e:
    print(f"Error initializing RAG system: {e}")
    sql_generator = None

def generate_sql(question, table_headers):
    """Generate SQL using the RAG system directly."""
    if sql_generator is None:
        return "❌ Error: RAG system not initialized"
    
    try:
        start_time = time.time()
        
        # Generate SQL using RAG system
        result = sql_generator.generate_sql(question, table_headers)
        
        processing_time = time.time() - start_time
        
        return f"""
**Generated SQL:**
```sql
{result['sql_query']}
```

**Model Used:** {result['model_used']}
**Processing Time:** {processing_time:.2f}s
**Status:** {result['status']}
**Retrieved Examples:** {len(result['retrieved_examples'])} examples used for RAG
"""
    except Exception as e:
        return f"❌ Error: {str(e)}"

def batch_generate_sql(questions_text, table_headers):
    """Generate SQL for multiple questions."""
    if sql_generator is None:
        return "❌ Error: RAG system not initialized"
    
    try:
        # Parse questions
        questions = [q.strip() for q in questions_text.split("\n") if q.strip()]
        
        output = f"**Batch Results:**\n"
        output += f"Total Queries: {len(questions)}\n"
        successful_count = 0
        
        for i, question in enumerate(questions):
            try:
                start_time = time.time()
                result = sql_generator.generate_sql(question, table_headers)
                processing_time = time.time() - start_time
                
                output += f"\n**Query {i+1}:** {question}\n"
                output += f"```sql\n{result['sql_query']}\n```\n"
                output += f"Model: {result['model_used']} | Time: {processing_time:.2f}s\n"
                
                if result['status'] == 'success':
                    successful_count += 1
                    
            except Exception as e:
                output += f"\n**Query {i+1}:** {question}\n"
                output += f"❌ Error: {str(e)}\n"
        
        output += f"\n**Successful:** {successful_count}/{len(questions)}"
        return output
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

def check_system_health():
    """Check the health of the RAG system."""
    try:
        if sql_generator is None:
            return "❌ System Status: RAG system not initialized"
        
        # Get model info
        model_info = sql_generator.get_model_info()
        
        return f"""
**System Health:**
- **Status:** ✅ Healthy
- **System Loaded:** ✅ Yes
- **System Loading:** ❌ No
- **Error:** None
- **Timestamp:** {time.strftime('%Y-%m-%d %H:%M:%S')}

**Model Info:**
{json.dumps(model_info, indent=2) if model_info else 'Not available'}
"""
    except Exception as e:
        return f"❌ Health check error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Text-to-SQL RAG with CodeLlama", theme=gr.themes.Soft()) as demo:
    gr.Markdown("#Text-to-SQL RAG with CodeLlama")
    gr.Markdown("Generate SQL queries from natural language using **RAG (Retrieval-Augmented Generation)** and **CodeLlama** models.")
    gr.Markdown("**Features:** RAG-enhanced generation, CodeLlama integration, Vector-based retrieval, Advanced prompt engineering")
    
    with gr.Tab("Single Query"):
        with gr.Row():
            with gr.Column(scale=1):
                question_input = gr.Textbox(
                    label="Question",
                    placeholder="e.g., Show me all employees with salary greater than 50000",
                    lines=3
                )
                table_headers_input = gr.Textbox(
                    label="Table Headers (comma-separated)",
                    placeholder="e.g., id, name, salary, department",
                    value="id, name, salary, department"
                )
                generate_btn = gr.Button("Generate SQL", variant="primary", size="lg")
            
            with gr.Column(scale=1):
                output = gr.Markdown(label="Result")
    
    with gr.Tab("Batch Queries"):
        with gr.Row():
            with gr.Column(scale=1):
                batch_questions = gr.Textbox(
                    label="Questions (one per line)",
                    placeholder="Show me all employees\nCount total employees\nAverage salary by department",
                    lines=5
                )
                batch_headers = gr.Textbox(
                    label="Table Headers (comma-separated)",
                    placeholder="e.g., id, name, salary, department",
                    value="id, name, salary, department"
                )
                batch_btn = gr.Button("Generate Batch SQL", variant="primary", size="lg")
            
            with gr.Column(scale=1):
                batch_output = gr.Markdown(label="Batch Results")
    
    with gr.Tab("System Health"):
        with gr.Row():
            health_btn = gr.Button("Check System Health", variant="secondary", size="lg")
            health_output = gr.Markdown(label="Health Status")
    
    # Event handlers
    generate_btn.click(
        generate_sql,
        inputs=[question_input, table_headers_input],
        outputs=output
    )
    
    batch_btn.click(
        batch_generate_sql,
        inputs=[batch_questions, batch_headers],
        outputs=batch_output
    )
    
    health_btn.click(
        check_system_health,
        outputs=health_output
    )
    
    gr.Markdown("---")
    gr.Markdown("""
    ## How It Works
    
    1. **RAG System**: Retrieves relevant SQL examples from vector database
    2. **CodeLlama**: Generates SQL using retrieved examples as context
    3. **Vector Search**: Finds similar questions and their SQL solutions
    4. **Enhanced Generation**: Combines retrieval + generation for better accuracy
    
    ## Technology Stack
    
    - **Backend**: Direct RAG system integration
    - **LLM**: CodeLlama-7B-Python-GGUF (primary)
    - **Vector DB**: ChromaDB with sentence transformers
    - **Frontend**: Gradio interface
    - **Hosting**: Hugging Face Spaces
    
    ## Performance
    
    - **Model**: CodeLlama-7B-Python-GGUF
    - **Response Time**: < 5 seconds
    - **Accuracy**: High (RAG-enhanced)
    - **Cost**: Free (local inference)
    """)

# Launch the interface
if __name__ == "__main__":
    demo.launch()