"""Snowflake connection adapter.

Implements the DBConnection protocol for contract data loading using Snowflake's
connector and private key authentication.
"""

from __future__ import annotations

import logging
import os
from collections.abc import Callable, Generator
from contextlib import contextmanager
from enum import IntEnum
from functools import lru_cache, wraps
from typing import ParamSpec, TypeVar

import snowflake.connector
from pydantic import BaseModel, ConfigDict, Field
from snowflake.connector import (
    SnowflakeConnection as RawSnowflakeConnection,
)
from snowflake.connector.cursor import SnowflakeCursor
from snowflake.connector.errors import (
    OperationalError as SnowflakeOperationalError,
)
from snowflake.connector.s3_storage_client import SnowflakeS3RestClient

from src.connectors.db import Connection
from src.connectors.snowflake.utils import get_private_key
from src.errors import TransientError

# Disable Snowflake connector logging
logging.getLogger('snowflake.connector.connection').propagate = False

APPLICATION_NAME = 'lambda-abacus-earnings-transfer'
_NON_LOCAL_ENVS = ('qa', 'uat', 'prod')


class SnowflakeConfig(BaseModel):
    """Configuration model for a Snowflake database connection."""

    model_config = ConfigDict(frozen=True, populate_by_name=True)

    # Connection
    host: str | None = Field(default=None, description='Your Snowflake host.')
    account: str = Field(
        ...,
        min_length=1,
        description='Your Snowflake account ID (e.g., xy12345.region.aws).',
    )
    user: str = Field(..., min_length=1, description='Username for authentication.')
    private_key: str | None = Field(
        default=None,
        description='The private key content (PEM).',
    )
    private_key_path: str | None = Field(
        default=None,
        description='The absolute path to the .p8/.pem file.',
    )
    private_key_passphrase: str | None = Field(
        default=None, description='The passphrase for the encrypted private key.'
    )
    warehouse: str = Field(
        ..., min_length=1, description='The name of the compute warehouse to use.'
    )
    database: str = Field(
        ..., min_length=1, description='The name of the database to connect to.'
    )
    schema_name: str = Field(
        ...,
        min_length=1,
        description='The schema within the database to use.',
        alias='schema',
    )
    role: str | None = Field(
        default=None, description='The access role to use for the connection.'
    )
    connection_timeout: int = Field(
        default=10, description='Connection timeout in seconds.'
    )
    # Settings
    autocommit: bool = Field(
        default=False, description='Whether transactions should commit automatically.'
    )


class SnowflakeConnection(Connection):
    """Snowflake database connection adapter."""

    def __init__(self, conn: RawSnowflakeConnection):
        """Initialize the adapter with a raw Snowflake connection.

        Args:
            conn: An active snowflake connection.
        """
        self._conn = conn

    def close(self) -> None:
        """Close the connection."""
        self._conn.close()

    def commit(self) -> None:
        """Commit the current transaction to the database."""
        self._conn.commit()

    def cursor(self) -> SnowflakeCursor:
        """Return a raw Snowflake cursor object."""
        return self._conn.cursor()

    def rollback(self) -> None:
        """Roll back the current transaction, discarding changes."""
        self._conn.rollback()


class SnowflakeConnectionFactory:
    """Factory for creating Snowflake database connections with configuration management."""

    def __init__(self, config: SnowflakeConfig) -> None:
        """Initialize the connection factory with Snowflake configuration.

        Args:
            config: SnowflakeConfig instance containing database connection parameters.
        """
        self._config = config

    @contextmanager
    def connection(self) -> Generator[SnowflakeConnection, None, None]:
        """Create and yield a Snowflake database connection as a context manager.

        The connection is automatically closed when the context exits.

        Yields:
            SnowflakeConnection: Snowflake database adapter.

        Raises:
            TransientError: If connection fails due to transient database errors.
            SnowflakeOperationalError: If connection fails due to non-transient errors.
        """
        with snowflake_connection(self._config) as conn:
            yield SnowflakeConnection(conn)


class SnowflakeErrorCode(IntEnum):
    """Known Snowflake error codes, mostly related to transient errors."""

    # Query issued in another transaction is running for too long
    QUERY_TIMEOUT = 57002
    # Warehouse is suspended or cannot be found (often seen during warm-up)
    WAREHOUSE_SUSPENDED = 604
    # Connection time-out or network issue
    CONNECTION_ERROR = 252001
    # Statement cancelled by the user/system (can happen during resource contention)
    STATEMENT_CANCELLED = 607
    # Transient network/service error
    INTERNAL_SERVICE_ERROR = 300001


P = ParamSpec('P')
T = TypeVar('T')


TRANSIENT_ERROR_CODES: frozenset[int] = frozenset(
    {
        SnowflakeErrorCode.QUERY_TIMEOUT,
        SnowflakeErrorCode.WAREHOUSE_SUSPENDED,
        SnowflakeErrorCode.CONNECTION_ERROR,
        SnowflakeErrorCode.STATEMENT_CANCELLED,
        SnowflakeErrorCode.INTERNAL_SERVICE_ERROR,
    }
)


