"""Snowflake connection configuration.

This module defines Pydantic models for configuring a Snowflake connection
backed by a Snowpark :class:`snowflake.snowpark.Session`.

You typically use these models through :class:`relationalai.config.Config` by
adding a Snowflake connection entry under ``connections``. The connection entry
can be provided as either:

- A Python dict (recommended for config files), or
- An instance of one of the authenticator classes (useful when building configs
    programmatically).

The authenticator is selected via the ``authenticator`` field in the config.
Supported values and their corresponding classes are:

- ``"username_password"`` → :class:`~relationalai.config.UsernamePasswordAuth`
- ``"username_password_mfa"`` → :class:`~relationalai.config.UsernamePasswordMFAAuth`
- ``"externalbrowser"`` → :class:`~relationalai.config.ExternalBrowserAuth`
- ``"jwt"`` → :class:`~relationalai.config.JWTAuth` (key-pair authentication)
- ``"oauth"`` → :class:`~relationalai.config.OAuthAuth`
- ``"programmatic_access_token"`` → :class:`~relationalai.config.ProgrammaticAccessTokenAuth`

These authenticator classes are re-exported from :mod:`relationalai.config` for
convenience.

All Snowflake authenticators share common fields like ``account`` and
``warehouse``, and optional fields like ``role``, ``database``, and ``schema``.
Calling :meth:`~relationalai.config.connections.snowflake.SnowflakeConnectionBase.get_session`
creates and caches a Snowpark session on the connection instance; use
:meth:`relationalai.config.connections.base.BaseConnection.clear_session_cache`
to force a fresh session.

Examples
--------
Configure a Snowflake connection using a dict:

>>> from relationalai.config import Config
>>> cfg = Config(connections={
...     "sf": {
...         "type": "snowflake",
...         "authenticator": "username_password",
...         "account": "my_account",
...         "warehouse": "my_warehouse",
...         "user": "my_user",
...         "password": "my_password",
...     }
... })

Configure a Snowflake connection using an authenticator class:

>>> from relationalai.config import Config, UsernamePasswordAuth
>>> cfg = Config(connections={
...     "sf": UsernamePasswordAuth(
...         account="my_account",
...         warehouse="my_warehouse",
...         user="my_user",
...         password="my_password",
...     )
... })
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from pathlib import Path
from typing import Literal, Any, Annotated, Dict, Union, cast, TYPE_CHECKING
from pydantic import Field, SecretStr, field_validator, model_validator

if TYPE_CHECKING:
    import snowflake.snowpark

from ...util.docutils import include_in_docs

from .base import BaseConnection


# include this module in documentation
__include_in_docs__ = True


def _load_private_key_bytes(
    private_key_path: str,
    private_key_passphrase: str | SecretStr | None,
) -> bytes:
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import serialization

    password = None
    if isinstance(private_key_passphrase, SecretStr):
        password = private_key_passphrase.get_secret_value().encode()
    elif private_key_passphrase is not None:
        password = private_key_passphrase.encode()

    with open(private_key_path, "rb") as key_file:
        private_key = serialization.load_pem_private_key(
            key_file.read(),
            password=password,
            backend=default_backend(),
        )

    return private_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )


@include_in_docs
class SnowflakeConnectionBase(BaseConnection, ABC):
    """Configure a Snowflake Snowpark session.

    This abstract base class defines common Snowflake connection fields and
    provides :meth:`get_session` to create a :class:`~snowflake.snowpark.Session`.
    Use one of the concrete authenticators (for example
    :class:`~relationalai.config.connections.snowflake.UsernamePasswordAuth`).

    The session is cached on the instance; call
    :meth:`~relationalai.config.connections.base.BaseConnection.clear_session_cache`
    to force a fresh session on the next access.

    Attributes
    ----------
    account : str
        Snowflake account identifier.
    warehouse : str
        Snowflake warehouse name.
    rai_app_name : str, optional
        RAI application name in Snowflake (default: "RELATIONALAI").
    role : str, optional
        Role to set on the session.
    database : str, optional
        Default database.
    schema_ : str, optional
        Default schema (set via ``schema=...``).

    Notes
    -----
    Do not instantiate this class directly.
    """

    type: Literal["snowflake"] = "snowflake"

    # Required fields common to all Snowflake authenticators
    account: str = Field(..., description="Snowflake account identifier")
    warehouse: str = Field(..., description="Snowflake warehouse name")

    # Optional fields common to all Snowflake authenticators
    rai_app_name: str = Field(default="RELATIONALAI", description="RAI application name in Snowflake")
    role: str | None = Field(default=None, description="Snowflake role")
    database: str | None = Field(default=None, description="Default database")
    schema_: str | None = Field(default=None, alias="schema", description="Default schema")

    @abstractmethod
    def _get_connection_params(self) -> Dict[str, Any]:
        """Return authenticator-specific connection parameters."""
        pass

    def _create_new_session(self) -> snowflake.snowpark.Session:
        """Create a new Snowflake session using connection params."""
        from snowflake.snowpark import Session
        connection_params = self._get_connection_params()
        self._add_common_params(connection_params)
        return Session.builder.configs(connection_params).create()

    @include_in_docs
    def get_session(self) -> snowflake.snowpark.Session:
        """Return a Snowpark session for this connection.

        The session is created on first access from this connection's configuration
        and cached on the instance. Call
        :meth:`~relationalai.config.connections.base.BaseConnection.clear_session_cache`
        to force a fresh session on the next access.

        Returns
        -------
        snowflake.snowpark.Session
            The configured Snowpark session.
        """
        from snowflake.snowpark.context import get_active_session

        if self._cached_session is None:
            try:
                self._cached_session = get_active_session()
            except Exception:
                self._cached_session = self._create_new_session()

        return self._cached_session

    def _add_common_params(self, connection_params: Dict[str, Any]) -> None:
        if self.role:
            connection_params["role"] = self.role
        if self.database:
            connection_params["database"] = self.database
        if self.schema_:
            connection_params["schema"] = self.schema_


@include_in_docs
class UsernamePasswordAuth(SnowflakeConnectionBase):
    """Configure Snowflake access using a username and password.

    Use this authenticator when you sign in to Snowflake with a password.
    You can instantiate this class directly (for example, when building a
    :func:`~relationalai.config.create_config` programmatically), or provide a config dict
    and let :func:`~relationalai.config.create_config` validate and create it.

    Instantiating this class is optional; you can also pass an equivalent dict to
    :func:`~relationalai.config.create_config`.

    When loading from a config dict or file, set ``authenticator="username_password"``.

    Parameters
    ----------
    account : str
        Snowflake account identifier.
    warehouse : str
        Snowflake warehouse name.
    user : str
        Snowflake username.
    password : str or SecretStr
        Snowflake password.

    Examples
    --------
    Create a programmatic config using this authenticator class:

    >>> from relationalai.config import Config, UsernamePasswordAuth
    >>> cfg = Config(connections={
    ...     "sf": UsernamePasswordAuth(
    ...         account="my_account",
    ...         warehouse="my_warehouse",
    ...         user="my_user",
    ...         password="my_password",
    ...     )
    ... })

    Create a programmatic config using a plain dict (no authenticator import):

    >>> from relationalai.config import Config
    >>> cfg = Config(connections={
    ...     "sf": {
    ...         "type": "snowflake",
    ...         "authenticator": "username_password",
    ...         "account": "my_account",
    ...         "warehouse": "my_warehouse",
    ...         "user": "my_user",
    ...         "password": "my_password",
    ...     }
    ... })
    """

    authenticator: Literal["username_password"] = "username_password"

    # Required Snowflake fields
    user: str = Field(..., description="Snowflake username")
    password: str | SecretStr = Field(..., description="Snowflake password")

    @field_validator('password', mode='before')
    @classmethod
    def convert_password_to_secret(cls, v: Any) -> SecretStr:
        if isinstance(v, str):
            return SecretStr(v)
        return v

    def _get_connection_params(self) -> Dict[str, Any]:
        return {
            "user": self.user,
            "password": cast(SecretStr, self.password).get_secret_value(),
            "account": self.account,
            "warehouse": self.warehouse,
        }


@include_in_docs
class UsernamePasswordMFAAuth(SnowflakeConnectionBase):
    """Configure Snowflake access using a username/password and an MFA passcode.

    Use this authenticator when your Snowflake sign-in requires a time-based MFA
    code. When loading from a config dict or file, set
    ``authenticator="username_password_mfa"``.

    Instantiating this class is optional; you can also pass an equivalent dict to
    :func:`~relationalai.config.create_config`.

    Parameters
    ----------
    account : str
        Snowflake account identifier.
    warehouse : str
        Snowflake warehouse name.
    user : str
        Snowflake username.
    password : str or SecretStr
        Snowflake password.
    passcode : str, optional
        MFA passcode (typically a 6-digit code).

    Examples
    --------
    Create a programmatic config using this authenticator class:

    >>> from relationalai.config import Config, UsernamePasswordMFAAuth
    >>> cfg = Config(connections={
    ...     "sf": UsernamePasswordMFAAuth(
    ...         account="my_account",
    ...         warehouse="my_warehouse",
    ...         user="my_user",
    ...         password="my_password",
    ...         passcode="123456",
    ...     )
    ... })

    Create a programmatic config using a plain dict (no authenticator import):

    >>> from relationalai.config import Config
    >>> cfg = Config(connections={
    ...     "sf": {
    ...         "type": "snowflake",
    ...         "authenticator": "username_password_mfa",
    ...         "account": "my_account",
    ...         "warehouse": "my_warehouse",
    ...         "user": "my_user",
    ...         "password": "my_password",
    ...         "passcode": "123456",
    ...     }
    ... })
    """

    authenticator: Literal["username_password_mfa"] = "username_password_mfa"

    # Required Snowflake fields
    user: str = Field(..., description="Snowflake username")
    password: str | SecretStr = Field(..., description="Snowflake password")
    passcode: str | None = Field(default=None, description="MFA passcode (6-digit code)")

    @field_validator('password', mode='before')
    @classmethod
    def convert_password_to_secret(cls, v: Any) -> SecretStr:
        """Convert plain str to SecretStr for internal use."""
        if isinstance(v, str):
            return SecretStr(v)
        return v

    def _get_connection_params(self) -> Dict[str, Any]:
        params = {
            "user": self.user,
            "password": cast(SecretStr, self.password).get_secret_value(),
            "account": self.account,
            "warehouse": self.warehouse,
            "authenticator": self.authenticator,
        }
        if self.passcode is not None:
            params["passcode"] = self.passcode
        return params


@include_in_docs
class ExternalBrowserAuth(SnowflakeConnectionBase):
    """Configure Snowflake access using browser-based sign-in.

    Use this authenticator when you sign in to Snowflake via an interactive
    browser flow (for example, when your account uses SSO). When loading from a
    config dict or file, set ``authenticator="externalbrowser"``.

    Instantiating this class is optional; you can also pass an equivalent dict to
    :func:`~relationalai.config.create_config`.

    Parameters
    ----------
    account : str
        Snowflake account identifier.
    warehouse : str
        Snowflake warehouse name.
    user : str
        Snowflake username.

    Notes
    -----
    This authenticator requires an interactive browser session.

    Examples
    --------
    Create a programmatic config using this authenticator class:

    >>> from relationalai.config import Config, ExternalBrowserAuth
    >>> cfg = Config(connections={
    ...     "sf": ExternalBrowserAuth(
    ...         account="my_account",
    ...         warehouse="my_warehouse",
    ...         user="my_user",
    ...     )
    ... })

    Create a programmatic config using a plain dict (no authenticator import):

    >>> from relationalai.config import Config
    >>> cfg = Config(connections={
    ...     "sf": {
    ...         "type": "snowflake",
    ...         "authenticator": "externalbrowser",
    ...         "account": "my_account",
    ...         "warehouse": "my_warehouse",
    ...         "user": "my_user",
    ...     }
    ... })
    """

    authenticator: Literal["externalbrowser"] = "externalbrowser"

    # Required Snowflake fields
    user: str = Field(..., description="Snowflake username")

    def _get_connection_params(self) -> Dict[str, Any]:
        return {
            "user": self.user,
            "account": self.account,
            "warehouse": self.warehouse,
            "authenticator": self.authenticator,
        }


@include_in_docs
class JWTAuth(SnowflakeConnectionBase):
    """Configure Snowflake access using key-pair (JWT) authentication.

    Use this authenticator when you sign in to Snowflake with a private key
    instead of a password. When loading from a config dict or file, set
    ``authenticator="jwt"`` or the legacy alias ``"snowflake_jwt"``.

    Instantiating this class is optional; you can also pass an equivalent dict to
    :func:`~relationalai.config.create_config`.

    Parameters
    ----------
    account : str
        Snowflake account identifier.
    warehouse : str
        Snowflake warehouse name.
    user : str
        Snowflake username.
    private_key_path : str
        Path to a PEM-encoded private key file.
    private_key_passphrase : str or SecretStr, optional
        Passphrase for the private key file, if encrypted.

    Notes
    -----
    The private key file is read from ``private_key_path`` when creating the
    Snowpark session.

    Examples
    --------
    Create a programmatic config using this authenticator class:

    >>> from relationalai.config import Config, JWTAuth
    >>> cfg = Config(connections={
    ...     "sf": JWTAuth(
    ...         account="my_account",
    ...         warehouse="my_warehouse",
    ...         user="my_user",
    ...         private_key_path="/path/to/key.pem",
    ...     )
    ... })

    Create a programmatic config using a plain dict (no authenticator import):

    >>> from relationalai.config import Config
    >>> cfg = Config(connections={
    ...     "sf": {
    ...         "type": "snowflake",
    ...         "authenticator": "jwt",
    ...         "account": "my_account",
    ...         "warehouse": "my_warehouse",
    ...         "user": "my_user",
    ...         "private_key_path": "/path/to/key.pem",
    ...     }
    ... })
    """

    authenticator: Literal["jwt", "snowflake_jwt"] = "jwt"

    # Required Snowflake fields
    user: str = Field(..., description="Snowflake username")
    private_key_path: str = Field(..., description="Path to private key file for JWT (.p8/.pem)")

    # Optional Snowflake fields
    private_key_passphrase: str | SecretStr | None = Field(None, description="Passphrase for private key (optional)")

    @field_validator("private_key_passphrase", mode="before")
    @classmethod
    def convert_passphrase_to_secret(cls, v: Any) -> SecretStr | None:
        """Convert plain str to SecretStr for internal use."""
        if v is None:
            return None
        if isinstance(v, str):
            return SecretStr(v)
        return v

    def _get_connection_params(self) -> Dict[str, Any]:
        raise NotImplementedError("JWTAuth uses custom _create_new_session() implementation")

    def _create_new_session(self) -> snowflake.snowpark.Session:
        from snowflake.snowpark import Session

        pkb = _load_private_key_bytes(self.private_key_path, self.private_key_passphrase)

        connection_params: Dict[str, Any] = {
            "user": self.user,
            "account": self.account,
            "warehouse": self.warehouse,
            "private_key": pkb,
        }

        self._add_common_params(connection_params)
        return Session.builder.configs(connection_params).create()

@include_in_docs
class OAuthAuth(SnowflakeConnectionBase):
    """Configure Snowflake access using an OAuth access token.

    Use this authenticator when you already have an OAuth token for Snowflake.
    When loading from a config dict or file, set ``authenticator="oauth"``.

    Instantiating this class is optional; you can also pass an equivalent dict to
    :func:`~relationalai.config.create_config`.

    Parameters
    ----------
    account : str
        Snowflake account identifier.
    warehouse : str
        Snowflake warehouse name.
    token : str or SecretStr
        OAuth access token.

    Examples
    --------
    Create a programmatic config using this authenticator class:

    >>> from relationalai.config import Config, OAuthAuth
    >>> cfg = Config(connections={
    ...     "sf": OAuthAuth(
    ...         account="my_account",
    ...         warehouse="my_warehouse",
    ...         token="my_oauth_token",
    ...     )
    ... })

    Create a programmatic config using a plain dict (no authenticator import):

    >>> from relationalai.config import Config
    >>> cfg = Config(connections={
    ...     "sf": {
    ...         "type": "snowflake",
    ...         "authenticator": "oauth",
    ...         "account": "my_account",
    ...         "warehouse": "my_warehouse",
    ...         "token": "my_oauth_token",
    ...     }
    ... })
    """

    authenticator: Literal["oauth"] = "oauth"

    # Required Snowflake fields
    token: str | SecretStr = Field(..., description="OAuth access token")

    @field_validator('token', mode='before')
    @classmethod
    def convert_token_to_secret(cls, v: Any) -> SecretStr:
        """Convert plain str to SecretStr for internal use."""
        if isinstance(v, str):
            return SecretStr(v)
        return v

    def _get_connection_params(self) -> Dict[str, Any]:
        return {
            "account": self.account,
            "warehouse": self.warehouse,
            "authenticator": "oauth",
            "token": cast(SecretStr, self.token).get_secret_value(),
        }


@include_in_docs
class ProgrammaticAccessTokenAuth(SnowflakeConnectionBase):
    """Configure Snowflake access using a Programmatic Access Token.

    Use this authenticator when you have a Snowflake Programmatic Access Token.
    When loading from a config dict or file, set
    ``authenticator="programmatic_access_token"``.

    Instantiating this class is optional; you can also pass an equivalent dict to
    :func:`~relationalai.config.create_config`.

    Parameters
    ----------
    account : str
        Snowflake account identifier.
    warehouse : str
        Snowflake warehouse name.
    token : str or SecretStr, optional
        Programmatic Access Token value. Either ``token`` or ``token_file_path``
        must be provided.
    token_file_path : str, optional
        Path to a file containing the Programmatic Access Token. Either ``token``
        or ``token_file_path`` must be provided.

    Examples
    --------
    Create a programmatic config using this authenticator class:

    >>> from relationalai.config import Config, ProgrammaticAccessTokenAuth
    >>> cfg = Config(connections={
    ...     "sf": ProgrammaticAccessTokenAuth(
    ...         account="my_account",
    ...         warehouse="my_warehouse",
    ...         token="my_pat",
    ...     )
    ... })

    Create a programmatic config using a plain dict (no authenticator import):

    >>> from relationalai.config import Config
    >>> cfg = Config(connections={
    ...     "sf": {
    ...         "type": "snowflake",
    ...         "authenticator": "programmatic_access_token",
    ...         "account": "my_account",
    ...         "warehouse": "my_warehouse",
    ...         "token": "my_pat",
    ...     }
    ... })
    """

    authenticator: Literal["programmatic_access_token"] = "programmatic_access_token"

    user: str = Field(..., description="Snowflake username")

    # Token fields — at least one must be provided
    token: str | SecretStr | None = Field(default=None, description="Programmatic Access Token value")
    token_file_path: str | None = Field(default=None, description="Path to file containing the Programmatic Access Token")

    @field_validator('token', mode='before')
    @classmethod
    def convert_token_to_secret(cls, v: Any) -> SecretStr | None:
        """Convert plain str to SecretStr for internal use."""
        if v is None:
            return None
        if isinstance(v, str):
            return SecretStr(v)
        return v

    @model_validator(mode='after')
    def require_token_or_file(self) -> ProgrammaticAccessTokenAuth:
        if self.token is None and self.token_file_path is None:
            raise ValueError("Either 'token' or 'token_file_path' must be provided for programmatic_access_token auth")
        return self

    def _get_connection_params(self) -> Dict[str, Any]:
        if self.token is not None:
            token_value = cast(SecretStr, self.token).get_secret_value()
        else:
            token_value = Path(cast(str, self.token_file_path)).expanduser().read_text().strip()
        return {
            "user": self.user,
            "account": self.account,
            "warehouse": self.warehouse,
            "authenticator": "PROGRAMMATIC_ACCESS_TOKEN",
            "token": token_value,
        }


# Union type for all Snowflake authenticators (for discriminated unions in Pydantic)
SnowflakeAuthenticator = Annotated[
    Union[
        UsernamePasswordAuth,
        UsernamePasswordMFAAuth,
        ExternalBrowserAuth,
        JWTAuth,
        OAuthAuth,
        ProgrammaticAccessTokenAuth
    ],
    Field(discriminator="authenticator")
]
"""Snowflake authenticator configuration type.

`SnowflakeAuthenticator` is a Pydantic discriminated union used when parsing a
Snowflake connection config. The concrete type is selected from the
``authenticator`` config field (for example, ``"username_password"`` or ``"jwt"``).
"""

# Export both the base class (for isinstance checks and type hints) and the union (for Pydantic validation)
SnowflakeConnection = SnowflakeConnectionBase
"""Alias for :class:`SnowflakeConnectionBase`.

This is the recommended type to use when referring to Snowflake connections in
type hints, ``isinstance`` checks, and when calling
:meth:`~relationalai.config.config.RAIConfig.get_connection`.

Examples
--------
>>> from relationalai.config import Config
>>> from relationalai.config import SnowflakeConnection
>>> cfg = Config(
...     connections={
...         "sf": {
...             "type": "snowflake",
...             "account": "my_account",
...             "warehouse": "my_warehouse",
...             "user": "my_user",
...             "password": "my_password"
...         }
...     }
... )
>>> conn = cfg.get_connection(SnowflakeConnection, name="sf")
"""

