import gradio as gr
import requests
import pandas as pd
import base64
from io import BytesIO
import json
import time
import traceback

# 🔧 CONFIGURE: Replace with your EC2 public IP
API_URL = "http://3.14.153.135:5000/ask"
HEALTH_URL = "http://3.14.153.135:5000/health"
SIMPLE_TEST_URL = "http://3.14.153.135:5000/simple-test"
DEBUG_URL = "http://3.14.153.135:5000/debug"

def test_connection():
    try:
        print("=== TESTING CONNECTION ===")
        response = requests.get(HEALTH_URL, timeout=30)
        print(f"Connection test status: {response.status_code}")
        print(f"Connection test response: {response.text}")
        
        if response.status_code == 200:
            return f"✅ Connection successful! Status: {response.status_code}\nResponse: {response.text}"
        else:
            return f"⚠️ Connection returned status: {response.status_code}\nResponse: {response.text}"
    except Exception as e:
        error_msg = f"Connection error: {str(e)}"
        print(f"Connection error: {error_msg}")
        print(f"Traceback: {traceback.format_exc()}")
        return f"❌ {error_msg}"

def test_simple_endpoint():
    try:
        print("=== TESTING SIMPLE ENDPOINT ===")
        response = requests.post(
            SIMPLE_TEST_URL,
            json={"question": "test question"},
            timeout=30
        )
        
        print(f"Simple test status: {response.status_code}")
        print(f"Simple test response: {response.text}")
        
        if response.status_code == 200:
            result = response.json()
            return f"✅ Simple test successful! SQL: {result.get('sql', 'N/A')}", None
        else:
            return f"❌ Simple test failed: {response.status_code} - {response.text}", None
    except Exception as e:
        error_msg = f"Simple test error: {str(e)}"
        print(f"Simple test error: {error_msg}")
        print(f"Traceback: {traceback.format_exc()}")
        return f"❌ {error_msg}", None

def test_debug_endpoint():
    try:
        print("=== TESTING DEBUG ENDPOINT ===")
        response = requests.post(
            DEBUG_URL,
            json={"test": "data"},
            timeout=30
        )
        
        print(f"Debug test status: {response.status_code}")
        print(f"Debug test response: {response.text}")
        
        return f"Debug response: {response.text}", None
    except Exception as e:
        error_msg = f"Debug test error: {str(e)}"
        print(f"Debug test error: {error_msg}")
        print(f"Traceback: {traceback.format_exc()}")
        return f"❌ {error_msg}", None

