"""
Snowflake executor factory.

Wraps the ``ConnectorSnowflake`` from ``snowflake-data-validation`` to provide
connection management, SQL literal formatting, and stored-procedure execution.

Supports multiple connection modes:
1. Named connections (from TOML files)
2. Session token-based connections (for desktop app integration)
3. Connection wrapping (reuse existing connections)

Column types are captured from ``cursor.description`` using the Snowflake
connector's built-in ``FIELD_ID_TO_NAME`` mapping.
"""

from __future__ import annotations

import logging
from datetime import datetime
from decimal import Decimal
from typing import Any, Sequence

from snowflake.connector import SnowflakeConnection
from snowflake.connector.constants import FIELD_ID_TO_NAME
from snowflake.connector.cursor import DictCursor
from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase
from snowflake.snowflake_data_validation.snowflake.model.snowflake_named_connection import (
    SnowflakeNamedConnection,
)

from test_runner.common.database_dialect import DatabaseDialect
from test_runner.common.database_executor import (
    DatabaseExecutor,
    DatabaseExecutorFactory,
    LiteralFormatter,
)
from test_runner.common.models import (
    ColumnTypeMap,
    TestCaseResult,
    ValidationSteps,
)
from test_runner.validate.targets.snowflake_connection_config import (
    ConnectionMode,
    SessionTokenConfig,
)
from test_runner.validate.targets.snowflake_connection_manager import (
    SnowflakeConnectionManager,
)
from test_runner.replay.isolation import SnowflakeCloneIsolation
from test_runner.replay.source_introspector import DefaultIntrospector

LOGGER = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _describe_to_column_types(
    cursor_description: Sequence[Any],
) -> ColumnTypeMap:
    """Map Snowflake ``cursor.description`` to ``{COLUMN_NAME: type_name}``.

    Uses the connector's built-in ``FIELD_ID_TO_NAME`` mapping which
    translates integer type codes to Snowflake type names.
    """
    return {
        desc[0].upper(): FIELD_ID_TO_NAME.get(desc[1], "VARCHAR")
        for desc in cursor_description
    }


# ---------------------------------------------------------------------------
# Literal formatter (SRP: separated from factory)
# ---------------------------------------------------------------------------

class SnowflakeLiteralFormatter(LiteralFormatter):
    """Snowflake SQL literal formatting: TRUE/FALSE for bools."""

    def format_literal(self, value: Any) -> str:
        if value is None:
            return "NULL"
        if isinstance(value, bool):
            return "TRUE" if value else "FALSE"
        if isinstance(value, (int, float, Decimal)):
            return str(value)
        if isinstance(value, datetime):
            return f"'{value.isoformat()}'"
        escaped = str(value).replace("'", "''")
        return f"'{escaped}'"


# ---------------------------------------------------------------------------
# Executor
# ---------------------------------------------------------------------------

class SnowflakeExecutor(DatabaseExecutor):
    """Executes SQL on Snowflake via the data-validation ``ConnectorSnowflake``.

    Captures column types from ``cursor.description`` using the Snowflake
    connector's integer type codes mapped to type names.
    """

    def __init__(self, connector: ConnectorBase) -> None:
        self._connector = connector

    @property
    def connector(self) -> ConnectorBase:
        return self._connector

    def close(self) -> None:
        self._connector.close()

    def execute(self, sql: str, steps: ValidationSteps | None = None, fetch_stmts: list[str] | None = None) -> TestCaseResult:
        result = TestCaseResult()
        try:
            import snowflake.connector as sf_connector

            cursor = self._connector.connection.cursor(sf_connector.DictCursor)
            try:
                cursor.execute(sql)
                self._append_result_set(cursor, result)
            finally:
                cursor.close()
        except Exception as exc:
            result.success = False
            result.error = str(exc)
        return result

    def execute_statement(self, sql: str) -> None:
        self._connector.execute_statement(sql)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _append_result_set(cursor: DictCursor, result: TestCaseResult) -> None:
        """Read the current cursor result set and append it to *result*."""
        description = cursor.description
        if description is None:
            result.result_sets.append([])
            result.row_counts.append(0)
            result.column_types.append({})
            return

        rows = cursor.fetchall()
        result.result_sets.append(rows)
        result.row_counts.append(len(rows))
        result.column_types.append(_describe_to_column_types(description))


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------

