server.py

# Warning: This project may not work due to Hugging Face restrictions. Please check out the GitHub repo for the latest updates.

import gradio as gr
import pandas as pd
import os
import json
import re
from typing import Optional, Tuple, Dict, Any, List
import traceback
from datetime import datetime
import time

# Database imports
import mysql.connector
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.exc import SQLAlchemyError

# LangChain imports
from langchain_community.agent_toolkits.sql.base import create_sql_agent
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain.agents.agent_types import AgentType
from langchain_community.callbacks.manager import get_openai_callback
from langchain_google_genai import ChatGoogleGenerativeAI

# Environment setup
from dotenv import load_dotenv
load_dotenv()

class DatabaseManager:
    def __init__(self):
        self.db_connection = None
        self.db_context = None
        self.sql_agent = None
        self.connection_status = "Not Connected"
        self.db_type = None
        self.query_history = []  # Store query history
        self.max_history_items = 20  # Maximum number of history items to keep
        self.user_api_key = None  # Store user-provided API key
        
    def set_api_key(self, api_key: str) -> str:
        """Set user-provided API key"""
        if not api_key or not api_key.strip():
            self.user_api_key = None
            return "❌ API key cleared. Using environment variable if available."
        
        # Store the API key
        self.user_api_key = api_key.strip()
        return "✅ API key set successfully!"
    
    def get_api_key(self) -> str:
        """Get API key with priority to user-provided key"""
        if self.user_api_key:
            return self.user_api_key
        return os.getenv("GOOGLE_API_KEY", "")
        
    def connect_mysql(self, host: str, port: str, username: str, password: str, database: str) -> Tuple[str, str]:
        """Connect to MySQL database"""
        try:
            # Clean and validate inputs
            host = host.strip() if host else "localhost"
            port_num = int(port.strip()) if port and port.strip() else 3306
            username = username.strip() if username else ""
            password = str(password) if password else ""  # Ensure password is treated as string
            database = database.strip() if database else ""
            
            if not username or not database:
                return "❌ Missing required fields", "Please provide username and database name."
            
            # Test connection first with mysql.connector
            # Using raw credentials without URL encoding for direct connection
            conn = mysql.connector.connect(
                host=host,
                port=port_num,
                user=username,
                password=password,
                database=database,
                autocommit=True
            )
            conn.close()
            
            # Create SQLAlchemy engine with proper URL encoding
            from urllib.parse import quote_plus
            # Make sure to properly encode all special characters in password
            encoded_password = quote_plus(str(password))
            encoded_username = quote_plus(username)
            encoded_database = quote_plus(database)
            
            # Add binary_prefix=true to handle binary data warnings
            connection_string = f"mysql+pymysql://{encoded_username}:{encoded_password}@{host}:{port_num}/{encoded_database}?binary_prefix=true"
            engine = create_engine(connection_string, echo=False)
            
            # Test SQLAlchemy connection
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            
            # Create LangChain SQLDatabase
            self.db_connection = SQLDatabase(engine)
            self.db_type = "MySQL"
            self.connection_status = f"✅ Connected to MySQL: {host}:{port_num}/{database}"
            
            return self.connection_status, "Connection successful! You can now analyze the database."
            
        except Exception as e:
            error_msg = f"❌ MySQL Connection Failed: {str(e)}"
            self.connection_status = "Not Connected"
            return error_msg, f"Connection failed. Please check your credentials.\nError details: {str(e)}"
    
    def validate_sql_query(self, sql_query: str) -> Tuple[bool, str]:
        """
        Validate SQL query for common errors and security issues
        
        Args:
            sql_query: SQL query string to validate
            
        Returns:
            Tuple of (is_valid, message)
        """
        if not sql_query or not isinstance(sql_query, str):
            return False, "Invalid or empty SQL query"
            
        sql_query = sql_query.strip()
        
        # Check for basic SQL injection patterns
        dangerous_patterns = [
            "DROP TABLE", "DROP DATABASE", "DELETE FROM", "TRUNCATE TABLE",
            "ALTER TABLE", "UPDATE", "INSERT INTO", "CREATE TABLE", "GRANT",
            "REVOKE", "--", ";--", ";", "/*", "*/"
        ]
        
        for pattern in dangerous_patterns:
            if pattern.upper() in sql_query.upper():
                return False, f"Potentially harmful SQL detected: {pattern}"
        
        # Check for common SQL errors
        common_errors = [
            # NOT IN with NULL values
            (r"NOT\s+IN.*NULL", "Using NOT IN with NULL values can lead to unexpected results"),
            # BETWEEN for exclusive ranges
            (r"BETWEEN.*AND", "Check BETWEEN usage for correct inclusive/exclusive ranges"),
            # Potential data type mismatches
            (r"CAST\(|CONVERT\(", "Verify data type casting is correct"),
            # Potential quoting issues
            (r"[^']'[^']|[^']'$", "Check for proper quoting of identifiers")
        ]
        
        import re
        for pattern, message in common_errors:
            if re.search(pattern, sql_query, re.IGNORECASE):
                # This is just a warning, not an error
                return True, f"Warning: {message}"
        
        # Check for SELECT statement
        if not sql_query.upper().startswith("SELECT"):
            return False, "Only SELECT queries are allowed"
            
        return True, "Query validation passed"
    
    def fix_sql_query(self, sql_query: str, error_message: str, db_schema: Optional[dict] = None) -> str:
        """
        Use LLM to fix an invalid SQL query
        
        Args:
            sql_query: The original invalid SQL query
            error_message: The error message from validation or execution
            db_schema: Optional database schema information to help with correction
            
        Returns:
            Corrected SQL query
        """
        api_key = self.get_api_key()
        if not api_key:
            raise ValueError("No API key available. Please set a Google API key.")
            
        # Initialize LLM
        llm = ChatGoogleGenerativeAI(
            model="gemini-2.5-flash-preview-05-20",
            temperature=0,
            google_api_key=api_key
        )
        
        # Prepare schema information if available
        schema_info = ""
        if db_schema and isinstance(db_schema, dict):
            schema_info = "Database schema information:\n"
            for table, info in db_schema.items():
                schema_info += f"Table: {table}\n"
                if "columns" in info:
                    schema_info += "Columns:\n"
                    for col in info["columns"]:
                        schema_info += f"- {col['name']} ({col['type']})\n"
                schema_info += "\n"
        
        # Build prompt for the LLM
        prompt = f"""
        Fix the following SQL query that has errors:
        
        ```sql
        {sql_query}
        ```
        
        Error message:
        {error_message}
        
        {schema_info}
        
        Please provide ONLY the corrected SQL query with no additional text or explanation.
        The query should be a valid SELECT statement.
        """
        
        # Get the corrected query
        try:
            response = llm.invoke(prompt)
            corrected_query = response.content
            
            # Extract SQL from response if needed
            if "```sql" in corrected_query:
                corrected_query = corrected_query.split("```sql")[1].split("```")[0].strip()
            elif "```" in corrected_query:
                corrected_query = corrected_query.split("```")[1].strip()
                
            return corrected_query
        except Exception as e:
            # If correction fails, return the original query
            return sql_query
    
    def analyze_database(self) -> Tuple[str, str]:
        """Analyze database structure and create context"""
        if not self.db_connection:
            return "❌ No database connection", "Please connect to a database first."
        
        try:
            # Get database schema information
            inspector = inspect(self.db_connection._engine)
            tables = inspector.get_table_names()
            
            context_info = {
                "database_type": self.db_type,
                "total_tables": len(tables),
                "tables": {},
                "analysis_timestamp": datetime.now().isoformat()
            }
            
            # Analyze each table
            for table in tables[:10]:  # Limit to first 10 tables for performance
                try:
                    columns = inspector.get_columns(table)
                    primary_keys = inspector.get_pk_constraint(table)
                    foreign_keys = inspector.get_foreign_keys(table)
                    
                    # Get sample data count
                    with self.db_connection._engine.connect() as conn:
                        result = conn.execute(text(f"SELECT COUNT(*) FROM {table}"))
                        row_count = result.scalar()
                    
                    context_info["tables"][table] = {
                        "columns": [{"name": col["name"], "type": str(col["type"])} for col in columns],
                        "primary_keys": primary_keys["constrained_columns"] if primary_keys else [],
                        "foreign_keys": [{"columns": fk["constrained_columns"], "refers_to": f"{fk['referred_table']}.{fk['referred_columns']}"} for fk in foreign_keys],
                        "row_count": row_count
                    }
                except Exception as table_error:
                    context_info["tables"][table] = {"error": str(table_error)}
            
            self.db_context = context_info
            
            # Initialize Gemini LLM
            api_key = self.get_api_key()
            if not api_key:
                return "❌ Analysis Failed", "Please set a Google API key in the settings or environment variables"
            
            llm = ChatGoogleGenerativeAI(
                model="gemini-2.5-flash-preview-05-20",
                temperature=0,
                google_api_key=api_key
            )
            
            # Create SQL agent
            toolkit = SQLDatabaseToolkit(db=self.db_connection, llm=llm)
            self.sql_agent = create_sql_agent(
                llm=llm,
                toolkit=toolkit,
                agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                verbose=True,
                handle_parsing_errors=True
            )
            
            summary = f"""
            ✅ Database Analysis Complete!
            
            📊 Database: {self.db_type}
            📋 Tables Found: {len(tables)}
            🔍 Analyzed Tables: {min(len(tables), 10)}
            
            Ready for natural language queries!
            """
            
            detailed_info = json.dumps(context_info, indent=2)
            return summary, f"Database context saved. You can now ask questions!\n\nDetailed Analysis:\n{detailed_info}"
            
        except Exception as e:
            error_msg = f"❌ Analysis Failed: {str(e)}"
            return error_msg, f"Error during analysis: {traceback.format_exc()}"
    
    def query_database(self, question: str) -> Tuple[str, str]:
        """Process natural language query and return results"""
        if not self.sql_agent:
            # Convert error to table format with clean RDBMS style
            df_error = pd.DataFrame({"Message": ["Please connect and analyze database first."]})
            table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)
            return "❌ Not Ready", table_html
        
        if not question.strip():
            # Convert error to table format with clean RDBMS style
            df_error = pd.DataFrame({"Message": ["Please enter a question about your database."]})
            table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)
            return "❌ Empty Query", table_html
        
        try:
            # Track query start time for overall performance
            start_time = time.time()
            
            # Process the query with the agent
            result = self.sql_agent.run(question)
            
            # Try to extract and execute the SQL query for tabular display
            try:
                # Look for SQL in the result
                if "SELECT" in result.upper():
                    # Extract SQL query (this is a simple extraction, could be improved)
                    lines = result.split('\n')
                    sql_lines = [line for line in lines if 'SELECT' in line.upper()]
                    
                    if sql_lines:
                        sql_query = sql_lines[0].strip()
                        # Clean up the SQL query
                        sql_query = sql_query.replace('sql', '').replace('```', '').strip()
                        
                        # Validate the SQL query before execution
                        is_valid, validation_message = self.validate_sql_query(sql_query)
                        
                        # If query is invalid, try to fix it
                        correction_applied = False
                        if not is_valid:
                            # Get schema information for the correction agent
                            schema_info = self.db_context["tables"] if self.db_context else None
                            
                            # Try to fix the query
                            corrected_query = self.fix_sql_query(sql_query, validation_message, schema_info)
                            
                            # Validate the corrected query
                            is_valid_corrected, validation_message_corrected = self.validate_sql_query(corrected_query)
                            
                            if is_valid_corrected:
                                sql_query = corrected_query
                                validation_message = validation_message_corrected
                                correction_applied = True
                                is_valid = True
                            else:
                                # If correction also failed, return both errors in table format
                                error_msg = f"The generated SQL query failed validation: {validation_message}\n\nAttempted correction also failed: {validation_message_corrected}\n\nOriginal result:\n{result}"
                                df_error = pd.DataFrame({"Error": [error_msg]})
                                table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)
                                return "❌ Query Validation Failed", table_html
                        
                        # If there's a warning but query is valid, add it to the result
                        warning_message = ""
                        if validation_message.startswith("Warning:"):
                            warning_message = f"\n\n⚠️ {validation_message}"
                        
                        # Add correction notice if applicable
                        if correction_applied:
                            warning_message += f"\n\n🔧 Query was automatically corrected. Original query had issues: {validation_message}"
                        
                        # Execute the query to get structured data
                        try:
                            # Measure query performance
                            performance_metrics = self.measure_query_performance(sql_query)
                            
                            if performance_metrics.get("success", False):
                                # Get the data from the metrics
                                with self.db_connection._engine.connect() as conn:
                                    df = pd.read_sql(sql_query, conn)
                                
                                # Calculate overall processing time
                                total_time_ms = round((time.time() - start_time) * 1000, 2)
                                
                                # Add query to history
                                history_item = {
                                    "question": question,
                                    "sql_query": sql_query,
                                    "execution_time_ms": performance_metrics["execution_time_ms"],
                                    "total_time_ms": total_time_ms,
                                    "row_count": performance_metrics["row_count"],
                                    "complexity": performance_metrics["complexity"]["level"],
                                    "timestamp": datetime.now().isoformat()
                                }
                                self.add_to_query_history(history_item)
                                
                                # Generate performance and complexity insights
                                complexity = performance_metrics["complexity"]
                                perf_insights = f"\n\n📊 Query Metrics:\n"
                                perf_insights += f"• Execution time: {performance_metrics['execution_time_ms']}ms\n"
                                perf_insights += f"• Total processing time: {total_time_ms}ms\n"
                                perf_insights += f"• Rows returned: {performance_metrics['row_count']}\n"
                                perf_insights += f"• Complexity: {complexity['level']}\n"
                                
                                if complexity["insights"]:
                                    perf_insights += "\n🔍 Insights:\n"
                                    for insight in complexity["insights"]:
                                        perf_insights += f"• {insight}\n"
                                
                                if not df.empty:
                                    # Format table in RDBMS style
                                    table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0)
                                    # Add custom styling to make it look more like RDBMS output
                                    table_html = f"""
                                    <style>
                                    .table-bordered {{
                                        border-collapse: collapse;
                                        width: 100%;
                                        font-family: 'Courier New', Courier, monospace;
                                    }}
                                    .table-bordered th {{
                                        background-color: #f2f2f2;
                                        color: #333;
                                        font-weight: bold;
                                        text-align: left;
                                        padding: 8px;
                                        border: 1px solid #ddd;
                                    }}
                                    .table-bordered td {{
                                        padding: 8px;
                                        border: 1px solid #ddd;
                                    }}
                                    .table-striped tbody tr:nth-of-type(odd) {{
                                        background-color: rgba(0,0,0,.05);
                                    }}
                                    </style>
                                    {table_html}
                                    """
                                    return f"✅ Query Successful\n\n{result}{warning_message}{perf_insights}", table_html
                            else:
                                # If performance measurement failed, continue with normal execution
                                with self.db_connection._engine.connect() as conn:
                                    df = pd.read_sql(sql_query, conn)
                                    
                                if not df.empty:
                                    # Format table in RDBMS style
                                    table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0)
                                    # Add custom styling
                                    table_html = f"""
                                    <style>
                                    .table-bordered {{
                                        border-collapse: collapse;
                                        width: 100%;
                                        font-family: 'Courier New', Courier, monospace;
                                    }}
                                    .table-bordered th {{
                                        background-color: #f2f2f2;
                                        color: #333;
                                        font-weight: bold;
                                        text-align: left;
                                        padding: 8px;
                                        border: 1px solid #ddd;
                                    }}
                                    .table-bordered td {{
                                        padding: 8px;
                                        border: 1px solid #ddd;
                                    }}
                                    .table-striped tbody tr:nth-of-type(odd) {{
                                        background-color: rgba(0,0,0,.05);
                                    }}
                                    </style>
                                    {table_html}
                                    """
                                    return f"✅ Query Successful\n\n{result}{warning_message}", table_html
                                    
                        except SQLAlchemyError as exec_error:
                            # If execution fails, try to fix the query again with the specific error
                            if not correction_applied:
                                schema_info = self.db_context["tables"] if self.db_context else None
                                corrected_query = self.fix_sql_query(sql_query, str(exec_error), schema_info)
                                
                                # Try executing the corrected query
                                try:
                                    # Measure performance of corrected query
                                    performance_metrics = self.measure_query_performance(corrected_query)
                                    
                                    if performance_metrics.get("success", False):
                                        # Get the data from the metrics
                                        with self.db_connection._engine.connect() as conn:
                                            df = pd.read_sql(corrected_query, conn)
                                        
                                        # Calculate overall processing time
                                        total_time_ms = round((time.time() - start_time) * 1000, 2)
                                        
                                        # Add query to history
                                        history_item = {
                                            "question": question,
                                            "sql_query": corrected_query,
                                            "execution_time_ms": performance_metrics["execution_time_ms"],
                                            "total_time_ms": total_time_ms,
                                            "row_count": performance_metrics["row_count"],
                                            "complexity": performance_metrics["complexity"]["level"],
                                            "timestamp": datetime.now().isoformat(),
                                            "corrected": True,
                                            "original_query": sql_query
                                        }
                                        self.add_to_query_history(history_item)
                                        
                                        # Generate performance and complexity insights
                                        complexity = performance_metrics["complexity"]
                                        perf_insights = f"\n\n📊 Query Metrics:\n"
                                        perf_insights += f"• Execution time: {performance_metrics['execution_time_ms']}ms\n"
                                        perf_insights += f"• Total processing time: {total_time_ms}ms\n"
                                        perf_insights += f"• Rows returned: {performance_metrics['row_count']}\n"
                                        perf_insights += f"• Complexity: {complexity['level']}\n"
                                        
                                        if complexity["insights"]:
                                            perf_insights += "\n🔍 Insights:\n"
                                            for insight in complexity["insights"]:
                                                perf_insights += f"• {insight}\n"
                                        
                                        if not df.empty:
                                            # Format table in RDBMS style
                                            table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0)
                                            # Add custom styling
                                            table_html = f"""
                                            <style>
                                            .table-bordered {{
                                                border-collapse: collapse;
                                                width: 100%;
                                                font-family: 'Courier New', Courier, monospace;
                                            }}
                                            .table-bordered th {{
                                                background-color: #f2f2f2;
                                                color: #333;
                                                font-weight: bold;
                                                text-align: left;
                                                padding: 8px;
                                                border: 1px solid #ddd;
                                            }}
                                            .table-bordered td {{
                                                padding: 8px;
                                                border: 1px solid #ddd;
                                            }}
                                            .table-striped tbody tr:nth-of-type(odd) {{
                                                background-color: rgba(0,0,0,.05);
                                            }}
                                            </style>
                                            {table_html}
                                            """
                                            return f"✅ Query Successful (after correction)\n\n{result}\n\n🔧 Query was automatically corrected due to execution error: {str(exec_error)}{perf_insights}", table_html
                                    else:
                                        # If performance measurement failed, continue with normal execution
                                        with self.db_connection._engine.connect() as conn:
                                            df = pd.read_sql(corrected_query, conn)
                                        
                                        if not df.empty:
                                            # Format table in RDBMS style
                                            table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0)
                                            # Add custom styling
                                            table_html = f"""
                                            <style>
                                            .table-bordered {{
                                                border-collapse: collapse;
                                                width: 100%;
                                                font-family: 'Courier New', Courier, monospace;
                                            }}
                                            .table-bordered th {{
                                                background-color: #f2f2f2;
                                                color: #333;
                                                font-weight: bold;
                                                text-align: left;
                                                padding: 8px;
                                                border: 1px solid #ddd;
                                            }}
                                            .table-bordered td {{
                                                padding: 8px;
                                                border: 1px solid #ddd;
                                            }}
                                            .table-striped tbody tr:nth-of-type(odd) {{
                                                background-color: rgba(0,0,0,.05);
                                            }}
                                            </style>
                                            {table_html}
                                            """
                                            return f"✅ Query Successful (after correction)\n\n{result}\n\n🔧 Query was automatically corrected due to execution error: {str(exec_error)}", table_html
                                except Exception:
                                    # If correction fails, return the original error
                                    pass
                            
                            # Return the execution error in table format
                            error_msg = f"The query failed to execute:\n\n{str(exec_error)}\n\nOriginal result:\n{result}"
                            df_error = pd.DataFrame({"Error": [error_msg]})
                            table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)
                            # Add custom styling
                            table_html = f"""
                            <style>
                            .table-bordered {{
                                border-collapse: collapse;
                                width: 100%;
                                font-family: 'Courier New', Courier, monospace;
                            }}
                            .table-bordered th {{
                                background-color: #f2f2f2;
                                color: #333;
                                font-weight: bold;
                                text-align: left;
                                padding: 8px;
                                border: 1px solid #ddd;
                            }}
                            .table-bordered td {{
                                padding: 8px;
                                border: 1px solid #ddd;
                            }}
                            .table-striped tbody tr:nth-of-type(odd) {{
                                background-color: rgba(0,0,0,.05);
                            }}
                            </style>
                            {table_html}
                            """
                            return "❌ SQL Execution Error", table_html
                        
            except SQLAlchemyError as sql_error:
                # Handle SQL execution errors
                error_details = str(sql_error)
                error_msg = f"❌ SQL Execution Error"
                details = f"The query failed to execute:\n\n{error_details}\n\nOriginal result:\n{result}"
                df_error = pd.DataFrame({"Error": [details]})
                table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)
                # Add custom styling
                table_html = f"""
                <style>
                .table-bordered {{
                    border-collapse: collapse;
                    width: 100%;
                    font-family: 'Courier New', Courier, monospace;
                }}
                .table-bordered th {{
                    background-color: #f2f2f2;
                    color: #333;
                    font-weight: bold;
                    text-align: left;
                    padding: 8px;
                    border: 1px solid #ddd;
                }}
                .table-bordered td {{
                    padding: 8px;
                    border: 1px solid #ddd;
                }}
                .table-striped tbody tr:nth-of-type(odd) {{
                    background-color: rgba(0,0,0,.05);
                }}
                </style>
                {table_html}
                """
                return error_msg, table_html
            except Exception as table_error:
                # If table extraction fails, just return the text result
                pass
            
            # If we got here, we just have the text result without structured data
            # Convert to table format with RDBMS style
            df_text = pd.DataFrame({"Result": [result]})
            table_html = df_text.to_html(index=False, classes="table table-bordered table-striped", border=0)
            # Add custom styling
            table_html = f"""
            <style>
            .table-bordered {{
                border-collapse: collapse;
                width: 100%;
                font-family: 'Courier New', Courier, monospace;
            }}
            .table-bordered th {{
                background-color: #f2f2f2;
                color: #333;
                font-weight: bold;
                text-align: left;
                padding: 8px;
                border: 1px solid #ddd;
            }}
            .table-bordered td {{
                padding: 8px;
                border: 1px solid #ddd;
            }}
            .table-striped tbody tr:nth-of-type(odd) {{
                background-color: rgba(0,0,0,.05);
            }}
            </style>
            {table_html}
            """
            
            # Add to history
            history_item = {
                "question": question,
                "result": result,
                "timestamp": datetime.now().isoformat()
            }
            self.add_to_query_history(history_item)
            
            return f"✅ Query Successful", table_html
            
        except Exception as e:
            # Convert exception to table format with RDBMS style
            error_msg = f"❌ Query Failed: {str(e)}"
            details = f"Error processing query: {traceback.format_exc()}"
            df_error = pd.DataFrame({"Error": [details]})
            table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)
            # Add custom styling
            table_html = f"""
            <style>
            .table-bordered {{
                border-collapse: collapse;
                width: 100%;
                font-family: 'Courier New', Courier, monospace;
            }}
            .table-bordered th {{
                background-color: #f2f2f2;
                color: #333;
                font-weight: bold;
                text-align: left;
                padding: 8px;
                border: 1px solid #ddd;
            }}
            .table-bordered td {{
                padding: 8px;
                border: 1px solid #ddd;
            }}
            .table-striped tbody tr:nth-of-type(odd) {{
                background-color: rgba(0,0,0,.05);
            }}
            </style>
            {table_html}
            """
            return error_msg, table_html

    def analyze_query_complexity(self, sql_query: str) -> Dict[str, Any]:
        """
        Analyze SQL query complexity and provide insights
        
        Args:
            sql_query: SQL query to analyze
            
        Returns:
            Dictionary with complexity metrics and insights
        """
        if not sql_query or not isinstance(sql_query, str):
            return {"error": "Invalid query provided"}
            
        sql_query = sql_query.strip().upper()
        
        # Initialize complexity metrics
        complexity = {
            "level": "Simple",
            "score": 0,
            "joins": 0,
            "tables": [],
            "aggregations": False,
            "grouping": False,
            "ordering": False,
            "limiting": False,
            "subqueries": 0,
            "complex_functions": [],
            "insights": []
        }
        
        # Count number of JOINs
        join_count = len(re.findall(r'\bJOIN\b', sql_query))
        complexity["joins"] = join_count
        if join_count > 0:
            complexity["score"] += join_count * 2
            if join_count >= 3:
                complexity["insights"].append(f"Query uses {join_count} joins, which may impact performance")
        
        # Detect tables used
        from_clause = re.search(r'\bFROM\b\s+(.*?)(?:\bWHERE\b|\bGROUP\b|\bHAVING\b|\bORDER\b|\bLIMIT\b|$)', sql_query)
        if from_clause:
            # Extract table names from FROM clause
            tables_text = from_clause.group(1).strip()
            # Handle JOIN syntax in FROM clause
            tables = re.findall(r'([a-zA-Z0-9_]+)(?:\s+(?:AS\s+)?[a-zA-Z0-9_]+)?', tables_text)
            complexity["tables"] = list(set(tables))  # Remove duplicates
        
        # Check for aggregations
        agg_functions = ["COUNT", "SUM", "AVG", "MIN", "MAX"]
        for func in agg_functions:
            if re.search(rf'\b{func}\s*\(', sql_query):
                complexity["aggregations"] = True
                complexity["score"] += 1
                break
        
        # Check for GROUP BY
        if re.search(r'\bGROUP\s+BY\b', sql_query):
            complexity["grouping"] = True
            complexity["score"] += 2
        
        # Check for ORDER BY
        if re.search(r'\bORDER\s+BY\b', sql_query):
            complexity["ordering"] = True
            complexity["score"] += 1
        
        # Check for LIMIT
        if re.search(r'\bLIMIT\b', sql_query):
            complexity["limiting"] = True
            complexity["score"] += 0.5
        
        # Check for subqueries
        subquery_count = len(re.findall(r'\(\s*SELECT', sql_query))
        complexity["subqueries"] = subquery_count
        if subquery_count > 0:
            complexity["score"] += subquery_count * 3
            complexity["insights"].append(f"Query contains {subquery_count} subqueries, which may affect performance")
        
        # Check for complex functions
        complex_funcs = ["CASE", "COALESCE", "NULLIF", "CAST", "CONVERT", "SUBSTRING", "CONCAT", "DATE_FORMAT", "EXTRACT"]
        for func in complex_funcs:
            if re.search(rf'\b{func}\b', sql_query):
                complexity["complex_functions"].append(func)
                complexity["score"] += 1
        
        # Determine complexity level
        if complexity["score"] <= 2:
            complexity["level"] = "Simple"
        elif complexity["score"] <= 5:
            complexity["level"] = "Moderate"
        elif complexity["score"] <= 10:
            complexity["level"] = "Complex"
        else:
            complexity["level"] = "Very Complex"
            complexity["insights"].append("This is a highly complex query that may benefit from optimization")
        
        # Add insights based on complexity
        if complexity["level"] in ["Complex", "Very Complex"] and not complexity["limiting"]:
            complexity["insights"].append("Consider adding a LIMIT clause to prevent large result sets")
        
        if complexity["joins"] >= 2 and not any(idx for idx in complexity["insights"] if "index" in idx.lower()):
            complexity["insights"].append("Ensure proper indexes exist on join columns")
        
        return complexity

    def add_to_query_history(self, query_data: Dict[str, Any]) -> None:
        """
        Add a query to the history
        
        Args:
            query_data: Dictionary containing query information
        """
        # Add timestamp if not present
        if "timestamp" not in query_data:
            query_data["timestamp"] = datetime.now().isoformat()
            
        # Add to history (at the beginning for most recent first)
        self.query_history.insert(0, query_data)
        
        # Trim history if needed
        if len(self.query_history) > self.max_history_items:
            self.query_history = self.query_history[:self.max_history_items]
    
    def get_query_history(self) -> List[Dict[str, Any]]:
        """
        Get the query history
        
        Returns:
            List of query history items
        """
        return self.query_history
    
    def clear_query_history(self) -> None:
        """Clear the query history"""
        self.query_history = []
        
    def measure_query_performance(self, sql_query: str) -> Dict[str, Any]:
        """
        Measure the performance of a SQL query
        
        Args:
            sql_query: SQL query to execute and measure
            
        Returns:
            Dictionary with performance metrics
        """
        if not self.db_connection:
            return {"error": "No database connection"}
            
        metrics = {
            "query": sql_query,
            "execution_time_ms": 0,
            "row_count": 0,
            "success": False,
            "error": None
        }
        
        try:
            # Measure execution time
            start_time = time.time()
            
            with self.db_connection._engine.connect() as conn:
                result = conn.execute(text(sql_query))
                # Convert to DataFrame to get row count
                df = pd.DataFrame(result.fetchall(), columns=result.keys())
                
            end_time = time.time()
            
            # Calculate metrics
            metrics["execution_time_ms"] = round((end_time - start_time) * 1000, 2)
            metrics["row_count"] = len(df)
            metrics["success"] = True
            
            # Add complexity analysis
            metrics["complexity"] = self.analyze_query_complexity(sql_query)
            
            return metrics
            
        except Exception as e:
            metrics["error"] = str(e)
            return metrics

    def generate_schema_diagram(self, include_all_tables: bool = False) -> str:
        """
        Generate a Mermaid ER diagram for the database schema
        
        Args:
            include_all_tables: Whether to include all tables or just a subset
            
        Returns:
            Mermaid diagram code
        """
        # Return a message that this functionality is not available
        return "This functionality has been removed"

