"""Query executor for running validated queries against PostgreSQL.

This module provides the QueryExecutor class that executes queries with
timeout handling, row limits, and proper result formatting.
"""

from __future__ import annotations

import asyncio
import logging
import time
from datetime import date, datetime, timedelta
from datetime import time as time_type
from decimal import Decimal
from typing import TYPE_CHECKING, Any
from uuid import UUID

import sqlglot
import sqlglot.errors
from sqlglot import exp

from prismiq.query import QueryBuilder
from prismiq.sql_validator import SQLValidationError, SQLValidator
from prismiq.types import (
    DatabaseSchema,
    QueryDefinition,
    QueryExecutionError,
    QueryResult,
    QueryTimeoutError,
    QueryValidationError,
)

_logger = logging.getLogger(__name__)


def qualify_table_schemas(
    sql: str,
    schema_name: str,
    known_tables: frozenset[str],
) -> str:
    """Add schema qualification to known table references in SQL.

    Rewrites unqualified table names to schema-qualified form. Only qualifies
    tables whose names appear in the known set (from schema introspection).
    CTE references and subquery aliases are left unqualified because their
    names are typically absent from the known_tables set. Note: if a CTE
    shadows a real table name, the CTE reference will be incorrectly qualified.

    Raises SQLValidationError if the SQL cannot be parsed (fail-closed) or
    if any table is explicitly qualified with a schema other than
    ``schema_name`` (prevents cross-tenant escape).

    Args:
        sql: SQL query string.
        schema_name: PostgreSQL schema name to qualify tables with.
        known_tables: Set of lowercase table names from the schema.

    Returns:
        SQL with schema-qualified table references.

    Raises:
        SQLValidationError: If the SQL cannot be parsed or references a
            foreign schema.

    Example:
        >>> qualify_table_schemas(
        ...     'SELECT * FROM "users"',
        ...     "org_123",
        ...     frozenset({"users"}),
        ... )
        'SELECT * FROM "org_123"."users"'
    """
    try:
        parsed = sqlglot.parse_one(sql, dialect="postgres")
    except sqlglot.errors.SqlglotError as exc:
        # Fail closed: reject queries we cannot schema-qualify. Returning
        # unqualified SQL would bypass the foreign-schema check below and
        # allow cross-tenant access via explicit schema qualifiers.
        _logger.error(
            "sqlglot failed to parse SQL for schema qualification; "
            "rejecting query to prevent potential tenant isolation bypass",
            exc_info=True,
        )
        raise SQLValidationError(
            "Failed to schema-qualify SQL for tenant isolation",
            errors=[f"Schema qualification parse error: {exc}"],
        ) from exc

    for table in parsed.find_all(exp.Table):
        if not table.name:
            continue
        if table.db:
            # Table already has an explicit schema qualifier — reject if
            # it points to a different schema (cross-tenant escape attempt).
            explicit_schema = (
                table.db.this if isinstance(table.db, exp.Identifier) else str(table.db)
            )
            if explicit_schema != schema_name:
                raise SQLValidationError(
                    f'Query references schema "{explicit_schema}" but only '
                    f'"{schema_name}" is allowed',
                    errors=[
                        f'Table "{explicit_schema}"."{table.name}" references a foreign schema'
                    ],
                )
        elif table.name.lower() in known_tables:
            table.set("db", exp.Identifier(this=schema_name, quoted=True))

    return parsed.sql(dialect="postgres")


if TYPE_CHECKING:
    from asyncpg import Pool  # type: ignore[import-not-found]


def serialize_value(value: Any) -> Any:
    """Convert database values to JSON-serializable Python types."""
    if value is None:
        return None
    if isinstance(value, Decimal):
        # Convert Decimal to float for JSON serialization
        return float(value)
    if isinstance(value, datetime | date):
        return value.isoformat()
    if isinstance(value, time_type):
        return value.isoformat()
    if isinstance(value, timedelta):
        return value.total_seconds()
    if isinstance(value, UUID):
        return str(value)
    if isinstance(value, bytes):
        return value.hex()
    if isinstance(value, list | tuple):
        return [serialize_value(v) for v in value]
    if isinstance(value, dict):
        return {k: serialize_value(v) for k, v in value.items()}
    return value


