"""
Database schema optimization with indexing for CHRONOS.

This module provides optimized database models and index strategies to
improve query performance for vulnerability scanning and risk assessment.
"""

from typing import List, Optional, Dict, Any
from datetime import datetime
import logging

logger = logging.getLogger(__name__)


# Database Index Definitions
DATABASE_INDEXES = {
    'vulnerabilities': [
        {
            'name': 'idx_vulnerability_severity',
            'fields': ['severity'],
            'description': 'Fast filtering by severity level'
        },
        {
            'name': 'idx_vulnerability_cwe',
            'fields': ['cwe_id'],
            'description': 'Fast lookup by CWE identifier'
        },
        {
            'name': 'idx_vulnerability_timestamp',
            'fields': ['discovered_at'],
            'description': 'Time-range queries for vulnerability history'
        },
        {
            'name': 'idx_vulnerability_source',
            'fields': ['source_file', 'line_number'],
            'description': 'Fast location lookups in code'
        },
        {
            'name': 'idx_vulnerability_status',
            'fields': ['status', 'severity'],
            'description': 'Composite index for status and severity filtering'
        },
        {
            'name': 'idx_vulnerability_project_severity',
            'fields': ['project_id', 'severity', 'discovered_at'],
            'description': 'Fast project-specific vulnerability queries'
        }
    ],
    'scans': [
        {
            'name': 'idx_scan_project_date',
            'fields': ['project_id', 'scanned_at'],
            'description': 'Project scan history'
        },
        {
            'name': 'idx_scan_status',
            'fields': ['status'],
            'description': 'Find in-progress or failed scans'
        },
        {
            'name': 'idx_scan_risk_score',
            'fields': ['risk_score'],
            'description': 'Find scans by risk level'
        }
    ],
    'remediation_plans': [
        {
            'name': 'idx_remediation_priority',
            'fields': ['priority'],
            'description': 'Sort remediations by priority'
        },
        {
            'name': 'idx_remediation_status_duedate',
            'fields': ['status', 'due_date'],
            'description': 'Track overdue remediations'
        }
    ],
    'network_monitors': [
        {
            'name': 'idx_monitor_ip_port',
            'fields': ['target_ip', 'target_port'],
            'description': 'Fast network service lookup'
        },
        {
            'name': 'idx_monitor_protocol',
            'fields': ['protocol'],
            'description': 'Protocol-specific queries'
        }
    ],
    'cryptography_findings': [
        {
            'name': 'idx_crypto_algorithm',
            'fields': ['algorithm'],
            'description': 'Fast algorithm vulnerability lookup'
        },
        {
            'name': 'idx_crypto_key_size',
            'fields': ['key_size'],
            'description': 'Identify weak key sizes'
        },
        {
            'name': 'idx_crypto_strength',
            'fields': ['strength_score'],
            'description': 'Sort by cryptographic strength'
        }
    ]
}


class IndexManager:
    """Manage database indexes for optimal performance."""
    
    def __init__(self, db_connection):
        self.db = db_connection
        self.applied_indexes: List[str] = []
        
    def create_all_indexes(self) -> None:
        """Create all recommended indexes."""
        for table, indexes in DATABASE_INDEXES.items():
            for index_spec in indexes:
                self.create_index(table, index_spec)
                
    def create_index(self, table: str, index_spec: Dict[str, Any]) -> None:
        """Create a single index.
        
        Args:
            table: Table name
            index_spec: Index specification dict
        """
        try:
            fields_str = ', '.join(index_spec['fields'])
            sql = f"""
                CREATE INDEX IF NOT EXISTS {index_spec['name']}
                ON {table} ({fields_str})
            """
            self.db.execute(sql)
            self.applied_indexes.append(index_spec['name'])
            logger.info(f"Created index: {index_spec['name']} on {table}")
        except Exception as e:
            logger.error(f"Failed to create index {index_spec['name']}: {e}")
            
    def analyze_query_performance(self, query: str) -> Dict[str, Any]:
        """Analyze query performance with EXPLAIN.
        
        Args:
            query: SQL query to analyze
            
        Returns:
            Query plan analysis
        """
        try:
            explain_query = f"EXPLAIN ANALYZE {query}"
            result = self.db.execute(explain_query).fetchall()
            return {
                'query': query,
                'plan': result,
                'uses_index': any('Index' in str(r) for r in result)
            }
        except Exception as e:
            logger.error(f"Failed to analyze query: {e}")
            return {'error': str(e)}


