backend/core/query_optimizer.py

"""
Query Analysis and Optimization System for AegisLM SaaS Backend.

Production-ready query analysis with optimization recommendations,
index suggestions, and performance improvements.
"""

import asyncio
import re
from typing import List, Dict, Optional, Any, Tuple, Set
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession
import logging
import json
from dataclasses import dataclass
from enum import Enum

from .database import async_engine
from .performance_monitor import performance_monitor
from .config import settings

logger = logging.getLogger(__name__)


class QueryType(Enum):
    """Query types for analysis."""
    SELECT = "SELECT"
    INSERT = "INSERT"
    UPDATE = "UPDATE"
    DELETE = "DELETE"
    JOIN = "JOIN"
    AGGREGATE = "AGGREGATE"
    SUBQUERY = "SUBQUERY"


@dataclass
class QueryAnalysis:
    """Query analysis results."""
    query: str
    query_type: QueryType
    complexity_score: float
    tables_involved: List[str]
    indexes_used: List[str]
    missing_indexes: List[Dict[str, Any]]
    optimization_suggestions: List[str]
    estimated_cost: Optional[float]
    execution_plan: Optional[Dict[str, Any]]


@dataclass
class IndexRecommendation:
    """Index recommendation."""
    table_name: str
    columns: List[str]
    index_type: str
    estimated_impact: str
    reason: str


