import sqlite3
import re
from typing import List, Dict, Any, Optional, Union, Tuple
from agents import function_tool

# Database connection manager with thread safety
import threading

class SQLiteConnectionManager:
    """Manager for SQLite database connections with thread safety."""
    
    def __init__(self, db_path: str = ":memory:"):
        """Initialize connection manager.
        
        Args:
            db_path: Path to SQLite database file
        """
        self.db_path = db_path
        self.connections = {}
        self.lock = threading.Lock()
    
    def get_connection(self) -> sqlite3.Connection:
        """Get or create database connection for the current thread."""
        thread_id = threading.get_ident()
        
        with self.lock:
            if thread_id not in self.connections:
                self.connections[thread_id] = sqlite3.connect(self.db_path)
                self.connections[thread_id].execute("PRAGMA foreign_keys = ON")
        
        return self.connections[thread_id]
    
    def close(self):
        """Close all database connections."""
        with self.lock:
            for connection in self.connections.values():
                connection.close()
            self.connections.clear()

# Create global connection manager with file-based database
import tempfile
import os

# Create a temporary file for the database
temp_db = tempfile.NamedTemporaryFile(suffix='.db', delete=False)
temp_db_path = temp_db.name
temp_db.close()

# Create global connection manager
conn_manager = SQLiteConnectionManager(db_path=temp_db_path)

# Clean up the temporary database when the module is unloaded
def cleanup_db():
    try:
        os.unlink(temp_db_path)
    except:
        pass

import atexit
atexit.register(cleanup_db)

# Whitelist of allowed SQL operations
ALLOWED_OPERATIONS = {
    'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'CREATE', 'DROP', 'ALTER',
    'BEGIN', 'COMMIT', 'ROLLBACK', 'PRAGMA'
}

# Whitelist of allowed table operations
ALLOWED_TABLE_OPS = {
    'CREATE TABLE', 'DROP TABLE', 'ALTER TABLE'
}

def validate_sql(sql: str) -> Tuple[bool, Optional[str]]:
    """Validate SQL statement against whitelist.
    
    Args:
        sql: SQL statement to validate
        
    Returns:
        Tuple of (is_valid, error_message)
    """
    # Normalize SQL for validation
    normalized_sql = sql.strip().upper()
    
    # Check for empty or whitespace-only SQL
    if not normalized_sql:
        return False, "Empty SQL statement"
    
    # Extract the first word (SQL operation)
    first_word = normalized_sql.split()[0] if normalized_sql else ""
    
    # Check if operation is in whitelist
    if first_word not in ALLOWED_OPERATIONS:
        return False, f"Operation '{first_word}' is not in whitelist. Allowed operations: {', '.join(sorted(ALLOWED_OPERATIONS))}"
    
    # Additional validation for specific operations
    if first_word in ['CREATE', 'DROP', 'ALTER']:
        # Check if it's a table operation
        if not normalized_sql.startswith('CREATE TABLE') and \
           not normalized_sql.startswith('DROP TABLE') and \
           not normalized_sql.startswith('ALTER TABLE'):
            return False, f"Only table operations are allowed for {first_word}. Use {first_word} TABLE syntax."
    
    # Prevent dangerous operations
    dangerous_patterns = [
        r'DROP\s+DATABASE',
        r'ATTACH\s+DATABASE',
        r'DETACH\s+DATABASE',
        r'VACUUM',
        r'PRAGMA\s+(?!foreign_keys|table_info)',
        r'REINDEX',
        r'ANALYZE'
    ]
    
    for pattern in dangerous_patterns:
        if re.search(pattern, normalized_sql, re.IGNORECASE):
            return False, f"Dangerous operation detected: {pattern}"
    
    return True, None