# Optimized Query Patterns

class OptimizedQueries:
    """Collection of optimized query patterns for common operations."""
    
    @staticmethod
    def get_critical_vulnerabilities(project_id: str, limit: int = 100) -> str:
        """Get critical vulnerabilities efficiently.
        
        Uses index: idx_vulnerability_project_severity
        """
        return f"""
            SELECT v.id, v.title, v.cwe_id, v.severity, v.source_file
            FROM vulnerabilities v
            WHERE v.project_id = '{project_id}'
              AND v.severity = 'CRITICAL'
            ORDER BY v.discovered_at DESC
            LIMIT {limit}
        """
        
    @staticmethod
    def get_scan_history(project_id: str, days: int = 30) -> str:
        """Get scan history for trend analysis.
        
        Uses index: idx_scan_project_date
        """
        return f"""
            SELECT s.id, s.scanned_at, s.risk_score, COUNT(v.id) as vuln_count
            FROM scans s
            LEFT JOIN vulnerabilities v ON s.id = v.scan_id
            WHERE s.project_id = '{project_id}'
              AND s.scanned_at >= DATETIME('now', '-{days} days')
            GROUP BY s.id
            ORDER BY s.scanned_at DESC
        """
        
    @staticmethod
    def get_open_remediations(project_id: str) -> str:
        """Get open remediations sorted by priority.
        
        Uses index: idx_remediation_status_duedate
        """
        return f"""
            SELECT r.id, r.vulnerability_id, r.priority, r.due_date,
                   r.status, v.title, v.severity
            FROM remediation_plans r
            JOIN vulnerabilities v ON r.vulnerability_id = v.id
            WHERE r.project_id = '{project_id}'
              AND r.status IN ('OPEN', 'IN_PROGRESS')
            ORDER BY r.priority DESC, r.due_date ASC
        """
        
    @staticmethod
    def get_weak_cryptography(project_id: str) -> str:
        """Find weak cryptographic implementations.
        
        Uses index: idx_crypto_strength
        """
        return f"""
            SELECT c.id, c.algorithm, c.key_size, c.strength_score,
                   c.source_file, c.line_number
            FROM cryptography_findings c
            WHERE c.project_id = '{project_id}'
              AND c.strength_score < 3.0
            ORDER BY c.strength_score ASC
        """
        
    @staticmethod
    def get_network_services_by_protocol(
        project_id: str,
        protocol: str
    ) -> str:
        """Get network services by protocol.
        
        Uses index: idx_monitor_protocol
        """
        return f"""
            SELECT DISTINCT nm.target_ip, nm.target_port, nm.service_name,
                   nm.version, COUNT(*) as issue_count
            FROM network_monitors nm
            LEFT JOIN vulnerabilities v ON nm.id = v.source_id
            WHERE nm.project_id = '{project_id}'
              AND nm.protocol = '{protocol}'
            GROUP BY nm.target_ip, nm.target_port
            ORDER BY issue_count DESC
        """