class SnowflakeExecutorFactory(DatabaseExecutorFactory):
    """Factory for Snowflake executors backed by data-validation connectors.
    
    Supports multiple connection modes:
    1. Named connections (from TOML files)
    2. Session token connections (for desktop app integration)
    3. Connection wrapping (reuse existing connections)
    """

    def __init__(self, connection_manager: SnowflakeConnectionManager | None = None) -> None:
        self._connection_manager = connection_manager or SnowflakeConnectionManager()

    @property
    def dialect(self) -> DatabaseDialect:
        return DatabaseDialect.SNOWFLAKE

    def build_config(self, raw: dict[str, Any]) -> SnowflakeNamedConnection | SessionTokenConfig:
        """Build configuration from raw dict.
        
        Supports multiple connection modes:
        - "name": Named connection from TOML files
        - "session_token": Session token-based connection  
        - Other modes: Passed through to SnowflakeNamedConnection
        
        Args:
            raw: Configuration dictionary with 'mode' key
            
        Returns:
            Either SnowflakeNamedConnection or SessionTokenConfig based on mode
            
        Raises:
            ValueError: If session_token configuration is invalid
        """
        mode = raw.get("mode", ConnectionMode.NAME.value)
        
        if mode == ConnectionMode.SESSION_TOKEN.value:
            # Parse and validate session token config (raises ValueError if invalid)
            return SessionTokenConfig.from_dict(raw)
        
        # Named connection or other modes - delegate to library
        return SnowflakeNamedConnection(
            mode=mode,
            name=raw.get("name", ""),
        )

    def create_literal_formatter(self) -> SnowflakeLiteralFormatter:
        return SnowflakeLiteralFormatter()

    def create_isolation(self, **kwargs):
        return SnowflakeCloneIsolation(
            source_database=kwargs.get("source_database", ""),
            procedure=kwargs.get("procedure", ""),
            params_hash=kwargs.get("params_hash", ""),
        )

    def create_introspector(self):
        return DefaultIntrospector()

    def create_executor(self, config: SnowflakeNamedConnection | SessionTokenConfig) -> SnowflakeExecutor:
        """Create executor with a new connection from configuration.
        
        Args:
            config: Either SnowflakeNamedConnection or SessionTokenConfig
            
        Returns:
            SnowflakeExecutor wrapping a new connection
        """
        connector = self._connection_manager.create_connector(config)
        return SnowflakeExecutor(connector)

    def create_executor_from_connection(
        self,
        connection: SnowflakeConnection,
        *,
        owns_connection: bool = False
    ) -> SnowflakeExecutor:
        """Create executor that reuses an existing Snowflake connection.
        
        This method enables connection sharing, avoiding redundant authentication
        and connection overhead. Particularly useful for:
        - Session token-based connections (desktop app integration)
        - Connection pooling scenarios
        - Test scenarios where setup costs are high
        
        Args:
            connection: Active Snowflake connection to reuse.
            owns_connection: If True, the executor will close the connection
                           when closed. If False (default), connection management
                           is left to the caller.
        
        Returns:
            SnowflakeExecutor wrapping the provided connection.
        
        Example:
            >>> # Establish connection (e.g., via session tokens)
            >>> conn = snowflake.connector.connect(
            ...     session_token=session_token,
            ...     master_token=master_token,
            ...     account="myaccount",
            ...     user="myuser"
            ... )
            >>> 
            >>> # Create executor that reuses this connection
            >>> factory = SnowflakeExecutorFactory()
            >>> executor = factory.create_executor_from_connection(conn)
            >>> 
            >>> # Use executor for validation (shares the connection)
            >>> result = executor.execute("CALL my_procedure()")
            >>> 
            >>> # Executor close won't close connection (owns_connection=False)
            >>> executor.close()
            >>> conn.close()  # Caller manages connection lifecycle
        """
        connector = self._connection_manager.wrap_connection(
            connection,
            owns_connection=owns_connection
        )
        return SnowflakeExecutor(connector)