# Initialize the database manager
db_manager = DatabaseManager()

def create_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(title="AI Database Query Assistant", theme=gr.themes.Soft()) as demo:
        # Warning banner at the top
        gr.Markdown("""
        <div style="background-color: #FFF3CD; color: #856404; padding: 15px; border-radius: 5px; border: 1px solid #FFEEBA; margin-bottom: 20px; font-weight: bold; text-align: center;">
        ⚠️ WARNING: This project may not work due to Hugging Face restrictions. Please check out the GitHub repo for the latest updates. https://github.com/yash-8923/gradio.git
        </div>
        """)
        
        gr.Markdown("""
        # 🤖 AI Database Query Assistant
        
        Connect to your MySQL database and query it using natural language!
        
        ### Steps:
        1. **Connect** to your database
        2. **Analyze** your database structure  
        3. **Ask questions** in natural language
        """)
        
        # Connection Status
        connection_status = gr.Textbox(
            label="Connection Status", 
            value="Not Connected", 
            interactive=False
        )
        
        with gr.Tabs():
            # MySQL Connection Tab
            with gr.TabItem("MySQL Connection"):
                gr.Markdown("""
                **MySQL Connection Details:**
                - Enter your MySQL server connection details
                - Password will be securely handled (not stored)
                - Default port is 3306 if not specified
                - Special characters in passwords are supported
                """)
                
                with gr.Row():
                    mysql_host = gr.Textbox(
                        label="Host", 
                        value="localhost", 
                        placeholder="localhost or IP address"
                    )
                    mysql_port = gr.Textbox(
                        label="Port", 
                        value="3306", 
                        placeholder="3306"
                    )
                
                with gr.Row():
                    mysql_username = gr.Textbox(
                        label="Username", 
                        placeholder="root or your username"
                    )
                    mysql_password = gr.Textbox(
                        label="DB Password(optional)", 
                        type="password",
                        placeholder="Your MySQL password"
                    )
                
                mysql_database = gr.Textbox(
                    label="Database Name", 
                    placeholder="my_database"
                )
                
                mysql_connect_btn = gr.Button("Connect to MySQL", variant="primary")
                mysql_message = gr.Textbox(label="Connection Message", interactive=False)
            
                # API Key Section
                gr.Markdown("""
                ### 🔑 Google API Key
                Enter your Google API key for Gemini model. If not provided, will use environment variable.
                """)
                
                with gr.Row():
                    api_key_input = gr.Textbox(
                        label="Google API Key", 
                    type="password",
                        placeholder="Enter your Gemini API key here",
                        info="Get your API key from: https://makersuite.google.com/app/apikey"
                    )
                    api_key_btn = gr.Button("Set API Key", variant="secondary")
                
                api_key_message = gr.Textbox(label="API Key Status", interactive=False)
        
        # Database Analysis Section
        with gr.Tabs():
            with gr.TabItem("Database Analysis"):
                gr.Markdown("## 🔍 Database Analysis")
                analyze_btn = gr.Button("Analyze Database", variant="secondary", size="lg")
                
                with gr.Row():
                    analysis_status = gr.Textbox(label="Analysis Status", interactive=False)
                    analysis_details = gr.Textbox(label="Analysis Details", lines=10, interactive=False)
                
                # Schema Visualization
                gr.Markdown("### 📊 Database Schema Visualization")
                with gr.Row():
                    schema_table_select = gr.Dropdown(label="Select Table", choices=[], interactive=True)
                    visualize_schema_btn = gr.Button("Visualize Schema", variant="secondary")
                
                schema_output = gr.HTML(label="Schema Visualization")
                
                # Removed ER Diagram Visualization section
            
            # Query Section
            with gr.TabItem("Query Database"):
                gr.Markdown("## 💬 Ask Questions")
                
                question_input = gr.Textbox(
                    label="Your Question",
                    placeholder="Example: Show me all customers from New York, What are the top 5 selling products?",
                    lines=2
                )
                
                query_btn = gr.Button("Ask Question", variant="primary", size="lg")
                
                with gr.Row():
                    query_status = gr.Textbox(label="Query Result", lines=5, interactive=False)
                    query_output = gr.HTML(label="Data Output")
                
                # Example questions
                gr.Markdown("""
                ### 💡 Example Questions:
                - "Show me all users registered in the last month"
                - "What are the top 5 products by sales?"
                - "How many orders were placed yesterday?"
                - "Show me customers with more than 10 orders"
                - "What's the average order value?"
                """)
            
            # Query History Tab
            with gr.TabItem("Query History"):
                gr.Markdown("## 📜 Query History")
                
                with gr.Row():
                    refresh_history_btn = gr.Button("Refresh History", variant="secondary")
                    clear_history_btn = gr.Button("Clear History", variant="secondary")
                
                history_output = gr.HTML(label="Query History")
                
                # Reuse Query Section
                gr.Markdown("### 🔄 Reuse Previous Query")
                with gr.Row():
                    history_question_select = gr.Dropdown(label="Select Previous Question", choices=[], interactive=True)
                    reuse_query_btn = gr.Button("Use Selected Query", variant="primary")
        
        # Event handlers
        mysql_connect_btn.click(
            fn=lambda h, p, u, pw, d: db_manager.connect_mysql(h, p, u, pw, d) + (db_manager.connection_status,),
            inputs=[mysql_host, mysql_port, mysql_username, mysql_password, mysql_database],
            outputs=[mysql_message, connection_status]
        )
        
        # API Key event handler
        api_key_btn.click(
            fn=db_manager.set_api_key,
            inputs=[api_key_input],
            outputs=[api_key_message]
        )
        
        # Database analysis event handler
        def on_analyze_database():
            status, details = db_manager.analyze_database()
            
            # Update schema table dropdown if analysis was successful
            table_choices = []
            if "✅" in status and db_manager.db_context:
                table_choices = list(db_manager.db_context.get("tables", {}).keys())
            
            return status, details, gr.Dropdown(choices=table_choices)
        
        analyze_btn.click(
            fn=on_analyze_database,
            outputs=[analysis_status, analysis_details, schema_table_select]
        )
        
        # Schema visualization event handler
        def visualize_table_schema(table_name):
            if not table_name or not db_manager.db_context or table_name not in db_manager.db_context.get("tables", {}):
                return "<p>Please select a valid table</p>"
            
            table_info = db_manager.db_context["tables"][table_name]
            
            # Create HTML visualization
            html = f"<h3>Table: {table_name}</h3>"
            html += f"<p>Row count: {table_info.get('row_count', 'Unknown')}</p>"
            
            # Create table for columns
            html += "<table class='table table-bordered table-striped'>"
            html += "<thead><tr><th>Column</th><th>Type</th><th>Key</th></tr></thead>"
            html += "<tbody>"
            
            # Add columns
            primary_keys = table_info.get("primary_keys", [])
            foreign_keys_flat = []
            
            # Flatten foreign key references
            for fk in table_info.get("foreign_keys", []):
                for col in fk.get("columns", []):
                    foreign_keys_flat.append(col)
            
            for col in table_info.get("columns", []):
                col_name = col.get("name", "")
                col_type = col.get("type", "")
                
                # Determine key type
                key_type = ""
                if col_name in primary_keys:
                    key_type = "🔑 Primary"
                elif col_name in foreign_keys_flat:
                    key_type = "🔗 Foreign"
                
                html += f"<tr><td>{col_name}</td><td>{col_type}</td><td>{key_type}</td></tr>"
            
            html += "</tbody></table>"
            
            # Add foreign key relationships
            if table_info.get("foreign_keys"):
                html += "<h4>Foreign Key Relationships</h4>"
                html += "<ul>"
                for fk in table_info.get("foreign_keys", []):
                    cols = ", ".join(fk.get("columns", []))
                    refs = fk.get("refers_to", "")
                    html += f"<li>{cols} → {refs}</li>"
                html += "</ul>"
            
            return html
        
        visualize_schema_btn.click(
            fn=visualize_table_schema,
            inputs=[schema_table_select],
            outputs=[schema_output]
        )
        
        # Removed ER Diagram event handler
        
        # Query event handler
        query_btn.click(
            fn=db_manager.query_database,
            inputs=[question_input],
            outputs=[query_status, query_output]
        )
        
        # Query history event handlers
        def format_query_history():
            history = db_manager.get_query_history()
            if not history:
                return "<p>No queries in history</p>", gr.Dropdown(choices=[])
            
            # Format history as HTML table
            html = "<table class='table table-bordered table-striped'>"
            html += "<thead><tr><th>Time</th><th>Question</th><th>SQL Query</th><th>Execution Time</th><th>Rows</th><th>Complexity</th></tr></thead>"
            html += "<tbody>"
            
            # Collect questions for dropdown
            questions = []
            
            for i, item in enumerate(history):
                # Format timestamp
                timestamp = item.get("timestamp", "")
                if timestamp:
                    try:
                        dt = datetime.fromisoformat(timestamp)
                        timestamp = dt.strftime("%Y-%m-%d %H:%M:%S")
                    except:
                        pass
                
                question = item.get("question", "")
                sql_query = item.get("sql_query", "")
                exec_time = f"{item.get('execution_time_ms', 0)}ms" if "execution_time_ms" in item else "-"
                row_count = item.get("row_count", "-")
                complexity = item.get("complexity", "-")
                
                # Add question to dropdown options
                if question:
                    questions.append(question)
                
                # Format row with corrected query highlight
                row_class = " class='table-warning'" if item.get("corrected", False) else ""
                html += f"<tr{row_class}>"
                html += f"<td>{timestamp}</td>"
                html += f"<td>{question}</td>"
                html += f"<td><code>{sql_query}</code></td>"
                html += f"<td>{exec_time}</td>"
                html += f"<td>{row_count}</td>"
                html += f"<td>{complexity}</td>"
                html += "</tr>"
            
            html += "</tbody></table>"
            
            return html, gr.Dropdown(choices=questions)
        
        refresh_history_btn.click(
            fn=format_query_history,
            outputs=[history_output, history_question_select]
        )
        
        clear_history_btn.click(
            fn=lambda: (db_manager.clear_query_history(), "<p>History cleared</p>", gr.Dropdown(choices=[])),
            outputs=[history_output, history_question_select]
        )
        
        # Reuse query event handler
        def reuse_question(selected_question):
            if not selected_question:
                return gr.Textbox(value="")
            return gr.Textbox(value=selected_question)
        
        reuse_query_btn.click(
            fn=reuse_question,
            inputs=[history_question_select],
            outputs=[question_input]
        )
    
    return demo

if __name__ == "__main__":
    # Check for required environment variables
    if not os.getenv("GOOGLE_API_KEY"):
        print("⚠️  Warning: GOOGLE_API_KEY not found in environment variables")
        print("You will need to provide an API key in the interface or set the environment variable.")
        print("Get your API key from: https://makersuite.google.com/app/apikey")
    
    # Create and launch the interface
    demo = create_interface()
    demo.launch(
        #server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=False
    )