"""Snowflake-backed dataset."""
from __future__ import annotations

import dataclasses as dc
import logging
import os
import sys
from typing import Any
from typing import ClassVar
from typing import Literal

import pydantic.dataclasses as pydc
import tomli
from pydantic import validate_call

from .base import _DBDatasetArgs
from .base import DatasetBase
from .filter import FilterMixin
from tempora.utils.dataclass import pydc_config
from tempora.utils.db import close_db_connection
from tempora.utils.db import ConnectionParams
from tempora.utils.db import open_db_connection
from tempora.utils.exceptions import TemporaClientError


log = logging.getLogger(__name__)


@dc.dataclass(frozen=True)
class _SnowflakeDatasetArgs(_DBDatasetArgs):
    ...


@pydc.dataclass(frozen=True, eq=False, config=pydc_config)
class SnowflakeDataset(DatasetBase, _SnowflakeDatasetArgs):
    """Class representing a Snowflake-backed dataset.

    Args:
        database: The name of the Snowflake database.
        db_schema: The name of the schema for the database.
        table: The name of the table under the given schema.
    """
    filtered_dataset_cls: ClassVar[type[FilteredSnowflakeDataset]]

    def __post_init__(self, _op_dict: dict[str, Any] | None) -> None:
        super().__post_init__(_op_dict=_op_dict)


@pydc.dataclass(frozen=True, eq=False, config=pydc_config)
class FilteredSnowflakeDataset(FilterMixin, SnowflakeDataset):
    """Class representing a filtered Snowflake-based dataset."""

    def __post_init__(self, _op_dict: dict[str, Any] | None) -> None:
        super().__post_init__(_op_dict=_op_dict)


SnowflakeDataset.filtered_dataset_cls = FilteredSnowflakeDataset


@validate_call
def connect_snowflake(
    user: str | None = None,
    password: str | None = None,
    account: str | None = None,
    auth_type: Literal['password', 'mfa', 'pat'] = 'password',
    auth_token: str | None = None,
    role: str | None = None,
    warehouse: str | None = None,
    connection_name: str | None = None,
) -> None:
    """Open a connection to Snowflake.

    Args:
        user: Login name for the user. If omitted, this must be specified using the
            SNOWFLAKE_USER env var.
        password: Password for the user. If omitted, and the auth_type is 'password'
            or 'mfa', must be specified using the SNOWFLAKE_PASSWORD env var.
        account: The Snowflake account identifier. If omitted, this must be specified
            using the SNOWFLAKE_ACCOUNT env var.
        auth_type: Authentication mode, either 'password' (default username + password
            authentication), 'mfa' (username + password with MFA), or 'pat' (use a
            programmatic access token).
        auth_token: The authentication token to be used with 'mfa' or 'pat'. If not
            provided uses SNOWFLAKE_AUTH_TOKEN.
        role: The role to be used for authentication.
        warehouse: Name of the default warehouse to use.
        connection_name: Name of the Snowflake connection definition to use within
            the connections.toml file. Any passed in arguments will override the
            associated values from the connection definition, if one is specified.

    Returns:
        None
    """
    params = _get_snowflake_connection_params(
        user=user, password=password, account=account,
        auth_type=auth_type, auth_token=auth_token, role=role,
        warehouse=warehouse, connection_name=connection_name,
    )
    open_db_connection(dbms_name='Snowflake', params=params)
    return None


def _get_snowflake_connection_params(
    **kwargs: str | dict[str, str] | None,
) -> ConnectionParams:
    connection_params: ConnectionParams = {}
    connection_name = kwargs.get('connection_name')

    if connection_name:
        # Get the parameters of the named connection from the Snowflake config file
        assert isinstance(connection_name, str)
        snowflake_dir = os.getenv('SNOWFLAKE_HOME', '~/.snowflake')
        snowflake_dir = os.path.expanduser(snowflake_dir)

        if not os.path.exists(snowflake_dir):
            # Alternative paths if ~/.snowflake directory does not exist on the machine:
            # https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-connect#connecting-using-the-connections-toml-file
            if sys.platform.startswith('win32'):
                userprofile = os.environ['USERPROFILE']
                snowflake_dir = os.path.join(userprofile, r'AppData\Local\snowflake')
            elif sys.platform.startswith('linux'):
                config_dir = os.getenv('XDG_CONFIG_HOME', '~/.config')
                snowflake_dir = os.path.join(config_dir, 'snowflake')
            elif sys.platform.startswith('darwin'):
                snowflake_dir = '~/Library/Application Support/snowflake'
            else:
                raise TemporaClientError(f'Unknown OS type: {sys.platform}')
            snowflake_dir = os.path.expanduser(snowflake_dir)

        connections_file = os.path.join(snowflake_dir, 'connections.toml')
        if os.path.exists(connections_file):
            with open(connections_file, 'rb') as f:
                connections_dict = tomli.load(f)
                if connection_name in connections_dict:
                    connection_params = connections_dict[connection_name]
                else:
                    raise TemporaClientError(
                        f'Missing connection profile "{connection_name}" in '
                        f'Snowflake config file: {connections_file}',
                    )
        else:
            raise TemporaClientError(
                f'Missing Snowflake config file: {connections_file}',
            )

    optional_args = ['warehouse', 'role']
    if kwargs.get('auth_type') == 'password':
        # Username+password auth doesn't require an auth token
        optional_args.append('auth_token')
    elif kwargs.get('auth_type') == 'pat':
        # PAT doesn't require a password
        optional_args.append('password')

    # Process the arguments: try to read the environment var if passed arg is None
    for arg_name, arg_value in kwargs.items():
        if arg_name != 'connection_name':
            if not arg_value and arg_name not in optional_args:
                env_var = f'SNOWFLAKE_{arg_name.upper()}'
                env_value = os.environ.get(env_var)
                if not env_value and not connection_params.get(arg_name):
                    msg = (
                        f'No value passed for "{arg_name}" argument and environment '
                        f'variable {env_var} is not set.'
                    )
                    log.error(msg)
                    raise TemporaClientError(msg)
                else:
                    arg_value = env_value
            connection_params[arg_name] = arg_value

    return connection_params


def disconnect_snowflake() -> None:
    """Close the currently open Snowflake connection."""
    return close_db_connection(dbms_name='Snowflake')
