"""
Snowflake database connection implementation.

This module provides the SnowflakeDataSource class for connecting to and
executing queries against Snowflake databases using the snowflake-connector-python
library. This implementation is thread-safe using thread-local storage for connections.
"""

import os
import threading

from collections.abc import Generator
from contextlib import contextmanager
from typing import Any


# Snowflake error number for "Session no longer exists. New login required."
_SF_ERRNO_SESSION_EXPIRED = 390111

from dependency_injector.wiring import Provide, inject

from data_exchange_agent.__version__ import __version__
from data_exchange_agent.config.manager import ConfigManager
from data_exchange_agent.config.sections.connections.cloud_storages.snowflake import (
    SnowflakeConnectionConfig,
    SnowflakeConnectionNameConfig,
)
from data_exchange_agent.constants import config_keys, container_keys
from data_exchange_agent.constants.connection_types import SPCS_CONNECTION_NAME
from data_exchange_agent.interfaces.data_source import DataSourceInterface
from data_exchange_agent.utils.sf_logger import SFLogger
from shared.telemetry import (
    detect_runtime_environment,
    get_telemetry_client,
    initialize_telemetry,
)


# SPCS environment constants
# Note: SNOWFLAKE_HOST, SNOWFLAKE_ACCOUNT, SNOWFLAKE_DATABASE, SNOWFLAKE_SCHEMA are
# automatically provided by SPCS. SNOWFLAKE_WAREHOUSE is NOT provided by SPCS but
# can be set manually or via QUERY_WAREHOUSE parameter in CREATE SERVICE.
SPCS_TOKEN_FILE_PATH = "/snowflake/session/token"
SPCS_ENV_HOST = "SNOWFLAKE_HOST"
SPCS_ENV_ACCOUNT = "SNOWFLAKE_ACCOUNT"
SPCS_ENV_DATABASE = "SNOWFLAKE_DATABASE"
SPCS_ENV_SCHEMA = "SNOWFLAKE_SCHEMA"
SPCS_ENV_WAREHOUSE = "SNOWFLAKE_WAREHOUSE"