class QueryAnalyzer:
    """Advanced query analyzer and optimizer."""
    
    def __init__(self):
        self.query_patterns = {
            'slow_patterns': [
                r'SELECT.*\s+FROM\s+\w+\s+WHERE\s+.*LIKE\s+.*%',  # Leading wildcard LIKE
                r'SELECT.*\s+FROM\s+\w+\s+WHERE\s+.*OR\s+',           # OR conditions
                r'SELECT.*\s+FROM\s+\w+\s+ORDER BY\s+.*\s+LIMIT\s+',  # ORDER BY + LIMIT
                r'SELECT.*\s+FROM\s+\w+\s+WHERE\s+.*IN\s+\(.*SELECT',  # Subquery in IN
                r'SELECT.*\s+FROM\s+\w+\s+WHERE\s+.*NOT\s+IN',        # NOT IN
                r'SELECT.*\s+FROM\s+\w+\s+WHERE\s+.*!=\s*',           # Not equal operator
            ],
            'join_patterns': [
                r'JOIN\s+\w+\s+ON\s+.*=.*',                           # JOIN conditions
                r'LEFT\s+JOIN',                                       # LEFT JOIN
                r'RIGHT\s+JOIN',                                      # RIGHT JOIN
                r'FULL\s+OUTER\s+JOIN',                               # FULL OUTER JOIN
            ],
            'aggregate_patterns': [
                r'COUNT\(', r'SUM\(', r'AVG\(', r'MIN\(', r'MAX\(',  # Aggregate functions
                r'GROUP\s+BY',                                       # GROUP BY
                r'HAVING\s+',                                        # HAVING clause
            ]
        }
    
    async def analyze_query(self, query: str) -> QueryAnalysis:
        """Perform comprehensive query analysis."""
        # Normalize query
        normalized_query = self._normalize_query(query)
        
        # Determine query type
        query_type = self._determine_query_type(normalized_query)
        
        # Calculate complexity
        complexity_score = self._calculate_complexity(normalized_query)
        
        # Extract tables
        tables = self._extract_tables(normalized_query)
        
        # Get execution plan
        execution_plan = await self._get_execution_plan(query)
        
        # Analyze index usage
        indexes_used, missing_indexes = await self._analyze_indexes(normalized_query, tables, execution_plan)
        
        # Generate optimization suggestions
        suggestions = self._generate_optimization_suggestions(
            normalized_query, query_type, complexity_score, execution_plan, missing_indexes
        )
        
        # Extract estimated cost
        estimated_cost = self._extract_cost(execution_plan)
        
        return QueryAnalysis(
            query=query,
            query_type=query_type,
            complexity_score=complexity_score,
            tables_involved=tables,
            indexes_used=indexes_used,
            missing_indexes=missing_indexes,
            optimization_suggestions=suggestions,
            estimated_cost=estimated_cost,
            execution_plan=execution_plan
        )
    
    def _normalize_query(self, query: str) -> str:
        """Normalize query for analysis."""
        # Remove extra whitespace
        normalized = ' '.join(query.split())
        # Convert to uppercase for keywords, but keep string literals
        return normalized
    
    def _determine_query_type(self, query: str) -> QueryType:
        """Determine the primary query type."""
        query_upper = query.upper()
        
        if 'JOIN' in query_upper:
            return QueryType.JOIN
        elif any(pattern in query_upper for pattern in ['COUNT(', 'SUM(', 'AVG(', 'GROUP BY']):
            return QueryType.AGGREGATE
        elif query_upper.startswith('SELECT'):
            return QueryType.SELECT
        elif query_upper.startswith('INSERT'):
            return QueryType.INSERT
        elif query_upper.startswith('UPDATE'):
            return QueryType.UPDATE
        elif query_upper.startswith('DELETE'):
            return QueryType.DELETE
        elif 'SELECT' in query_upper and '(' in query_upper:
            return QueryType.SUBQUERY
        else:
            return QueryType.SELECT
    
    def _calculate_complexity(self, query: str) -> float:
        """Calculate query complexity score."""
        score = 0.0
        
        # Base score for query type
        if 'JOIN' in query.upper():
            score += 2.0
        if 'SUBQUERY' in query.upper() or '(' in query:
            score += 1.5
        if any(pattern in query.upper() for pattern in ['COUNT(', 'SUM(', 'AVG(', 'GROUP BY']):
            score += 1.0
        
        # Complexity from conditions
        if 'WHERE' in query.upper():
            conditions = query.upper().split('WHERE')[1].split('ORDER BY')[0].split('GROUP BY')[0]
            score += conditions.count('AND') * 0.5
            score += conditions.count('OR') * 0.8
            score += conditions.count('LIKE') * 0.3
            score += conditions.count('IN') * 0.4
        
        # Complexity from functions
        functions = ['COUNT(', 'SUM(', 'AVG(', 'MIN(', 'MAX(', 'COALESCE(', 'CASE WHEN']
        for func in functions:
            score += query.upper().count(func) * 0.2
        
        return min(score, 10.0)  # Cap at 10
    
    def _extract_tables(self, query: str) -> List[str]:
        """Extract table names from query."""
        tables = []
        
        # Simple regex-based extraction (can be improved)
        from_pattern = re.search(r'FROM\s+(\w+)', query, re.IGNORECASE)
        if from_pattern:
            tables.append(from_pattern.group(1))
        
        # Extract JOIN tables
        join_patterns = re.findall(r'JOIN\s+(\w+)', query, re.IGNORECASE)
        tables.extend(join_patterns)
        
        return list(set(tables))  # Remove duplicates
    
    async def _get_execution_plan(self, query: str) -> Optional[Dict[str, Any]]:
        """Get query execution plan."""
        try:
            async with async_engine.begin() as conn:
                # Use EXPLAIN (ANALYZE, BUFFERS) for detailed plan
                explain_query = f"EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) {query}"
                result = await conn.execute(text(explain_query))
                plan_data = result.fetchone()[0]
                
                if plan_data and len(plan_data) > 0:
                    return plan_data[0]  # First element is the plan
        except Exception as e:
            logger.error(f"Failed to get execution plan: {e}")
        
        return None
    
    async def _analyze_indexes(self, query: str, tables: List[str], 
                             execution_plan: Optional[Dict[str, Any]]) -> Tuple[List[str], List[Dict[str, Any]]]:
        """Analyze index usage and suggest missing indexes."""
        indexes_used = []
        missing_indexes = []
        
        try:
            async with async_engine.begin() as conn:
                # Get existing indexes for involved tables
                for table in tables:
                    result = await conn.execute(text("""
                        SELECT indexname, indexdef 
                        FROM pg_indexes 
                        WHERE tablename = :table_name
                    """), {"table_name": table})
                    
                    table_indexes = {row.indexname: row.indexdef for row in result.fetchall()}
                    
                    # Check if indexes are used in execution plan
                    if execution_plan:
                        used_indexes = self._extract_indexes_from_plan(execution_plan, table)
                        indexes_used.extend(used_indexes)
                    
                    # Suggest missing indexes based on WHERE clauses
                    where_clause = self._extract_where_clause(query)
                    if where_clause:
                        suggestions = self._suggest_indexes_for_where(table, where_clause, table_indexes)
                        missing_indexes.extend(suggestions)
        
        except Exception as e:
            logger.error(f"Failed to analyze indexes: {e}")
        
        return list(set(indexes_used)), missing_indexes
    
    def _extract_indexes_from_plan(self, plan: Dict[str, Any], table_name: str) -> List[str]:
        """Extract index names from execution plan."""
        indexes = []
        
        def traverse_plan(node):
            if isinstance(node, dict):
                if 'Index Name' in node and table_name in str(node.get('Relation Name', '')):
                    indexes.append(node['Index Name'])
                
                # Recursively check child nodes
                for key, value in node.items():
                    if key in ['Plans', 'Plan']:
                        if isinstance(value, list):
                            for child in value:
                                traverse_plan(child)
                        else:
                            traverse_plan(value)
        
        traverse_plan(plan)
        return indexes
    
    def _extract_where_clause(self, query: str) -> Optional[str]:
        """Extract WHERE clause from query."""
        where_match = re.search(r'WHERE\s+(.+?)(?:\s+ORDER\s+BY|\s+GROUP\s+BY|\s+LIMIT|$)', query, re.IGNORECASE)
        return where_match.group(1) if where_match else None
    
    def _suggest_indexes_for_where(self, table: str, where_clause: str, 
                                 existing_indexes: Dict[str, str]) -> List[Dict[str, Any]]:
        """Suggest indexes based on WHERE clause."""
        suggestions = []
        
        # Extract columns from WHERE clause
        columns = re.findall(r'(\w+)\s*=', where_clause)
        columns.extend(re.findall(r'(\w+)\s+IN\s+', where_clause))
        columns.extend(re.findall(r'(\w+)\s+LIKE\s+', where_clause))
        
        # Remove duplicates and common non-indexable columns
        columns = list(set([col for col in columns if col.lower() not in ['is', 'are', 'not', 'null']]))
        
        if columns:
            # Check if similar index already exists
            for col in columns:
                col_lower = col.lower()
                existing = any(col_lower in idx_def.lower() for idx_def in existing_indexes.values())
                
                if not existing:
                    suggestions.append({
                        'table': table,
                        'columns': [col],
                        'type': 'btree',
                        'reason': f'Column "{col}" used in WHERE clause but no suitable index found',
                        'estimated_impact': 'medium'
                    })
        
        # Check for composite index opportunities
        if len(columns) >= 2:
            # Suggest composite index for multiple conditions
            suggestions.append({
                'table': table,
                'columns': columns[:2],  # First two columns
                'type': 'btree',
                'reason': f'Multiple columns in WHERE clause could benefit from composite index',
                'estimated_impact': 'high'
            })
        
        return suggestions
    
    def _generate_optimization_suggestions(self, query: str, query_type: QueryType, 
                                        complexity: float, execution_plan: Optional[Dict[str, Any]], 
                                        missing_indexes: List[Dict[str, Any]]) -> List[str]:
        """Generate optimization suggestions."""
        suggestions = []
        
        # High complexity suggestions
        if complexity > 7.0:
            suggestions.append("Consider breaking this complex query into simpler parts")
        
        # Missing indexes
        for index_rec in missing_indexes:
            if index_rec['estimated_impact'] == 'high':
                suggestions.append(f"Create index on {index_rec['table']}.{', '.join(index_rec['columns'])}")
        
        # Pattern-based suggestions
        query_upper = query.upper()
        
        # LIKE with leading wildcard
        if re.search(r'LIKE\s+\'\%', query_upper):
            suggestions.append("Avoid leading wildcards in LIKE queries - consider full-text search")
        
        # OR conditions
        if query_upper.count(' OR ') > 2:
            suggestions.append("Multiple OR conditions - consider using UNION ALL or IN clauses")
        
        # NOT IN
        if 'NOT IN' in query_upper:
            suggestions.append("NOT IN can be slow - consider LEFT JOIN/IS NULL pattern")
        
        # SELECT *
        if 'SELECT *' in query_upper:
            suggestions.append("Avoid SELECT * - specify only needed columns")
        
        # Execution plan based suggestions
        if execution_plan:
            if 'Seq Scan' in str(execution_plan):
                suggestions.append("Sequential scan detected - consider adding indexes")
            
            if 'Sort' in str(execution_plan):
                suggestions.append("Sorting operation detected - ensure proper indexes for ORDER BY")
        
        # JOIN suggestions
        if query_type == QueryType.JOIN:
            if 'LEFT JOIN' in query_upper:
                suggestions.append("LEFT JOIN can be expensive - ensure it's necessary")
            
            if query_upper.count('JOIN') > 3:
                suggestions.append("Multiple JOINs - consider query restructuring")
        
        return suggestions
    
    def _extract_cost(self, execution_plan: Optional[Dict[str, Any]]) -> Optional[float]:
        """Extract estimated cost from execution plan."""
        if execution_plan and 'Total Cost' in execution_plan:
            return float(execution_plan['Total Cost'])
        return None
    
    async def generate_index_recommendations(self, table_name: Optional[str] = None) -> List[IndexRecommendation]:
        """Generate index recommendations for tables."""
        recommendations = []
        
        try:
            async with async_engine.begin() as conn:
                if table_name:
                    tables = [table_name]
                else:
                    # Get all user tables
                    result = await conn.execute(text("""
                        SELECT tablename FROM pg_tables 
                        WHERE schemaname = 'public' 
                        ORDER BY tablename
                    """))
                    tables = [row.tablename for row in result.fetchall()]
                
                for table in tables:
                    # Analyze table for index opportunities
                    table_recommendations = await self._analyze_table_for_indexes(table, conn)
                    recommendations.extend(table_recommendations)
        
        except Exception as e:
            logger.error(f"Failed to generate index recommendations: {e}")
        
        return recommendations
    
    async def _analyze_table_for_indexes(self, table: str, conn) -> List[IndexRecommendation]:
        """Analyze a specific table for index opportunities."""
        recommendations = []
        
        try:
            # Get table statistics
            result = await conn.execute(text(f"""
                SELECT 
                    n_tup_ins as inserts,
                    n_tup_upd as updates,
                    n_tup_del as deletes,
                    n_live_tup as live_tuples,
                    n_dead_tup as dead_tuples
                FROM pg_stat_user_tables 
                WHERE schemaname = 'public' AND tablename = '{table}'
            """))
            
            stats = result.fetchone()
            if not stats:
                return recommendations
            
            # High update/delete tables might need different indexing strategy
            if stats.updates > stats.inserts * 2:
                recommendations.append(IndexRecommendation(
                    table_name=table,
                    columns=['id'],  # Assuming primary key
                    index_type='btree',
                    estimated_impact='medium',
                    reason='High update activity - ensure primary key index is optimized'
                ))
            
            # Check for foreign key opportunities
            result = await conn.execute(text(f"""
                SELECT 
                    tc.constraint_name,
                    kcu.column_name
                FROM information_schema.table_constraints AS tc 
                JOIN information_schema.key_column_usage AS kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                WHERE tc.constraint_type = 'FOREIGN KEY' 
                    AND tc.table_schema = 'public'
                    AND tc.table_name = '{table}'
            """))
            
            for row in result.fetchall():
                recommendations.append(IndexRecommendation(
                    table_name=table,
                    columns=[row.column_name],
                    index_type='btree',
                    estimated_impact='high',
                    reason=f'Foreign key column {row.column_name} should be indexed'
                ))
        
        except Exception as e:
            logger.error(f"Failed to analyze table {table}: {e}")
        
        return recommendations


