"""
Index Advisor service for PostgreSQL index recommendations.

This service analyzes queries and workloads to recommend optimal indexes.
It uses HypoPG for hypothetical index testing when available.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any

try:
    from pglast import parse_sql
    HAS_PGLAST = True
except ImportError:
    HAS_PGLAST = False

from .hypopg_service import HypoPGService
from .sql_driver import DbConnPool, SqlDriver

logger = logging.getLogger(__name__)


@dataclass
class IndexRecommendation:
    """Represents an index recommendation."""
    table: str
    columns: list[str]
    using: str = "btree"
    estimated_size_bytes: int = 0
    estimated_improvement: float | None = None
    reason: str = ""
    create_statement: str = ""

    @property
    def definition(self) -> str:
        """Generate the CREATE INDEX statement."""
        if self.create_statement:
            return self.create_statement
        columns_str = ", ".join(self.columns)
        return f"CREATE INDEX ON {self.table} USING {self.using} ({columns_str})"


@dataclass
class WorkloadAnalysisResult:
    """Result of workload analysis."""
    recommendations: list[IndexRecommendation] = field(default_factory=list)
    analyzed_queries: int = 0
    total_improvement: float | None = None
    error: str | None = None


class IndexAdvisor:
    """
    Service for analyzing queries and recommending indexes.

    Uses a combination of:
    - Query parsing to extract referenced columns
    - pg_stat_statements for workload analysis
    - HypoPG for hypothetical index testing
    """

    def __init__(self, pool: DbConnPool):
        """
        Initialize the index advisor.

        Args:
            pool: Database connection pool
        """
        self.pool = pool
        self.driver = SqlDriver(pool)
        self.hypopg = HypoPGService(pool)

    async def analyze_query(
        self,
        query: str,
        max_recommendations: int = 5,
    ) -> WorkloadAnalysisResult:
        """
        Analyze a single query and recommend indexes.

        Args:
            query: SQL query to analyze
            max_recommendations: Maximum number of recommendations

        Returns:
            WorkloadAnalysisResult with recommendations
        """
        result = WorkloadAnalysisResult(analyzed_queries=1)

        try:
            # Parse the query to extract table and column info
            columns_by_table = self._extract_columns_from_query(query)

            if not columns_by_table:
                result.error = "Could not extract columns from query"
                return result

            # Check HypoPG availability
            hypopg_status = await self.hypopg.check_status()

            # Generate candidate indexes
            candidates = self._generate_candidate_indexes(columns_by_table)

            # If HypoPG is available, test the candidates
            if hypopg_status.is_installed:
                result.recommendations = await self._evaluate_with_hypopg(
                    query, candidates, max_recommendations
                )
            else:
                # Without HypoPG, just return the candidates
                result.recommendations = candidates[:max_recommendations]
                for rec in result.recommendations:
                    rec.reason = "Recommended based on query structure (HypoPG not available for testing)"

            # Calculate total improvement if we have it
            if result.recommendations:
                improvements = [r.estimated_improvement for r in result.recommendations if r.estimated_improvement]
                if improvements:
                    result.total_improvement = max(improvements)

        except Exception as e:
            logger.error(f"Error analyzing query: {e}")
            result.error = str(e)

        return result

    async def analyze_queries(
        self,
        queries: list[str],
        max_recommendations: int = 10,
    ) -> WorkloadAnalysisResult:
        """
        Analyze multiple queries and recommend indexes.

        Args:
            queries: List of SQL queries to analyze
            max_recommendations: Maximum number of recommendations

        Returns:
            WorkloadAnalysisResult with recommendations
        """
        result = WorkloadAnalysisResult(analyzed_queries=len(queries))

        try:
            # Collect columns from all queries
            all_columns_by_table: dict[str, set[str]] = {}

            for query in queries:
                columns_by_table = self._extract_columns_from_query(query)
                for table, columns in columns_by_table.items():
                    if table not in all_columns_by_table:
                        all_columns_by_table[table] = set()
                    all_columns_by_table[table].update(columns)

            if not all_columns_by_table:
                result.error = "Could not extract columns from any query"
                return result

            # Generate candidate indexes
            candidates = self._generate_candidate_indexes(
                {t: list(c) for t, c in all_columns_by_table.items()}
            )

            # Check HypoPG availability
            hypopg_status = await self.hypopg.check_status()

            if hypopg_status.is_installed:
                # Evaluate each candidate against all queries
                scored_candidates = []

                for candidate in candidates:
                    total_improvement = 0
                    queries_improved = 0

                    for query in queries:
                        try:
                            test_result = await self.hypopg.explain_with_hypothetical_index(
                                query,
                                candidate.table,
                                candidate.columns,
                                candidate.using,
                            )
                            if test_result.get("would_use_index") and test_result.get("improvement_percentage"):
                                total_improvement += test_result["improvement_percentage"]
                                queries_improved += 1
                        except Exception:
                            continue

                    if queries_improved > 0:
                        candidate.estimated_improvement = total_improvement / queries_improved
                        candidate.reason = f"Improves {queries_improved}/{len(queries)} queries"
                        scored_candidates.append(candidate)

                # Sort by improvement
                scored_candidates.sort(
                    key=lambda x: x.estimated_improvement or 0,
                    reverse=True
                )
                result.recommendations = scored_candidates[:max_recommendations]
            else:
                result.recommendations = candidates[:max_recommendations]
                for rec in result.recommendations:
                    rec.reason = "Recommended based on query structure"

        except Exception as e:
            logger.error(f"Error analyzing queries: {e}")
            result.error = str(e)

        return result

    async def analyze_workload(
        self,
        min_calls: int = 50,
        min_avg_time_ms: float = 5.0,
        limit: int = 100,
        max_recommendations: int = 10,
    ) -> WorkloadAnalysisResult:
        """
        Analyze workload from pg_stat_statements and recommend indexes.

        Args:
            min_calls: Minimum call count to consider
            min_avg_time_ms: Minimum average execution time in ms
            limit: Maximum queries to analyze
            max_recommendations: Maximum index recommendations

        Returns:
            WorkloadAnalysisResult with recommendations
        """
        result = WorkloadAnalysisResult()

        # Check if pg_stat_statements is available
        try:
            check_result = await self.driver.execute_query(
                "SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements'"
            )
            if not check_result:
                result.error = (
                    "pg_stat_statements extension is required for workload analysis.\n"
                    "Install with: CREATE EXTENSION pg_stat_statements;"
                )
                return result
        except Exception as e:
            result.error = f"Error checking pg_stat_statements: {e}"
            return result

        # Get top queries from pg_stat_statements
        try:
            queries_result = await self.driver.execute_query(f"""
                SELECT
                    query,
                    calls,
                    mean_exec_time,
                    total_exec_time
                FROM pg_stat_statements
                WHERE calls >= {min_calls}
                  AND mean_exec_time >= {min_avg_time_ms}
                  AND query NOT LIKE '%pg_%'
                  AND query NOT LIKE '%$%'
                  AND query ~* '^\\s*(SELECT|UPDATE|DELETE)'
                ORDER BY total_exec_time DESC
                LIMIT {limit}
            """)

            if not queries_result:
                result.error = "No queries found matching criteria"
                return result

            queries = [row.cells.get("query") for row in queries_result if row.cells.get("query")]
            result.analyzed_queries = len(queries)

            # Analyze the collected queries
            return await self.analyze_queries(queries, max_recommendations)

        except Exception as e:
            logger.error(f"Error analyzing workload: {e}")
            result.error = str(e)
            return result

    async def get_existing_indexes(self, table: str) -> list[dict[str, Any]]:
        """
        Get existing indexes for a table.

        Args:
            table: Table name

        Returns:
            List of index information
        """
        result = await self.driver.execute_query(f"""
            SELECT
                i.relname as index_name,
                am.amname as access_method,
                array_agg(a.attname ORDER BY x.ordinality) as columns,
                ix.indisunique as is_unique,
                ix.indisprimary as is_primary,
                pg_relation_size(i.oid) as size_bytes
            FROM pg_index ix
            JOIN pg_class t ON t.oid = ix.indrelid
            JOIN pg_class i ON i.oid = ix.indexrelid
            JOIN pg_am am ON am.oid = i.relam
            JOIN pg_attribute a ON a.attrelid = t.oid
            JOIN LATERAL unnest(ix.indkey) WITH ORDINALITY AS x(attnum, ordinality)
                ON a.attnum = x.attnum
            WHERE t.relname = '{table}'
            GROUP BY i.relname, am.amname, ix.indisunique, ix.indisprimary, i.oid
        """)

        if not result:
            return []

        return [row.cells for row in result]

    async def get_index_health(self, schema: str = "public") -> dict[str, Any]:
        """
        Analyze index health for a schema.

        Args:
            schema: Schema name

        Returns:
            Dictionary with index health information
        """
        health = {
            "duplicate_indexes": [],
            "unused_indexes": [],
            "bloated_indexes": [],
            "invalid_indexes": [],
        }

        # Find duplicate indexes
        try:
            dup_result = await self.driver.execute_query(f"""
                SELECT
                    pg_size_pretty(sum(pg_relation_size(idx))::bigint) as size,
                    array_agg(idx) as indexes,
                    (array_agg(idx))[1] as index1,
                    (array_agg(idx))[2] as index2
                FROM (
                    SELECT
                        indexrelid::regclass as idx,
                        indrelid::regclass as tbl,
                        indkey::text as key
                    FROM pg_index
                    WHERE indrelid::regclass::text LIKE '{schema}.%%'
                ) sub
                GROUP BY tbl, key
                HAVING count(*) > 1
            """)
            if dup_result:
                health["duplicate_indexes"] = [row.cells for row in dup_result]
        except Exception as e:
            logger.warning(f"Error finding duplicate indexes: {e}")

        # Find unused indexes
        try:
            unused_result = await self.driver.execute_query(f"""
                SELECT
                    schemaname,
                    relname as table_name,
                    indexrelname as index_name,
                    idx_scan as scans,
                    pg_size_pretty(pg_relation_size(indexrelid)) as size
                FROM pg_stat_user_indexes
                WHERE schemaname = '{schema}'
                  AND idx_scan = 0
                  AND indexrelname NOT LIKE '%%_pkey'
                ORDER BY pg_relation_size(indexrelid) DESC
            """)
            if unused_result:
                health["unused_indexes"] = [row.cells for row in unused_result]
        except Exception as e:
            logger.warning(f"Error finding unused indexes: {e}")

        # Find invalid indexes
        try:
            invalid_result = await self.driver.execute_query("""
                SELECT
                    c.relname as index_name,
                    n.nspname as schema_name,
                    t.relname as table_name
                FROM pg_class c
                JOIN pg_index i ON c.oid = i.indexrelid
                JOIN pg_class t ON i.indrelid = t.oid
                JOIN pg_namespace n ON c.relnamespace = n.oid
                WHERE NOT i.indisvalid
            """)
            if invalid_result:
                health["invalid_indexes"] = [row.cells for row in invalid_result]
        except Exception as e:
            logger.warning(f"Error finding invalid indexes: {e}")

        return health

    def _extract_columns_from_query(self, query: str) -> dict[str, list[str]]:
        """
        Extract table and column information from a query.

        Args:
            query: SQL query

        Returns:
            Dictionary mapping table names to column lists
        """
        if not HAS_PGLAST:
            logger.warning("pglast not installed, cannot parse query for column extraction")
            return {}

        try:
            parsed = parse_sql(query)
            columns_by_table: dict[str, set[str]] = {}

            # Simple extraction - this could be enhanced
            for stmt in parsed:
                self._extract_from_node(stmt, columns_by_table)

            return {t: list(c) for t, c in columns_by_table.items()}
        except Exception as e:
            logger.warning(f"Could not parse query: {e}")
            return {}

    def _extract_from_node(self, node: Any, columns: dict[str, set[str]]) -> None:
        """Recursively extract column references from AST node."""
        if node is None:
            return

        # Handle different node types based on node structure

        if hasattr(node, "__dict__"):
            for key, value in node.__dict__.items():
                if key.startswith("_"):
                    continue

                if isinstance(value, (list, tuple)):
                    for item in value:
                        self._extract_from_node(item, columns)
                else:
                    self._extract_from_node(value, columns)

    def _generate_candidate_indexes(
        self,
        columns_by_table: dict[str, list[str]]
    ) -> list[IndexRecommendation]:
        """
        Generate candidate indexes from column information.

        Args:
            columns_by_table: Dictionary mapping tables to columns

        Returns:
            List of index recommendations
        """
        candidates = []

        for table, columns in columns_by_table.items():
            if not columns:
                continue

            # Single column indexes
            for col in columns:
                candidates.append(IndexRecommendation(
                    table=table,
                    columns=[col],
                    reason="Single column index on frequently used column"
                ))

            # Multi-column index if we have multiple columns
            if len(columns) > 1:
                # Limit to 3 columns
                combo_columns = columns[:3]
                candidates.append(IndexRecommendation(
                    table=table,
                    columns=combo_columns,
                    reason="Composite index for multi-column filter"
                ))

        return candidates

    async def _evaluate_with_hypopg(
        self,
        query: str,
        candidates: list[IndexRecommendation],
        max_recommendations: int,
    ) -> list[IndexRecommendation]:
        """
        Evaluate candidate indexes using HypoPG.

        Args:
            query: Query to test
            candidates: List of candidate indexes
            max_recommendations: Maximum recommendations to return

        Returns:
            Evaluated and sorted recommendations
        """
        evaluated = []

        # Reset any existing hypothetical indexes
        await self.hypopg.reset()

        for candidate in candidates:
            try:
                test_result = await self.hypopg.explain_with_hypothetical_index(
                    query,
                    candidate.table,
                    candidate.columns,
                    candidate.using,
                )

                if test_result.get("would_use_index"):
                    candidate.estimated_improvement = test_result.get("improvement_percentage")
                    candidate.estimated_size_bytes = test_result.get("hypothetical_index", {}).get("estimated_size", 0)
                    candidate.create_statement = test_result.get("hypothetical_index", {}).get("definition", "")
                    candidate.reason = f"Estimated {candidate.estimated_improvement:.1f}% improvement" if candidate.estimated_improvement else "Would be used by query"
                    evaluated.append(candidate)

            except Exception as e:
                logger.warning(f"Error evaluating candidate {candidate}: {e}")

        # Sort by improvement
        evaluated.sort(key=lambda x: x.estimated_improvement or 0, reverse=True)

        return evaluated[:max_recommendations]