class SnowflakeDataSource(DataSourceInterface):
    """
    A thread-safe Snowflake implementation of the DataSourceInterface.

    This class provides functionality to connect to and execute queries against a Snowflake database.
    It manages connections using configuration from a Snowflake config file. Each thread gets its
    own connection to ensure thread safety.

    Attributes:
        connection_name (str | None): Name of connection configuration to use from config file
        _local (threading.local): Thread-local storage for per-thread connections

    """

    @inject
    def __init__(
        self,
        connection_name: str = None,
        logger: SFLogger = Provide[container_keys.SF_LOGGER],
        program_config: ConfigManager = Provide[container_keys.PROGRAM_CONFIG],
    ) -> None:
        """
        Initialize a new thread-safe SnowflakeDataSource.

        Args:
            connection_name (str, optional): Name of connection configuration to use.
                If None, uses default connection from config file.
            logger (SFLogger): Logger instance
            program_config (ConfigManager): Program configuration

        """
        self.logger = logger
        self.connection_name = connection_name
        if not self.connection_name:
            target_connections = program_config[config_keys.CONNECTIONS__TARGET]
            target_connections = {
                k: v
                for k, v in target_connections.items()
                if isinstance(v, SnowflakeConnectionConfig)
            }

            # Get the first Snowflake connection name configuration
            snowflake_connection_name_config = next(
                (
                    v
                    for v in target_connections.values()
                    if isinstance(v, SnowflakeConnectionNameConfig)
                ),
                None,
            )

            if snowflake_connection_name_config:
                self.connection_name = snowflake_connection_name_config.connection_name

        # Thread-local storage for per-thread connections
        self._local = threading.local()

    def __enter__(self) -> "SnowflakeDataSource":
        """
        Enter the runtime context for the SnowflakeDataSource.

        Creates a connection for the current thread if one doesn't exist.

        Returns:
            SnowflakeDataSource: The SnowflakeDataSource instance

        """
        self.create_connection()
        return self

    def __exit__(
        self, exc_type: type | None, exc_val: Exception | None, exc_tb: Any
    ) -> None:
        """
        Exit the runtime context for the SnowflakeDataSource.

        Closes the connection for the current thread.

        Args:
            exc_type: Exception type if an exception was raised
            exc_val: Exception value if an exception was raised
            exc_tb: Exception traceback if an exception was raised

        """
        self.close_connection()

    def create_connection(self) -> None:
        """
        Create a new connection to Snowflake for the current thread if one doesn't exist.

        Uses connection details from ~/.snowflake/config.toml file.
        If connection_name is specified, uses that named connection config,
        otherwise uses the default connection config.

        This method is thread-safe - each thread gets its own connection.
        """
        # Simply call _get_thread_connection to ensure connection exists
        self._get_thread_connection()

    @contextmanager
    def get_cursor(self, dict_cursor: bool = True):
        """
        Context manager for getting a thread-safe cursor.

        Args:
            dict_cursor (bool): Whether to use a dictionary cursor (default: True)

        Yields:
            Cursor object for executing queries

        """
        from snowflake.connector import DictCursor

        connection = self._get_thread_connection()
        cursor_class = DictCursor if dict_cursor else None
        cursor = (
            connection.cursor(cursor_class) if cursor_class else connection.cursor()
        )
        try:
            yield cursor
        finally:
            cursor.close()

    def execute_statement(
        self, statement: str, params: tuple | None = None
    ) -> Generator[dict, None, None]:
        """
        Execute a SQL statement against Snowflake in a thread-safe manner.

        Creates a connection for the current thread if one doesn't exist, then executes
        the statement and yields the results. Each thread gets its own connection.

        If the session has been terminated server-side (error 390111), the stale
        connection is discarded, a new one is created, and the statement is retried once.

        Args:
            statement (str): The SQL statement to execute
            params (tuple | None): Optional bind parameters for the statement.
                Uses Snowflake's ``%s`` (pyformat) placeholder style.

        Yields:
            Results from executing the SQL statement

        """
        from snowflake.connector import errors as sf_errors

        def _run(cursor):
            if params:
                yield from cursor.execute(statement, params)
            else:
                yield from cursor.execute(statement)

        try:
            with self.get_cursor() as cursor:
                yield from _run(cursor)
        except sf_errors.ProgrammingError as e:
            if e.errno != _SF_ERRNO_SESSION_EXPIRED:
                raise
            self.logger.warning(
                "Snowflake session expired (390111). Reconnecting and retrying."
            )
            self._local.connection = None
            with self.get_cursor() as cursor:
                yield from _run(cursor)

    def execute_multiple_statements(
        self, statements: list[str]
    ) -> list[list[dict[str, Any]]]:
        """
        Execute multiple statements on the same connection (same thread).

        This method is useful when you need to execute multiple statements
        in sequence and want to reuse the same connection.

        If the session has been terminated server-side (error 390111), the stale
        connection is discarded, a new one is created, and all statements are retried once.

        Args:
            statements (list[str]): List of SQL statements to execute

        Returns:
            list[list[dict]]: List of results for each statement

        """
        from snowflake.connector import errors as sf_errors

        def _run():
            results = []
            with self.get_cursor() as cursor:
                for statement in statements:
                    result = list(cursor.execute(statement))
                    results.append(result)
            return results

        try:
            return _run()
        except sf_errors.ProgrammingError as e:
            if e.errno != _SF_ERRNO_SESSION_EXPIRED:
                raise
            self.logger.warning(
                "Snowflake session expired (390111). Reconnecting and retrying."
            )
            self._local.connection = None
            return _run()

    def is_closed(self) -> bool:
        """
        Check if the Snowflake connection for the current thread is closed.

        Returns:
            bool: True if connection is closed or doesn't exist, False otherwise

        """
        if not hasattr(self._local, "connection") or self._local.connection is None:
            return True
        return self._local.connection.is_closed()

    def close_connection(self) -> None:
        """
        Close the Snowflake connection for the current thread if one exists.

        Closes the connection and removes it from thread-local storage.
        This is thread-safe - only the current thread's connection is closed.
        """
        if hasattr(self._local, "connection") and self._local.connection is not None:
            try:
                self._local.connection.close()
            except Exception as e:
                # Ignore errors during close, log error
                self.logger.error("Failed to close Snowflake connection.", exception=e)
            finally:
                self._local.connection = None

    def _get_thread_connection(self):
        """
        Get or create a connection for the current thread.

        Returns:
            snowflake.connector.SnowflakeConnection: Connection for current thread

        """
        conn = getattr(self._local, "connection", None)

        # Recreate the connection if it doesn't exist or has been closed.
        # Note: Snowflake connections can be closed externally; relying on `None` alone
        # isn't sufficient for long-running agents.
        should_create = conn is None
        if not should_create:
            try:
                should_create = conn.is_closed()
            except Exception:
                # If the connection object is in a bad state, recreate it.
                should_create = True

        if not should_create:
            return conn

        # Best-effort cleanup of any previous connection handle
        if conn is not None:
            try:
                conn.close()
            except Exception:
                pass

        self._local.connection = self._create_connection()

        # Initialize telemetry automatically for this connection
        self._initialize_telemetry_if_needed()
        return self._local.connection

    def _create_connection(self):
        """Create a new connection to Snowflake for the current thread if one doesn't exist."""
        import snowflake.connector

        if self.connection_name is None:
            # Use default connection from snowflake config file (~/.snowflake/config.toml)
            return snowflake.connector.connect()
        elif self.connection_name == SPCS_CONNECTION_NAME:
            # Use Snowflake-provided SPCS credentials (OAuth token + environment variables)
            return self._create_spcs_connection()
        else:
            # Use specified named connection from config file
            return snowflake.connector.connect(connection_name=self.connection_name)

    def _initialize_telemetry_if_needed(self):
        """Initialize telemetry for the connection if it's not already initialized."""
        try:
            telemetry_client = get_telemetry_client()
            if telemetry_client is not None and telemetry_client.is_enabled():
                return

            env = os.getenv("ENVIRONMENT", "prod")
            runtime = detect_runtime_environment()
            initialize_telemetry(
                conn=self._local.connection,
                app_name="data_exchange_agent",
                app_version=__version__,
                env=env,
                runtime=runtime,
            )
            self.logger.debug(
                f"Telemetry initialized (env={env}, runtime={runtime.value})"
            )
        except Exception as e:
            # Don't fail connection creation if telemetry fails
            self.logger.debug(f"Telemetry initialization skipped: {e}")

    def _create_spcs_connection(self):
        """
        Create a Snowflake connection using SPCS-provided credentials.

        When running in Snowpark Container Services, Snowflake provides:
        - An OAuth token at /snowflake/session/token
        - Environment variables: SNOWFLAKE_HOST, SNOWFLAKE_ACCOUNT, SNOWFLAKE_DATABASE, SNOWFLAKE_SCHEMA

        Note: SNOWFLAKE_WAREHOUSE is NOT provided by SPCS but can be set manually
        or via QUERY_WAREHOUSE parameter in CREATE SERVICE/EXECUTE JOB SERVICE.

        Returns:
            snowflake.connector.SnowflakeConnection: Connection using SPCS credentials

        Raises:
            FileNotFoundError: If the SPCS token file is not found
            ValueError: If required environment variables are not set

        See: https://docs.snowflake.com/en/developer-guide/snowpark-container-services/spcs-execute-sql

        """
        import snowflake.connector

        # Read the OAuth token from the SPCS-provided file
        token = self._get_spcs_login_token()

        # Get required environment variables
        host = os.getenv(SPCS_ENV_HOST)
        account = os.getenv(SPCS_ENV_ACCOUNT)

        if not host:
            raise ValueError(f"SPCS environment variable {SPCS_ENV_HOST} is not set.")
        if not account:
            raise ValueError(
                f"SPCS environment variable {SPCS_ENV_ACCOUNT} is not set."
            )

        # Get optional environment variables for database, schema, and warehouse context
        database = os.getenv(SPCS_ENV_DATABASE)
        schema = os.getenv(SPCS_ENV_SCHEMA)
        warehouse = os.getenv(SPCS_ENV_WAREHOUSE)

        self.logger.info(
            f"Creating SPCS connection to Snowflake (host={host}, account={account}, "
            f"database={database or 'not set'}, schema={schema or 'not set'}, "
            f"warehouse={warehouse or 'not set'})"
        )

        connection_params = {
            "host": host,
            "account": account,
            "token": token,
            "authenticator": "oauth",
        }

        # Add optional database, schema, and warehouse if provided
        if database:
            connection_params["database"] = database
        if schema:
            connection_params["schema"] = schema
        if warehouse:
            connection_params["warehouse"] = warehouse

        return snowflake.connector.connect(**connection_params)

    def _get_spcs_login_token(self) -> str:
        """
        Read the SPCS OAuth token from the Snowflake-provided file.

        Returns:
            str: The OAuth token

        Raises:
            FileNotFoundError: If the token file does not exist

        """
        try:
            with open(SPCS_TOKEN_FILE_PATH) as f:
                return f.read()
        except FileNotFoundError:
            raise FileNotFoundError(
                f"SPCS token file not found at {SPCS_TOKEN_FILE_PATH}. "
                "Ensure this service is running in Snowpark Container Services."
            ) from None