def _try_transient_error(e: SnowflakeOperationalError) -> None:
    """Check if an OperationalError is caused by a transient condition.

    If transient, a custom TransientError is raised.

    Args:
        e: The caught snowflake.connector.errors.OperationalError exception.

    Raises:
        TransientError: If the error code matches one in TRANSIENT_ERROR_CODES.
    """
    if hasattr(e, 'errno') and e.errno in TRANSIENT_ERROR_CODES:
        raise TransientError(f'Snowflake operation failed: {e}') from e

    if 'failed to establish a connection' in str(e).lower():
        raise TransientError(f'Snowflake connection failed: {e}') from e


def handle_snowflake_errors(func: Callable[P, T]) -> Callable[P, T]:
    """Decorate database operations and handle transient errors.

    If a transient error occurs, it is converted into a TransientError to
    allow for retry logic. Non-transient errors are re-raised as is.
    """

    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        try:
            return func(*args, **kwargs)
        except SnowflakeOperationalError as e:
            _try_transient_error(e)
            raise  # Re-raise if it wasn't transient

    return wrapper


@contextmanager
def snowflake_connection(
    config: SnowflakeConfig,
) -> Generator[RawSnowflakeConnection, None, None]:
    """Context manager for connection objects.

    Args:
        config: Snowflake configuration object.

    Yields:
        RawSnowflakeConnection: An active Snowflake connection.

    Raises:
        TransientError: If the connection attempt fails due to a transient error.
        SnowflakeOperationalError: If the connection attempt fails due to a non-transient error.
    """
    private_key_bytes = get_private_key(
        config.private_key, config.private_key_path, config.private_key_passphrase
    )

    connection_timeout = config.connection_timeout or 10

    connect_kwargs: dict = {
        'account': config.account,
        'user': config.user,
        'private_key': private_key_bytes,
        'warehouse': config.warehouse,
        'database': config.database,
        'schema': config.schema_name,
        'role': config.role,
        'connection_timeout': connection_timeout,
        'autocommit': config.autocommit,
    }
    if config.host is not None:
        connect_kwargs['host'] = config.host

    try:
        conn = snowflake.connector.connect(**connect_kwargs)
    except SnowflakeOperationalError as e:
        _try_transient_error(e)
        raise

    try:
        yield conn
    finally:
        if conn:
            conn.close()


@lru_cache(maxsize=1)
def build_snowflake_config() -> SnowflakeConfig:
    """Build a SnowflakeConfig from environment variables.

    For non-local environments (qa, uat, prod), credentials are fetched
    from AWS Secrets Manager via LambdaSecretsManager.

    Returns:
        SnowflakeConfig with resolved credentials.
    """
    environment = os.environ.get('ENVIRONMENT', 'dev')
    passphrase = os.environ.get('SNOWFLAKE_KEY_PASSPHRASE')
    private_key = os.environ.get('SNOWFLAKE_PRIVATE_KEY')
    private_key_path = os.environ.get('SNOWFLAKE_PRIVATE_KEY_PATH')

    if environment in _NON_LOCAL_ENVS:
        from secrets_manager.lambda_ext import LambdaSecretsManager

        client = LambdaSecretsManager(
            environment=environment,
            service_name=APPLICATION_NAME,
        )
        passphrase = client.get_cred('SNOWFLAKE_KEY_PASSPHRASE')
        private_key = client.get_cred('SNOWFLAKE_PRIVATE_KEY')

    return SnowflakeConfig(
        host=os.environ.get('SNOWFLAKE_HOST'),
        account=os.environ['SNOWFLAKE_ACCOUNT'],
        user=os.environ['SNOWFLAKE_USER'],
        private_key=private_key,
        private_key_path=private_key_path,
        private_key_passphrase=passphrase,
        warehouse=os.environ['SNOWFLAKE_WAREHOUSE'],
        database=os.environ['SNOWFLAKE_DATABASE'],
        schema_name=os.environ['SNOWFLAKE_SCHEMA'],
        role=os.environ.get('SNOWFLAKE_ROLE'),
    )


# MONKEY PATCHES


def _get_bucket_accelerate_config_patch(
    self: SnowflakeS3RestClient, bucket_name: str
) -> bool:
    """Disable S3 transfer acceleration.

    Args:
        self: A SnowflakeS3RestClient instance.
        bucket_name: S3 bucket name.
    """
    return False


setattr(
    SnowflakeS3RestClient,
    '_get_bucket_accelerate_config',
    _get_bucket_accelerate_config_patch,
)


def _transfer_accelerate_config_patch(
    self: SnowflakeS3RestClient, use_accelerate_endpoint: bool | None = None
) -> bool:
    """Disable S3 transfer acceleration.

    Args:
        self: A SnowflakeS3RestClient instance.
        use_accelerate_endpoint: Boolean.
    """
    self.endpoint = (
        f'https://{self.s3location.bucket_name}.s3.{self.region_name}.amazonaws.com'
    )
    return False


setattr(
    SnowflakeS3RestClient,
    'transfer_accelerate_config',
    _transfer_accelerate_config_patch,
)