@function_tool
def execute_query(sql: str, params: Optional[Tuple] = None) -> Dict[str, Any]:
    """Execute a single SQL query.
    
    Args:
        sql: SQL statement to execute
        params: Parameters for parameterized queries
        
    Returns:
        Execution result with success status, result data, and error message
    """
    try:
        # Validate SQL
        is_valid, error_msg = validate_sql(sql)
        if not is_valid:
            return {
                "success": False,
                "error": f"Validation failed: {error_msg}",
                "sql": sql
            }
        
        # Connect to database
        conn = conn_manager.get_connection()
        cursor = conn.cursor()
        
        # Execute SQL
        if params:
            cursor.execute(sql, params)
        else:
            cursor.execute(sql)
        
        # Handle different types of SQL operations
        sql_upper = sql.strip().upper()
        
        if sql_upper.startswith('SELECT') or sql_upper.startswith('PRAGMA'):
            # Fetch results for SELECT and PRAGMA statements
            results = cursor.fetchall()
            columns = [desc[0] for desc in cursor.description] if cursor.description else []
            
            # Convert results to a more readable format
            formatted_results = []
            for row in results:
                if columns:
                    formatted_results.append(dict(zip(columns, row)))
                else:
                    formatted_results.append(row)
            
            return {
                "success": True,
                "operation": sql_upper.split()[0],
                "rows_affected": len(results),
                "columns": columns,
                "data": formatted_results,
                "sql": sql,
                "message": f"Successfully executed {sql_upper.split()[0]} query"
            }
        else:
            # For non-SELECT operations
            conn.commit()
            rows_affected = cursor.rowcount
            
            return {
                "success": True,
                "operation": sql_upper.split()[0],
                "rows_affected": rows_affected,
                "sql": sql,
                "message": f"Successfully executed {sql_upper.split()[0]} operation, affected {rows_affected} rows"
            }
            
    except sqlite3.Error as e:
        return {
            "success": False,
            "error": f"SQL execution error: {str(e)}",
            "sql": sql
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Unexpected error: {str(e)}",
            "sql": sql
        }

@function_tool
def execute_batch(sql_statements: List[str]) -> List[Dict[str, Any]]:
    """Execute multiple SQL statements in batch.
    
    Args:
        sql_statements: List of SQL statements to execute
        
    Returns:
        List of execution results for each statement
    """
    results = []
    
    try:
        # Connect to database
        conn = conn_manager.get_connection()
        
        # Execute each statement
        for sql in sql_statements:
            # Validate each SQL statement
            is_valid, error_msg = validate_sql(sql)
            if not is_valid:
                results.append({
                    "success": False,
                    "error": f"Validation failed: {error_msg}",
                    "sql": sql
                })
                continue
            
            try:
                cursor = conn.cursor()
                cursor.execute(sql)
                
                sql_upper = sql.strip().upper()
                
                if sql_upper.startswith('SELECT'):
                    results_data = cursor.fetchall()
                    columns = [desc[0] for desc in cursor.description] if cursor.description else []
                    
                    results.append({
                        "success": True,
                        "operation": "SELECT",
                        "rows_affected": len(results_data),
                        "columns": columns,
                        "data": results_data,
                        "sql": sql
                    })
                else:
                    rows_affected = cursor.rowcount
                    results.append({
                        "success": True,
                        "operation": sql_upper.split()[0],
                        "rows_affected": rows_affected,
                        "sql": sql
                    })
                    
            except sqlite3.Error as e:
                results.append({
                    "success": False,
                    "error": f"SQL execution error: {str(e)}",
                    "sql": sql
                })
        
        # Commit transaction if all statements succeeded
        if all(result.get("success", False) for result in results):
            conn.commit()
        else:
            conn.rollback()
            
    except Exception as e:
        # Rollback on any unexpected error
        if conn:
            conn.rollback()
        results.append({
            "success": False,
            "error": f"Batch execution failed: {str(e)}"
        })
    
    return results

@function_tool
def get_table_info(table_name: str) -> Dict[str, Any]:
    """Get information about a table's schema.
    
    Args:
        table_name: Name of the table to get information about
        
    Returns:
        Table schema information
    """
    sql = f"PRAGMA table_info({table_name})"
    return execute_query(sql)

@function_tool
def list_tables() -> Dict[str, Any]:
    """List all tables in the database.
    
    Returns:
        List of tables in the database
    """
    sql = "SELECT name FROM sqlite_master WHERE type='table'"
    return execute_query(sql)

def close_connection():
    """Close the database connection."""
    conn_manager.close()