def query_database(question):
    if not question.strip():
        return "", None, None, "⚠️ Please enter a question.", "", ""
    
    try:
        start_time = time.time()
        print("\n" + "="*50)
        print("=== NEW QUERY REQUEST ===")
        print(f"Time: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"Question: {question}")
        print("="*50)
        
        # Prepare the request
        headers = {
            'Content-Type': 'application/json',
            'User-Agent': 'HuggingFace-Space/1.0'
        }
        payload = {"question": question}
        
        print(f"Request URL: {API_URL}")
        print(f"Request headers: {headers}")
        print(f"Request payload: {payload}")
        
        # Make the request
        print("Sending request...")
        response = requests.post(
            API_URL, 
            json=payload, 
            headers=headers,
            timeout=300
        )
        
        elapsed_time = time.time() - start_time
        print(f"\n=== RESPONSE RECEIVED ===")
        print(f"Response time: {elapsed_time:.2f} seconds")
        print(f"Response status: {response.status_code}")
        print(f"Response headers: {dict(response.headers)}")
        print(f"Response text length: {len(response.text)} characters")
        print(f"Response text: {response.text}")
        
        # Check if response is empty
        if not response.text:
            error_msg = "Empty response from server"
            print(f"ERROR: {error_msg}")
            return "", None, None, f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "Empty response"
        
        # Try to parse JSON response
        try:
            result = response.json()
            print(f"Parsed JSON successfully: {type(result)}")
        except json.JSONDecodeError as e:
            error_msg = f"Invalid JSON response: {str(e)}"
            print(f"ERROR: {error_msg}")
            print(f"Response text (first 500 chars): {response.text[:500]}")
            return "", None, None, f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "JSON parse error"
        
        # Check HTTP status
        if response.status_code != 200:
            error_msg = result.get("error", f"HTTP {response.status_code}")
            print(f"ERROR: HTTP status {response.status_code}: {error_msg}")
            return "", None, None, f"❌ Server error: {error_msg}", f"Request time: {elapsed_time:.2f}s", "HTTP error"
        
        # Check for application error
        if "error" in result:
            error_msg = result['error']
            print(f"ERROR: Application error: {error_msg}")
            return "", None, None, f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "Application error"
        
        # Extract data
        sql = result.get("sql", "N/A")
        rows = result.get("result", [])
        chart_data = result.get("chart")
        
        print(f"\n=== EXTRACTED DATA ===")
        print(f"SQL: {sql[:100] if sql else 'None'}...")
        print(f"Rows count: {len(rows)}")
        print(f"Chart data received: {'Yes' if chart_data else 'No'}")
        print(f"Chart data length: {len(chart_data) if chart_data else 0}")
        
        # Create DataFrame
        df = pd.DataFrame(rows) if rows else pd.DataFrame()
        print(f"DataFrame shape: {df.shape}")
        
        # Process chart
        chart_image = None
        if chart_data:
            try:
                image_bytes = base64.b64decode(chart_data)
                chart_image = BytesIO(image_bytes)
                print("Chart decoded successfully")
            except Exception as e:
                print(f"Error decoding chart: {e}")
        
        # Prepare details
        details = f"Request time: {elapsed_time:.2f}s\n"
        details += f"Status code: {response.status_code}\n"
        details += f"Rows returned: {len(rows)}\n"
        details += f"Chart generated: {'Yes' if chart_data else 'No'}\n"
        details += f"Response size: {len(response.text)} bytes"
        
        print(f"=== REQUEST COMPLETED SUCCESSFULLY ===")
        
        return sql, df, chart_image, "✅ Query completed successfully!", details, "Success"
    
    except requests.exceptions.ConnectionError as e:
        error_msg = f"Connection failed: {str(e)}"
        print(f"CONNECTION ERROR: {error_msg}")
        print(f"Traceback: {traceback.format_exc()}")
        return "", None, None, f"❌ {error_msg}", "Connection error", "Connection error"
    except requests.exceptions.Timeout:
        error_msg = "Request timed out after 300 seconds"
        print(f"TIMEOUT ERROR: {error_msg}")
        return "", None, None, f"⏱️ {error_msg}", "Timeout error", "Timeout error"
    except requests.exceptions.RequestException as e:
        error_msg = f"Request exception: {str(e)}"
        print(f"REQUEST ERROR: {error_msg}")
        print(f"Traceback: {traceback.format_exc()}")
        return "", None, None, f"❌ {error_msg}", "Request error", "Request error"
    except Exception as e:
        error_msg = f"Unexpected error: {str(e)}"
        print(f"UNEXPECTED ERROR: {error_msg}")
        print(f"Traceback: {traceback.format_exc()}")
        return "", None, None, f"🚨 {error_msg}", f"Error: {str(e)}", "Unexpected error"

# 🎨 Theme: Enterprise Dark Blue
theme = gr.themes.Default(
    primary_hue="blue",
    secondary_hue="gray",
    neutral_hue="slate",
    font=["Inter", "sans-serif"]
).set(
    body_background_fill="*neutral_950",
    background_fill_secondary="*neutral_900"
)

# 🚀 UI Layout
with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
    gr.HTML("""
    <div style="text-align: center; padding: 20px;">
        <h1 style="color: #3B82F6; margin-bottom: 5px;">📊 Enterprise SQL Assistant</h1>
        <p style="color: #9CA3AF; font-size: 1.1em; margin-top: 0;">
            Ask questions about your data. Get SQL, results, and insights.
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🔍 Ask a Question")
            question_input = gr.Textbox(
                placeholder="E.g., How many members are there?",
                label="Natural Language Query",
                lines=3
            )
            submit_btn = gr.Button("🚀 Generate SQL & Results", variant="primary")
            
            gr.Markdown("### 🔧 Connection Tests")
            with gr.Accordion("Connection Tests", open=True):
                test_btn = gr.Button("Test Connection to EC2")
                connection_status = gr.Textbox(label="Connection Status", interactive=False)
                
                simple_test_btn = gr.Button("Test Simple Endpoint")
                simple_test_status = gr.Textbox(label="Simple Test Status", interactive=False)
                
                debug_test_btn = gr.Button("Test Debug Endpoint")
                debug_test_status = gr.Textbox(label="Debug Test Status", interactive=False)
            
            gr.Markdown("### 💡 Example Queries")
            gr.Markdown("""
            Try these example queries:
            
            - How many members are there?
            - What is the total transaction amount?
            - Show members with their account balances
            - Which member has the highest balance?
            - Show transaction trends over time
            - Count of members by status
            - Show me the first 10 rows
            - What is the average age of members?
            """)
        
        with gr.Column(scale=2):
            gr.Markdown("### 📄 Results")
            status = gr.Textbox(label="Status", interactive=False)
            
            with gr.Tabs():
                with gr.Tab("SQL Query"):
                    sql_output = gr.Code(label="", language="sql")
                
                with gr.Tab("Data Results"):
                    results_output = gr.Dataframe(
                        label="Query Results",
                        interactive=False,
                        wrap=True
                    )
                
                with gr.Tab("Visual Insights"):
                    chart_output = gr.Image(
                        label="Chart",
                        type="pil",
                        height=400,
                        elem_classes="chart-container"
                    )
            
            gr.Markdown("### 📊 Request Details")
            request_details = gr.Textbox(label="Request Details", interactive=False, lines=6)
            
            gr.Markdown("### 🔍 Error Details")
            error_details = gr.Textbox(label="Error Details", interactive=False, lines=4)
    
    # Events
    test_btn.click(
        fn=test_connection,
        outputs=[connection_status]
    )
    
    simple_test_btn.click(
        fn=test_simple_endpoint,
        outputs=[simple_test_status]
    )
    
    debug_test_btn.click(
        fn=test_debug_endpoint,
        outputs=[debug_test_status]
    )
    
    submit_btn.click(
        fn=query_database,
        inputs=question_input,
        outputs=[sql_output, results_output, chart_output, status, request_details, error_details]
    )

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