class QueryOptimizer:
    """Optimize query performance with caching and batching."""
    
    def __init__(self, db_connection):
        self.db = db_connection
        self.query_cache: Dict[str, tuple] = {}
        
    def batch_insert(self, table: str, records: List[Dict]) -> int:
        """Insert multiple records efficiently.
        
        Args:
            table: Table name
            records: List of record dicts
            
        Returns:
            Number of records inserted
        """
        if not records:
            return 0
            
        # Use INSERT multiple values syntax for better performance
        columns = records[0].keys()
        placeholders = ', '.join('?' * len(columns))
        column_names = ', '.join(columns)
        
        sql = f"""
            INSERT INTO {table} ({column_names})
            VALUES ({placeholders})
        """
        
        count = 0
        try:
            for record in records:
                values = [record.get(col) for col in columns]
                self.db.execute(sql, values)
            self.db.commit()
            count = len(records)
            logger.info(f"Batch inserted {count} records into {table}")
        except Exception as e:
            logger.error(f"Batch insert failed: {e}")
            self.db.rollback()
            
        return count
        
    def batch_update(
        self,
        table: str,
        updates: List[Dict[str, Any]],
        id_field: str = 'id'
    ) -> int:
        """Update multiple records efficiently.
        
        Args:
            table: Table name
            updates: List of update dicts (must include id_field)
            id_field: Field to identify records
            
        Returns:
            Number of records updated
        """
        if not updates:
            return 0
            
        count = 0
        try:
            for update_dict in updates:
                rec_id = update_dict.pop(id_field)
                set_clause = ', '.join(
                    f"{k} = ?" for k in update_dict.keys()
                )
                sql = f"UPDATE {table} SET {set_clause} WHERE {id_field} = ?"
                values = list(update_dict.values()) + [rec_id]
                self.db.execute(sql, values)
                count += 1
                
            self.db.commit()
            logger.info(f"Batch updated {count} records in {table}")
        except Exception as e:
            logger.error(f"Batch update failed: {e}")
            self.db.rollback()
            
        return count


class AggregationQueries:
    """Pre-aggregated queries for dashboard metrics."""
    
    @staticmethod
    def create_vulnerability_summary_view(db_connection) -> None:
        """Create materialized view for vulnerability summaries.
        
        This view pre-aggregates common statistics for fast dashboard
        queries without computing on every request.
        """
        sql = """
            CREATE VIEW IF NOT EXISTS vulnerability_summary AS
            SELECT
                p.id as project_id,
                p.name as project_name,
                COUNT(DISTINCT v.id) as total_vulnerabilities,
                SUM(CASE WHEN v.severity = 'CRITICAL' THEN 1 ELSE 0 END) 
                    as critical_count,
                SUM(CASE WHEN v.severity = 'HIGH' THEN 1 ELSE 0 END) 
                    as high_count,
                SUM(CASE WHEN v.severity = 'MEDIUM' THEN 1 ELSE 0 END) 
                    as medium_count,
                SUM(CASE WHEN v.severity = 'LOW' THEN 1 ELSE 0 END) 
                    as low_count,
                AVG(s.risk_score) as avg_risk_score,
                MAX(s.scanned_at) as last_scan_date
            FROM projects p
            LEFT JOIN scans s ON p.id = s.project_id
            LEFT JOIN vulnerabilities v ON s.id = v.scan_id
            GROUP BY p.id, p.name
        """
        try:
            db_connection.execute(sql)
            logger.info("Created vulnerability_summary view")
        except Exception as e:
            logger.warning(f"Failed to create view: {e}")
            
    @staticmethod
    def create_remediation_status_view(db_connection) -> None:
        """Create materialized view for remediation status."""
        sql = """
            CREATE VIEW IF NOT EXISTS remediation_status AS
            SELECT
                p.id as project_id,
                r.priority,
                COUNT(*) as count,
                SUM(CASE WHEN r.status = 'COMPLETE' THEN 1 ELSE 0 END) 
                    as completed,
                SUM(CASE WHEN r.status = 'OVERDUE' THEN 1 ELSE 0 END) 
                    as overdue
            FROM projects p
            LEFT JOIN remediation_plans r ON p.id = r.project_id
            GROUP BY p.id, r.priority
        """
        try:
            db_connection.execute(sql)
            logger.info("Created remediation_status view")
        except Exception as e:
            logger.warning(f"Failed to create view: {e}")


def apply_database_optimizations(db_connection) -> None:
    """Apply all database performance optimizations.
    
    Args:
        db_connection: Database connection object
    """
    logger.info("Applying database optimizations...")
    
    # Create indexes
    index_manager = IndexManager(db_connection)
    index_manager.create_all_indexes()
    
    # Create aggregation views
    AggregationQueries.create_vulnerability_summary_view(db_connection)
    AggregationQueries.create_remediation_status_view(db_connection)
    
    # Enable query optimization
    try:
        db_connection.execute("PRAGMA query_only = FALSE")
        db_connection.execute("ANALYZE")  # Update statistics
        logger.info("Database optimizations applied successfully")
    except Exception as e:
        logger.error(f"Failed to apply optimizations: {e}")