# Global query analyzer instance
query_analyzer = QueryAnalyzer()


# Query optimization service
class QueryOptimizationService:
    """Service for query optimization and analysis."""
    
    def __init__(self):
        self.analyzer = query_analyzer
    
    async def optimize_query(self, query: str) -> Dict[str, Any]:
        """Optimize a query and return recommendations."""
        try:
            # Analyze the query
            analysis = await self.analyzer.analyze_query(query)
            
            # Generate optimized version if possible
            optimized_query = self._generate_optimized_query(query, analysis)
            
            # Get index recommendations
            index_recs = await self.analyzer.generate_index_recommendations()
            
            return {
                "original_query": query,
                "optimized_query": optimized_query,
                "analysis": {
                    "query_type": analysis.query_type.value,
                    "complexity_score": analysis.complexity_score,
                    "tables_involved": analysis.tables_involved,
                    "estimated_cost": analysis.estimated_cost
                },
                "optimization_suggestions": analysis.optimization_suggestions,
                "missing_indexes": analysis.missing_indexes,
                "index_recommendations": [rec.__dict__ for rec in index_recs[:10]]  # Top 10
            }
        
        except Exception as e:
            logger.error(f"Query optimization failed: {e}")
            return {
                "error": str(e),
                "original_query": query
            }
    
    def _generate_optimized_query(self, query: str, analysis: QueryAnalysis) -> Optional[str]:
        """Generate an optimized version of the query."""
        optimized = query
        
        try:
            # Basic optimizations
            if 'SELECT *' in optimized.upper():
                # Replace * with specific columns (simplified)
                if analysis.tables_involved:
                    optimized = optimized.replace('SELECT *', f'SELECT id')  # Simplified
            
            # Remove unnecessary parentheses
            optimized = re.sub(r'\(\s*([^()]+)\s*\)', r'\1', optimized)
            
            # Simplify WHERE conditions (basic)
            if '1=1' in optimized:
                optimized = optimized.replace('AND 1=1', '').replace('WHERE 1=1', 'WHERE')
            
            # Return optimized if changed
            return optimized if optimized != query else None
        
        except Exception:
            return None  # Return None if optimization fails
    
    async def analyze_slow_queries(self, limit: int = 20) -> List[Dict[str, Any]]:
        """Analyze recent slow queries."""
        slow_queries = await performance_monitor.get_slow_queries(limit)
        
        analyzed_queries = []
        for query_data in slow_queries:
            try:
                analysis = await self.analyzer.analyze_query(query_data['query'])
                analyzed_queries.append({
                    "query_data": query_data,
                    "analysis": analysis.__dict__
                })
            except Exception as e:
                logger.error(f"Failed to analyze slow query: {e}")
        
        return analyzed_queries
    
    async def get_database_optimization_report(self) -> Dict[str, Any]:
        """Generate comprehensive database optimization report."""
        try:
            # Get index recommendations
            index_recs = await self.analyzer.generate_index_recommendations()
            
            # Analyze top slow queries
            slow_analysis = await self.analyze_slow_queries(10)
            
            # Get performance summary
            perf_summary = await performance_monitor.get_performance_summary()
            
            return {
                "timestamp": asyncio.get_event_loop().time(),
                "index_recommendations": [rec.__dict__ for rec in index_recs],
                "slow_queries_analysis": slow_analysis,
                "performance_summary": perf_summary,
                "optimization_priority": self._calculate_optimization_priority(index_recs, slow_analysis)
            }
        
        except Exception as e:
            logger.error(f"Failed to generate optimization report: {e}")
            return {"error": str(e)}
    
    def _calculate_optimization_priority(self, index_recs: List[IndexRecommendation], 
                                       slow_analysis: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Calculate optimization priority recommendations."""
        high_priority = []
        medium_priority = []
        low_priority = []
        
        # Prioritize index recommendations
        for rec in index_recs:
            if rec.estimated_impact == 'high' and 'foreign key' in rec.reason.lower():
                high_priority.append(f"Create index: {rec.table_name}({', '.join(rec.columns)})")
            elif rec.estimated_impact == 'high':
                medium_priority.append(f"Create index: {rec.table_name}({', '.join(rec.columns)})")
            else:
                low_priority.append(f"Create index: {rec.table_name}({', '.join(rec.columns)})")
        
        # Prioritize slow query fixes
        for analysis in slow_analysis:
            query_data = analysis['query_data']
            if query_data['execution_time'] > 5.0:  # Very slow queries
                high_priority.append(f"Optimize slow query: {query_data['execution_time']:.2f}s {query_data['query_type']}")
        
        return {
            "high_priority": high_priority[:5],  # Top 5
            "medium_priority": medium_priority[:10],  # Top 10
            "low_priority": low_priority[:10]  # Top 10
        }


# Global optimization service instance
query_optimizer = QueryOptimizationService()


if __name__ == "__main__":
    import sys
    
    async def main():
        command = sys.argv[1] if len(sys.argv) > 1 else "help"
        
        if command == "analyze":
            if len(sys.argv) < 3:
                print("Error: analyze requires a query string")
                sys.exit(1)
            
            query = ' '.join(sys.argv[2:])
            optimization = await query_optimizer.optimize_query(query)
            print(json.dumps(optimization, indent=2, default=str))
        
        elif command == "report":
            report = await query_optimizer.get_database_optimization_report()
            print(json.dumps(report, indent=2, default=str))
        
        elif command == "indexes":
            recommendations = await query_analyzer.generate_index_recommendations()
            print(f"Index recommendations: {len(recommendations)}")
            for rec in recommendations:
                print(f"  - {rec.table_name}({', '.join(rec.columns)}) - {rec.reason}")
        
        else:
            print("Usage: python query_optimizer.py <command> [args]")
            print("Commands: analyze <query>, report, indexes")
    
    asyncio.run(main())