class QueryExecutor:
    """Executes validated queries against a PostgreSQL database.

    Handles query validation, timeout enforcement, row limits,
    and result formatting.

    Example:
        >>> executor = QueryExecutor(pool, schema, query_timeout=30.0)
        >>> result = await executor.execute(query_definition)
        >>> print(result.row_count)
        100
    """

    def __init__(
        self,
        pool: Pool,
        schema: DatabaseSchema,
        query_timeout: float = 30.0,
        max_rows: int = 10000,
        schema_name: str | None = None,
    ) -> None:
        """Initialize the query executor.

        Args:
            pool: asyncpg connection pool.
            schema: Database schema for validation.
            query_timeout: Maximum query execution time in seconds.
            max_rows: Maximum number of rows to return.
            schema_name: PostgreSQL schema name for schema-qualified SQL generation.
                If None, tables are referenced without schema prefix.
        """
        self._pool = pool
        self._schema = schema
        self._query_timeout = query_timeout
        self._max_rows = max_rows
        self._schema_name = schema_name
        self._builder = QueryBuilder(schema, schema_name=schema_name)
        self._sql_validator = SQLValidator(schema)

    async def execute(self, query: QueryDefinition) -> QueryResult:
        """Execute a query and return results.

        Args:
            query: Query definition to execute.

        Returns:
            QueryResult with columns, rows, and execution metadata.

        Raises:
            QueryValidationError: If the query fails validation.
            QueryTimeoutError: If the query exceeds the timeout.
            QueryExecutionError: If the query execution fails.
        """
        # Sanitize filters - remove filters referencing non-existent columns
        query = self._builder.sanitize_filters(query)

        # Validate
        errors = self._builder.validate(query)
        if errors:
            raise QueryValidationError("Query validation failed", errors=errors)

        # Build SQL
        sql, params = self._builder.build(query)

        # Apply row limit if not already specified
        effective_limit = query.limit
        truncated = False
        if effective_limit is None or effective_limit > self._max_rows:
            # Add limit to params and update SQL
            # We need to rebuild with the limit
            limited_query = query.model_copy(update={"limit": self._max_rows + 1})
            sql, params = self._builder.build(limited_query)
            truncated = True

        # Execute with timeout
        start_time = time.perf_counter()
        try:
            rows = await self._execute_with_timeout(sql, params)
        except asyncio.TimeoutError as e:
            raise QueryTimeoutError(
                f"Query exceeded timeout of {self._query_timeout} seconds",
                timeout_seconds=self._query_timeout,
            ) from e
        except Exception as e:
            raise QueryExecutionError(str(e), sql=sql) from e

        execution_time_ms = (time.perf_counter() - start_time) * 1000

        # Check if result was truncated
        if truncated and len(rows) > self._max_rows:
            rows = rows[: self._max_rows]
            truncated = True
        else:
            truncated = False

        # Convert to QueryResult
        return self._format_result(rows, execution_time_ms, truncated)

    async def preview(self, query: QueryDefinition, limit: int = 100) -> QueryResult:
        """Execute a query with a smaller limit for quick preview.

        Args:
            query: Query definition to execute.
            limit: Maximum rows to return (default: 100).

        Returns:
            QueryResult with limited rows.
        """
        # Create a copy with the specified limit
        preview_query = query.model_copy(update={"limit": min(limit, self._max_rows)})
        return await self.execute(preview_query)

    async def explain(self, query: QueryDefinition) -> dict[str, Any]:
        """Run EXPLAIN ANALYZE on a query.

        Args:
            query: Query definition to analyze.

        Returns:
            Query plan as a dictionary.

        Raises:
            QueryValidationError: If the query fails validation.
            QueryExecutionError: If the explain fails.
        """
        # Validate first
        errors = self._builder.validate(query)
        if errors:
            raise QueryValidationError("Query validation failed", errors=errors)

        # Build SQL
        sql, params = self._builder.build(query)

        # Wrap with EXPLAIN ANALYZE
        explain_sql = f"EXPLAIN (ANALYZE, FORMAT JSON) {sql}"

        try:
            async with self._pool.acquire() as conn:
                result = await conn.fetchval(explain_sql, *params)
                if result and isinstance(result, list) and len(result) > 0:
                    return result[0]
                return {"plan": result}
        except Exception as e:
            raise QueryExecutionError(str(e), sql=explain_sql) from e

    async def execute_raw_sql(
        self,
        sql: str,
        params: dict[str, Any] | None = None,
    ) -> QueryResult:
        """Execute a raw SQL query with validation.

        Args:
            sql: Raw SQL query (must be SELECT only).
            params: Optional named parameters for the query.

        Returns:
            QueryResult with columns, rows, and execution metadata.

        Raises:
            SQLValidationError: If the SQL fails validation.
            QueryTimeoutError: If the query exceeds the timeout.
            QueryExecutionError: If the query execution fails.
        """
        # Validate the SQL
        validation = self._sql_validator.validate(sql)
        if not validation.valid:
            raise SQLValidationError(
                "SQL validation failed: " + "; ".join(validation.errors),
                errors=validation.errors,
            )

        # Use the sanitized SQL from validation
        safe_sql = validation.sanitized_sql
        assert safe_sql is not None  # Guaranteed by valid=True

        # Schema-qualify table references for multi-tenant isolation.
        # Adds schema prefix to unqualified table names so that queries are
        # fully qualified and independent of search_path. Only tables present
        # in schema introspection are qualified; CTE names and subquery aliases
        # are skipped as long as they don't share a name with a real table.
        if self._schema_name:
            qualified_sql = qualify_table_schemas(
                safe_sql,
                self._schema_name,
                self._sql_validator.allowed_tables,
            )
            if qualified_sql == safe_sql:
                _logger.debug("No tables were schema-qualified in raw SQL")
            safe_sql = qualified_sql

        # Apply row limit using a CTE wrapper
        limited_sql = f"WITH _cte AS ({safe_sql}) SELECT * FROM _cte LIMIT {self._max_rows + 1}"

        # Convert named params to positional for asyncpg
        # asyncpg uses $1, $2, etc. for positional params
        param_list: list[Any] = []
        if params:
            # Replace :name with $n
            import re

            param_index = 1
            param_mapping: dict[str, int] = {}

            def replace_param(match: re.Match[str]) -> str:
                nonlocal param_index
                name = match.group(1)
                if name not in param_mapping:
                    param_mapping[name] = param_index
                    param_index += 1
                return f"${param_mapping[name]}"

            limited_sql = re.sub(r":([a-zA-Z_][a-zA-Z0-9_]*)", replace_param, limited_sql)

            # Build param list in order
            param_list = [None] * len(param_mapping)
            for name, idx in param_mapping.items():
                if name in params:
                    param_list[idx - 1] = params[name]
                else:
                    raise SQLValidationError(
                        f"Missing parameter: {name}",
                        errors=[f"Parameter '{name}' referenced in SQL but not provided"],
                    )

        # Execute with timeout
        start_time = time.perf_counter()
        try:
            rows = await self._execute_with_timeout(limited_sql, param_list)
        except asyncio.TimeoutError as e:
            raise QueryTimeoutError(
                f"Query exceeded timeout of {self._query_timeout} seconds",
                timeout_seconds=self._query_timeout,
            ) from e
        except Exception as e:
            raise QueryExecutionError(str(e), sql=sql) from e

        execution_time_ms = (time.perf_counter() - start_time) * 1000

        # Check if result was truncated
        truncated = len(rows) > self._max_rows
        if truncated:
            rows = rows[: self._max_rows]

        return self._format_result(rows, execution_time_ms, truncated)

    async def _execute_with_timeout(self, sql: str, params: list[Any]) -> list[Any]:
        """Execute SQL with timeout."""
        async with self._pool.acquire() as conn:
            # Set search_path as a safety net for schema isolation. Both
            # QueryBuilder and qualify_table_schemas produce fully-qualified SQL,
            # but this ensures correct resolution if any unqualified reference
            # slips through. Uses parameterized set_config to avoid SQL injection.
            if self._schema_name:
                await conn.fetchval(
                    "SELECT set_config('search_path', $1, false)",
                    self._schema_name,
                )

            # Set statement timeout on the connection
            timeout_ms = str(int(self._query_timeout * 1000))
            await conn.fetchval(
                "SELECT set_config('statement_timeout', $1, false)",
                timeout_ms,
            )

            try:
                return await asyncio.wait_for(
                    conn.fetch(sql, *params),
                    timeout=self._query_timeout,
                )
            finally:
                # Reset each independently so a failure in one doesn't skip the other
                try:
                    await conn.execute("RESET statement_timeout")
                except Exception:
                    _logger.warning("Failed to reset statement_timeout", exc_info=True)
                if self._schema_name:
                    try:
                        await conn.execute("RESET search_path")
                    except Exception:
                        _logger.error(
                            "Failed to reset search_path; closing connection "
                            "to prevent tenant schema '%s' from leaking to the pool",
                            self._schema_name,
                            exc_info=True,
                        )
                        # Close the connection so asyncpg's pool discards it
                        # on context-manager exit instead of returning it to
                        # the pool.
                        try:
                            await conn.close()
                        except Exception:
                            _logger.error(
                                "Graceful close failed; terminating forcefully",
                                exc_info=True,
                            )
                            conn.terminate()

    def _format_result(
        self, rows: list[Any], execution_time_ms: float, truncated: bool
    ) -> QueryResult:
        """Format raw database rows into QueryResult."""
        if not rows:
            return QueryResult(
                columns=[],
                column_types=[],
                rows=[],
                row_count=0,
                truncated=False,
                execution_time_ms=execution_time_ms,
            )

        # Extract column names and types from first row
        first_row = rows[0]
        columns = list(first_row.keys())

        # Get column types (using Python type names)
        # Note: Must iterate over keys(), not the record itself (which yields values)
        column_types: list[str] = []
        for key in columns:
            value = first_row[key]
            if value is None:
                column_types.append("unknown")
            else:
                column_types.append(type(value).__name__)

        # Convert rows to lists with JSON-serializable values
        result_rows = [[serialize_value(v) for v in row.values()] for row in rows]

        return QueryResult(
            columns=columns,
            column_types=column_types,
            rows=result_rows,
            row_count=len(result_rows),
            truncated=truncated,
            execution_time_ms=execution_time_ms,
        )
