"""
Database optimization service for Trax platform.

Provides query optimization, indexing recommendations, and
performance monitoring for PostgreSQL database operations.
"""

import logging
from typing import Dict, List, Optional, Protocol, Any
from sqlalchemy import text, inspect
from sqlalchemy.orm import Session
from sqlalchemy.engine import Engine

logger = logging.getLogger(__name__)


class DatabaseOptimizerProtocol(Protocol):
    """Protocol for database optimization services."""
    def analyze_table_performance(self, session: Session, table_name: str) -> Dict[str, Any]: ...
    def suggest_indexes(self, table_schema: Dict[str, Any]) -> List[Dict[str, Any]]: ...
    def optimize_query(self, query: str) -> str: ...
    def get_connection_pool_stats(self) -> Dict[str, Any]: ...


class DatabaseOptimizer:
    """Database optimization and performance monitoring."""
    
    def __init__(self):
        self.optimization_enabled = True
        self.auto_indexing_enabled = True
        self.query_analysis_enabled = True
        
        logger.info("DatabaseOptimizer initialized")
    
    def analyze_table_performance(self, session: Session, table_name: str) -> Dict[str, Any]:
        """Analyze table performance and provide optimization recommendations."""
        try:
            # Get table statistics
            stats_query = text("""
                SELECT 
                    schemaname,
                    tablename,
                    attname,
                    n_distinct,
                    correlation,
                    most_common_vals,
                    most_common_freqs
                FROM pg_stats 
                WHERE tablename = :table_name
            """)
            
            result = session.execute(stats_query, {"table_name": table_name})
            stats = [dict(row._mapping) for row in result]
            
            # Get table size information
            size_query = text("""
                SELECT 
                    pg_size_pretty(pg_total_relation_size(:table_name)) as total_size,
                    pg_size_pretty(pg_relation_size(:table_name)) as table_size,
                    pg_size_pretty(pg_total_relation_size(:table_name) - pg_relation_size(:table_name)) as index_size,
                    pg_total_relation_size(:table_name) as total_size_bytes
                FROM pg_tables 
                WHERE tablename = :table_name
            """)
            
            size_result = session.execute(size_query, {"table_name": table_name})
            size_info = dict(size_result.fetchone()._mapping) if size_result.rowcount > 0 else {}
            
            # Get index information
            index_query = text("""
                SELECT 
                    indexname,
                    indexdef
                FROM pg_indexes 
                WHERE tablename = :table_name
            """)
            
            index_result = session.execute(index_query, {"table_name": table_name})
            indexes = [dict(row._mapping) for row in index_result]
            
            # Analyze performance
            analysis = {
                "table_name": table_name,
                "statistics": stats,
                "size_information": size_info,
                "indexes": indexes,
                "recommendations": self._generate_recommendations(stats, size_info, indexes)
            }
            
            logger.info(f"Table performance analysis completed for {table_name}")
            return analysis
            
        except Exception as e:
            logger.error(f"Error analyzing table performance for {table_name}: {e}")
            return {
                "table_name": table_name,
                "error": str(e),
                "recommendations": []
            }
    
    def suggest_indexes(self, table_schema: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Suggest indexes based on table schema and usage patterns."""
        suggestions = []
        
        columns = table_schema.get("columns", [])
        estimated_rows = table_schema.get("estimated_rows", 1000)
        
        # Analyze each column for indexing potential
        for column in columns:
            column_name = column.get("name", "")
            column_type = column.get("type", "").lower()
            is_primary_key = column.get("primary_key", False)
            is_indexed = column.get("indexed", False)
            
            # Skip if already indexed or is primary key
            if is_indexed or is_primary_key:
                continue
            
            # Suggest indexes based on column type and usage patterns
            suggestion = self._analyze_column_for_indexing(
                column_name, column_type, estimated_rows
            )
            
            if suggestion:
                suggestions.append(suggestion)
        
        # Suggest composite indexes for common query patterns
        composite_suggestions = self._suggest_composite_indexes(columns, estimated_rows)
        suggestions.extend(composite_suggestions)
        
        # Sort by priority (higher priority first)
        suggestions.sort(key=lambda x: x.get("priority", 0), reverse=True)
        
        logger.info(f"Generated {len(suggestions)} index suggestions")
        return suggestions
    
    def optimize_query(self, query: str) -> str:
        """Optimize SQL query for better performance."""
        # Basic query optimization
        optimized = query.strip()
        
        # Remove unnecessary whitespace
        optimized = " ".join(optimized.split())
        
        # Add query hints for common patterns
        if "SELECT *" in optimized.upper():
            # Suggest specific column selection
            logger.warning("Consider selecting specific columns instead of SELECT *")
        
        if "ORDER BY" in optimized.upper() and "LIMIT" not in optimized.upper():
            # Suggest adding LIMIT for large result sets
            logger.info("Consider adding LIMIT clause for large result sets")
        
        # Add performance hints
        if "WHERE" in optimized.upper():
            # Ensure WHERE clause uses indexed columns
            logger.info("Ensure WHERE clause columns are properly indexed")
        
        return optimized
    
    def get_connection_pool_stats(self) -> Dict[str, Any]:
        """Get connection pool statistics."""
        # This would typically be implemented with actual connection pool
        # For now, return mock data
        return {
            "pool_size": 10,
            "checked_in": 8,
            "checked_out": 2,
            "overflow": 0,
            "invalid": 0,
            "total_connections": 10
        }
    
    def _generate_recommendations(
        self, 
        stats: List[Dict[str, Any]], 
        size_info: Dict[str, Any], 
        indexes: List[Dict[str, Any]]
    ) -> List[str]:
        """Generate optimization recommendations."""
        recommendations = []
        
        # Analyze table size
        total_size_bytes = size_info.get("total_size_bytes", 0)
        if total_size_bytes > 100 * 1024 * 1024:  # 100MB
            recommendations.append("Consider partitioning large table for better performance")
        
        # Analyze column statistics
        low_cardinality_columns = []
        high_cardinality_columns = []
        
        for stat in stats:
            n_distinct = stat.get("n_distinct", 0)
            if n_distinct > 0 and n_distinct < 10:
                low_cardinality_columns.append(stat.get("attname"))
            elif n_distinct > 1000:
                high_cardinality_columns.append(stat.get("attname"))
        
        if low_cardinality_columns:
            recommendations.append(
                f"Consider partial indexes for low cardinality columns: {low_cardinality_columns}"
            )
        
        if high_cardinality_columns:
            recommendations.append(
                f"High cardinality columns good for indexing: {high_cardinality_columns}"
            )
        
        # Analyze existing indexes
        if len(indexes) > 5:
            recommendations.append("Consider removing unused indexes to improve write performance")
        
        if not indexes:
            recommendations.append("No indexes found - consider adding indexes for frequently queried columns")
        
        return recommendations
    
    def _analyze_column_for_indexing(
        self, 
        column_name: str, 
        column_type: str, 
        estimated_rows: int
    ) -> Optional[Dict[str, Any]]:
        """Analyze a column for indexing potential."""
        # Skip certain column types
        if column_type in ["text", "blob", "longtext"]:
            return None
        
        # Determine priority based on column characteristics
        priority = 0
        index_type = "btree"
        
        # High priority for common query columns
        if column_name in ["id", "created_at", "updated_at", "status", "user_id"]:
            priority = 10
        elif column_name.endswith("_id"):
            priority = 8
        elif column_name.endswith("_at"):
            priority = 7
        elif column_name in ["name", "title", "email"]:
            priority = 6
        
        # Adjust priority based on data type
        if column_type in ["integer", "bigint"]:
            priority += 2
        elif column_type in ["varchar", "char"]:
            priority += 1
        
        # Adjust priority based on table size
        if estimated_rows > 10000:
            priority += 2
        elif estimated_rows > 1000:
            priority += 1
        
        if priority > 0:
            return {
                "column_name": column_name,
                "index_type": index_type,
                "priority": priority,
                "reason": f"Column '{column_name}' shows good indexing potential",
                "estimated_impact": "high" if priority >= 8 else "medium" if priority >= 5 else "low"
            }
        
        return None
    
    def _suggest_composite_indexes(
        self, 
        columns: List[Dict[str, Any]], 
        estimated_rows: int
    ) -> List[Dict[str, Any]]:
        """Suggest composite indexes for common query patterns."""
        suggestions = []
        
        # Common composite index patterns
        patterns = [
            (["user_id", "created_at"], "User activity queries"),
            (["status", "created_at"], "Status-based filtering with time"),
            (["category_id", "status"], "Category-based filtering"),
            (["user_id", "status"], "User-specific status queries")
        ]
        
        column_names = [col.get("name", "") for col in columns]
        
        for pattern_columns, reason in patterns:
            if all(col in column_names for col in pattern_columns):
                suggestions.append({
                    "columns": pattern_columns,
                    "index_type": "btree",
                    "priority": 7,
                    "reason": reason,
                    "estimated_impact": "high"
                })
        
        return suggestions
    
    def get_optimization_tips(self) -> List[str]:
        """Get database optimization tips."""
        return [
            "Use appropriate indexes for frequently queried columns",
            "Avoid SELECT * - specify only needed columns",
            "Use LIMIT clauses for large result sets",
            "Consider partitioning large tables",
            "Monitor query performance with EXPLAIN ANALYZE",
            "Regularly update table statistics with ANALYZE",
            "Use connection pooling for better resource management",
            "Consider read replicas for read-heavy workloads"
        ]
    
    def get_performance_metrics(self, session: Session) -> Dict[str, Any]:
        """Get database performance metrics."""
        try:
            # Get database size
            size_query = text("""
                SELECT 
                    pg_size_pretty(pg_database_size(current_database())) as database_size,
                    pg_database_size(current_database()) as database_size_bytes
            """)
            
            size_result = session.execute(size_query)
            size_info = dict(size_result.fetchone()._mapping)
            
            # Get table count
            table_query = text("""
                SELECT COUNT(*) as table_count
                FROM information_schema.tables 
                WHERE table_schema = 'public'
            """)
            
            table_result = session.execute(table_query)
            table_count = table_result.fetchone()[0]
            
            # Get index count
            index_query = text("""
                SELECT COUNT(*) as index_count
                FROM pg_indexes 
                WHERE schemaname = 'public'
            """)
            
            index_result = session.execute(index_query)
            index_count = index_result.fetchone()[0]
            
            return {
                "database_size": size_info.get("database_size"),
                "database_size_bytes": size_info.get("database_size_bytes"),
                "table_count": table_count,
                "index_count": index_count,
                "index_table_ratio": index_count / table_count if table_count > 0 else 0
            }
            
        except Exception as e:
            logger.error(f"Error getting performance metrics: {e}")
            return {"error": str(e)}


def create_database_optimizer() -> DatabaseOptimizer:
    """Create a database optimizer instance."""
    return DatabaseOptimizer()
