"""
Functions for handling connections to things -- databases, git, ssh etc.

Commentary: This has become a bit untidy with growth. Really needs to be
refactored into a decent class model. One day.

For the database connectors, the connector is characterised by two aspects:

1.  The _dialect_ of the database. This defines the drivers that can be used,
    the flavour of SQL etc. For example, MySQL and MariaDB both support the
    `mysql` dialect. Similarly, `redshift` and `redshift-severless` both support
    the `redshift` dialect.

2.  How to obtain credentials to authenticate to the database. All databases are
    required to support the basic user ID / password mode. If a password (via an
    SSM parameter or reference to a secret in AWS Secrets Manager) is provided
    in the connection spec, that is what is used. Some database engines support
    an alternate mechanism to get credentials at run-time. e.g. AWS Redshift,
    Redshift Serverless and RDS Aurora databases can generate dynamic
    credentials via IAM. Unfortunately, the method is different in each case.
    These mechanisms relate to the connector _type_, not the dialect. For
    example, even though Redshift (provisioned) and Redshift Serverless support
    the same dialect, the have different, incompatible IAM based authentication
    mechanisms. The @db_credentials_getter() decorator, registers credential
    generators by connector type, where supported. The generated credentials
    are added into the connection spec passed to the dialect handler, which
    has no knowledge of how they were generated.

The basic process to find a Python driver is:

1.  Determine the dialect from the connection type (e.g. conn type
    `mysql-aurora` maps to a dialect of `mysql`.
2.  If the conn spec doesn't specify a subtype, use the default subtype
    for that dialect.
3.  Use the connector function that handles the (dialect, subtype)
    combination.

A note about database connection application names ...

This is a bit of a kludge that breaks the encapsulation of the connector
subsystem. For application name setting to be useful in practical terms, it
needs to include info about the job (e.g. the job_id, realm etc). This means
the connector now needs to have some knowledge of the job, which it didn't
previously. This then drives a bunch of fiddly changes at various different
points in lava. I hate it. But it's a useful feature.
"""

# ..............................................................................
# region Imports and constants
# ..............................................................................

from __future__ import annotations

import logging
import os
import re
import warnings
from collections.abc import Iterable
from contextlib import suppress
from typing import Any

import boto3
from jinja2 import Template

from lava.common import get_lava_param
from lava.config import LOGNAME, config
from lava.lavacore import IGNORE_FIELDS, LavaError
from lava.lib.aws import secman_get_json_secret
from lava.lib.misc import dict_check
from lava.lib.smb import LavaSMBConnection

__author__ = 'Murray Andrews, Alex Boul'

# pymysql has this filthy habit of using the warnings module for things that
# ought to be exceptions. We need to override that before importing it.
warnings.filterwarnings(action='error', module='pymysql.cursors')

LOG = logging.getLogger(name=LOGNAME)

CONNECTION_REQUIRED_FIELDS = {'conn_id', 'type', 'enabled'}
CONNECTION_OPTIONAL_FIELDS = {'description', 'owner', 'secret_id'}

# Note that `user` and `password` are not here because they can be generated by
# getting stuff from Secrets Manager, IAM etc.
SQL_CONNECTION_REQUIRED_FIELDS = CONNECTION_REQUIRED_FIELDS | {'host', 'port'}

# Map from the fields used for DBs by Secrets Manager to lava nomenclature.
SECMAN_DBLOGIN_MAP = {
    'username': 'user',
    'password': 'password',
    'host': 'host',
    'port': 'port',
    'dbname': 'database',
    'dbClusterIdentifier': 'description',
    'serviceName': 'service_name',
    'sid': 'sid',
}

SMB_CONNECTION_REQUIRED_FIELDS = CONNECTION_REQUIRED_FIELDS | {
    'remote_name',
    'host',
    'user',
    'password',
}
SMB_CONNECTION_OPTIONAL_FIELDS = CONNECTION_OPTIONAL_FIELDS | {
    'subtype',
    'domain',
    'my_name',
    'port',
    'use_ntlm_v2',
    'is_direct_tcp',
    'encrypt',
}

# ..............................................................................
# endregion Imports and constants
# ..............................................................................


# ------------------------------------------------------------------------------
def get_connection_spec(
    conn_id: str, realm: str, aws_session: boto3.Session = None
) -> dict[str, Any]:
    """
    Read a connection specification from DynamoDB and do some basic checking.

    :param conn_id:     Connection ID.
    :param realm:       Realm
    :param aws_session: A boto3 Session().

    :return:            The connection spec.

    """

    if not aws_session:
        aws_session = boto3.Session()

    conn_table_name = 'lava.' + realm + '.connections'
    try:
        conn_table = aws_session.resource('dynamodb').Table(conn_table_name)
    except Exception as e:
        raise LavaError(f'Connection {conn_id}: Cannot get DynamoDB table {conn_table_name} - {e}')

    # ----------------------------------------
    # Retrieve connection spec from DynamoDB
    try:
        conn_spec = conn_table.get_item(Key={'conn_id': conn_id})['Item']  # type:dict
    except KeyError:
        raise Exception(f'{conn_id}: No such connection')

    # Set defaults for some optional elements
    conn_spec.setdefault('enabled', False)

    # ----------------------------------------
    # Make sure the spec is ok
    try:
        dict_check(conn_spec, required=CONNECTION_REQUIRED_FIELDS, ignore=IGNORE_FIELDS)
    except ValueError as e:
        raise Exception(f'{conn_id}: Bad connection record: {e}')

    LOG.debug(f'Connection record: {conn_spec}')

    return conn_spec


# ..............................................................................
# region SQL connectors
# ..............................................................................

# Map (dialect, subtype) --> connection handler
_PYSQL_CONN_HANDLERS = {}

# Map SQL dialects to conn types. Dialects must be SQLAlchemy compatible.
# The dialect + subtype is used to select an appropriate handler for the database.
# For example, connection types 'postgres' and 'postgres-aurora` both provide
# the 'postgresql' dialect. The subtype will tell us which default driver to use
# if not overridden in the connection spec. e.g. The 'postgresql' dialect
# databases default to using pg8000. Authentication is a separate matter and
# is linked to the DB engine not the dialect (e.g. postgres is not identical to
# postgres-aurora for authentication purposes).
_SQL_DIALECT_TO_CONN_TYPE: dict[str, Iterable[str]] = {
    'postgresql': {'psql', 'postgresql', 'postgres', 'postgres-aurora', 'postgres-rds'},
    'redshift': {'redshift', 'redshift-serverless'},
    'oracle': {'oracle', 'oracle-rds'},
    'mysql': {'mysql', 'mariadb', 'mysql-aurora', 'mysql-rds', 'mariadb-rds'},
    'mssql': {'mssql'},
    'sqlite3': {'sqlite3'},
}

# Invert the map to conn_type --> dialect
_SQL_CONN_TYPE_TO_DIALECT: dict[str, str] = {
    conn_type: dialect
    for dialect in _SQL_DIALECT_TO_CONN_TYPE
    for conn_type in _SQL_DIALECT_TO_CONN_TYPE[dialect]
}

# These are the default sql subtypes if not specified in the connector.
# Keys are dialects, values are the name of the module that handles it.
_SQL_DIALECT_DEFAULT_SUBTYPE = {
    'postgresql': 'pg8000',
    'redshift': 'pg8000',
    'redshift-serverless': 'pg8000',
    'mssql': 'pyodbc',
    'mysql': 'pymysql',
    'oracle': 'cx_oracle',
    'sqlite3': 'sqlite3',
}


# ------------------------------------------------------------------------------
def pysql_connector(dialect: str, subtype: str):
    """
    Register connection handlers for Python based connectors.

    Usage:

    ```python
    @pysql_connector(dialect, subtype)
    a_func(conn_spec)
    ```

    :param dialect:     SQL dialect. e.g. postgresql
    :param subtype:     Subtype indicating the handler (e.g pg8000). This is
                        not necessarily the module name.

    """

    def decorate(func):
        """
        Register the handler function.

        :param func:    Function to 2yyregister.
        :return:        Unmodified function.

        """

        if dialect not in _SQL_DIALECT_TO_CONN_TYPE:
            raise ValueError(f'pysql_connector: Unknown dialect: {dialect}')

        if not subtype:
            raise ValueError('pysql_connector: subtype must be specified')

        key = (dialect, subtype)

        if key in _PYSQL_CONN_HANDLERS:
            raise Exception(f'{key} is already registered')
        _PYSQL_CONN_HANDLERS[key] = func
        return func

    return decorate


# ------------------------------------------------------------------------------
_DB_CREDENTIALS_GETTERS = {}


def db_credentials_getter(*args: str):
    """
    Get credentials for a database.

    Typically the getter uses an IAM based mechanism to get a temporary
    password.

    The decorated function must return a tuple (username, password).

    Usage:

    ```python
    @db_password_getter(db_type, ...)
    a_func(conn_spec, aws_session=...)
    ```

    :param args:        Data base types (matches lava connection type).
    """

    def decorate(func):
        """
        Register the passord getter function.

        :param func:    Function to 2yyregister.
        :return:        Unmodified function.
        """

        for db_type in args:
            if db_type in _DB_CREDENTIALS_GETTERS:
                raise Exception(f'Password getter {db_type} already registered')
            _DB_CREDENTIALS_GETTERS[db_type] = func
        return func

    return decorate


# ------------------------------------------------------------------------------
def expand_sql_conn_spec(conn_spec: dict[str, Any], aws_session: boto3.Session = None) -> None:
    """
    Expand and validate an SQL DB connection spec by adding stuff from SSM / Secrets Manager.

    :param conn_spec:       A connection spec. This is modified in place.
    :param aws_session:     A boto3 Session().

    """

    # ----------------------------------------
    # Get the password from SSM.
    if 'password' in conn_spec:
        conn_spec['password'] = get_lava_param(conn_spec['password'], aws_session=aws_session)

    # ----------------------------------------
    # Handle secrets manager by merging values not specified in the conn spec.
    if 'secret_id' in conn_spec:
        secret_data = secman_get_json_secret(conn_spec['secret_id'], aws_session=aws_session)
        # Select the bits of the secret we want and change names to lava style
        for secman_key, lava_key in SECMAN_DBLOGIN_MAP.items():
            if lava_key not in conn_spec and secman_key in secret_data:
                conn_spec[lava_key] = secret_data[secman_key]

    # ----------------------------------------
    # Validate the conn_spec
    dict_check(conn_spec, required=SQL_CONNECTION_REQUIRED_FIELDS, ignore=IGNORE_FIELDS)

    if not conn_spec['enabled']:
        raise LavaError('Connection not enabled')

    with suppress(KeyError):
        conn_spec['port'] = int(conn_spec['port'])

    # ----------------------------------------
    # Work out dialect and subtype
    conn_type = conn_spec['type'].lower()
    try:
        dialect = conn_spec['dialect'] = _SQL_CONN_TYPE_TO_DIALECT[conn_type]
    except KeyError:
        raise LavaError(f'Unsupported connection type - {conn_type}')

    try:
        conn_spec.setdefault('subtype', _SQL_DIALECT_DEFAULT_SUBTYPE[dialect])
    except KeyError:
        LOG.critical(
            f'No default subtype for SQL dialect {dialect}', extra={'event_type': 'worker'}
        )
        raise LavaError(f'No default subtype for SQL dialect {dialect}')

    # ----------------------------------------
    # If password not specified, see if we have a platform password generator
    if 'password' not in conn_spec:
        if conn_type not in _DB_CREDENTIALS_GETTERS:
            raise LavaError('Password must be specified')

        conn_spec['user'], conn_spec['password'] = _DB_CREDENTIALS_GETTERS[conn_type](
            conn_spec, aws_session=aws_session
        )

    # If we don't have a user and password now, we have a problem.
    dict_check(conn_spec, required=('user', 'password'))


# ------------------------------------------------------------------------------
def clean_application_name(application_name: str) -> str:
    """Clean unwanted stuff from an application name (database connection)."""

    return re.sub(r"""['"\\]""", '', application_name)


# ------------------------------------------------------------------------------
def make_application_name(conn_id: str, realm: str = None, job_id: str = None) -> str | None:
    """
    Construct an application name for use with connections.

    Some connectors, including some of the database ones, support the concept of
    an application name or source identifier.

    All of the parameters must be non-empty, or discoverable from environment
    vars, in order to yield a value.

    :param conn_id:     Conn ID
    :param realm:       Realm name. If not specified, use the value of the
                        LAVA_REALM environment variable if available.
    :param job_id:      Job ID. If not specified, use the value of the
                        LAVA_JOB_ID environment variable if available.
    :return:            A cleansed application name.
    """

    if not realm:
        realm = os.environ.get('LAVA_REALM')
    if not job_id:
        job_id = os.environ.get('LAVA_JOB_ID')
    template = config('CONN_APP_NAME')

    if not all((realm, job_id, conn_id, template)):
        return None

    return clean_application_name(
        Template(template).render(realm=realm, job_id=job_id, conn_id=conn_id)
    )


# ------------------------------------------------------------------------------
def get_sql_conn_spec(
    conn_id: str, realm: str, aws_session: boto3.Session = None
) -> dict[str, Any]:
    """
    Get an sql connection spec from DynamoDB.

    The spec is validated and any other components (e.g. password from SSM or
    AWS Secrets Manager) are filled in.

    :param conn_id:         Connection ID
    :param realm:           Realm
    :param aws_session:     A boto3 Session().

    :return:                The inflated, validated conn spec.

    """

    conn_spec = get_connection_spec(conn_id, realm, aws_session=aws_session)
    expand_sql_conn_spec(conn_spec, aws_session=aws_session)

    return conn_spec


# ------------------------------------------------------------------------------
def _get_pysql_connection(
    conn_spec: dict[str, Any], autocommit: bool = False, application_name: str = None
):
    """
    Get a Python connection to the specified SQL database.

    The specifics depend on the underlying database type as specified on the
    connection info.

    We are assuming a DBAPI 2.0 interface.

    :param conn_spec:       A connection spec. This must be pre-validated.
    :param autocommit:      If True, attempt to enable autocommit. This is
                            database and driver dependent as not all DBs
                            support it (e.g. sqlite3) If False, autocommit is
                            not enabled (the default state for DBAPI 2.0).
    :param application_name: If possible, set the the application name when
                            connecting to the database. This is supported
                            for Postgres-like DBs but generally not for other
                            database types.
    :return:                A live DB connection

    """

    conn_id = conn_spec['conn_id']
    conn_type = conn_spec['type'].lower()
    key = (conn_spec['dialect'], conn_spec['subtype'])

    # ----------------------------------------
    # Get the handler for this dialect/subtype combo
    try:
        connect_handler = _PYSQL_CONN_HANDLERS[key]
    except KeyError:
        raise LavaError(f'Unsupported connection type - {conn_type} = {key}')

    # ----------------------------------------
    # Try to connect
    LOG.debug(f'Connecting to DB {conn_id} ({conn_type}) = {key}')
    conn = connect_handler(conn_spec, autocommit=autocommit, application_name=application_name)
    LOG.debug(
        f'Connected to DB {conn_id} with app name {application_name} (autocommit={autocommit})'
    )

    return conn


# ------------------------------------------------------------------------------
def get_pysql_connection(
    conn_id: str,
    realm: str,
    autocommit: bool = False,
    aws_session: boto3.Session = None,
    application_name: str = None,
):
    """
    Get a Python connection to the specified SQL database.

    The specifics depend on the underlying database type as specified on the
    connection info.

    We are assuming a DBAPI 2.0 interface.

    :param conn_id:         Connection ID
    :param realm:           Realm
    :param autocommit:      If True, attempt to enable autocommit. This is
                            database and driver dependent as not all DBs
                            support it (e.g. sqlite3) If False, autocommit is
                            not enabled (the default state for DBAPI 2.0).
    :param aws_session:     A boto3 Session().
    :param application_name: If possible, set the the application name when
                            connecting to the database. This is supported
                            for Postgres-like DBs, and in various ways, for
                            some others but not all database types. If not
                            specified, we try to construct a fallback value
                            from some lava specific environment variables.

    :return:                A live DB connection

    """

    conn_spec = get_sql_conn_spec(conn_id, realm, aws_session=aws_session)
    return _get_pysql_connection(
        conn_spec,
        autocommit=autocommit,
        application_name=application_name or make_application_name(conn_id=conn_id, realm=realm),
    )


# ..............................................................................
# endregion SQL connectors
# ..............................................................................

# ..............................................................................
# region SMB connectors
# ..............................................................................

# Map (conn_type, subtype) --> connection handler
_SMB_CONN_HANDLERS = {}

# Map SMB subtypes to conn types.
_SMB_CONN_SUBTYPES: list[str] = ['pysmb', 'smbprotocol']

# This is the default SMB subtype if not specified in the connector.
_SMB_CONN_DEFAULT_SUBTYPE = 'pysmb'


# ------------------------------------------------------------------------------
def smb_connector(subtype: str):
    """
    Register connection handlers for Python based connectors.

    Usage:

    ```python
    @smb_connector(subtype)
    a_func(conn_spec)
    ```

    :param subtype:     Subtype to use 'pysmb' or 'smbprotocol' drivers.

    """

    def decorate(func):
        """
        Register the handler function.

        :param func:    Function to 2yyregister.
        :return:        Unmodified function.

        """

        if subtype not in _SMB_CONN_SUBTYPES:
            raise ValueError(f'smb_connector: Unknown subtype: {subtype}')

        if not subtype:
            raise ValueError('smb_connector: subtype must be specified')

        if subtype in _SMB_CONN_HANDLERS:
            raise Exception(f'{subtype} is already registered')
        _SMB_CONN_HANDLERS[subtype] = func
        return func

    return decorate


# ------------------------------------------------------------------------------
def expand_smb_conn_spec(
    conn_spec: dict[str, Any], realm: str, aws_session: boto3.Session = None
) -> None:
    """
    Expand and validate an SMB connection spec by adding stuff from SSM / Secrets Manager.

    :param conn_spec:       A connection spec.
    :param realm:           A lava realm name.
    :param aws_session:     A boto3 Session().

    :return:                The inflated, validated conn spec.
    """

    # ----------------------------------------
    # Get the password from SSM.
    if 'password' in conn_spec:
        conn_spec['password'] = get_lava_param(conn_spec['password'], aws_session=aws_session)

    # ----------------------------------------
    # Handle secrets manager by merging values not specified in the conn spec.
    if 'secret_id' in conn_spec:
        secret_data = secman_get_json_secret(conn_spec['secret_id'], aws_session=aws_session)
        # Select the bits of the secret we want and change names to lava style
        for secman_key, lava_key in SECMAN_DBLOGIN_MAP.items():
            if lava_key not in conn_spec and secman_key in secret_data:
                conn_spec[lava_key] = secret_data[secman_key]

    # ----------------------------------------
    # Validate the conn_spec
    conn_id = conn_spec['conn_id']
    conn_type = conn_spec['type'].lower()
    try:
        dict_check(
            conn_spec,
            required=SMB_CONNECTION_REQUIRED_FIELDS,
            optional=SMB_CONNECTION_OPTIONAL_FIELDS,
            ignore=IGNORE_FIELDS,
        )
    except Exception as e:
        raise LavaError(f'Connection {conn_id}: {e}')

    if conn_type != 'smb':
        raise LavaError(f'Connection {conn_id}: Must be of type "smb" not "{conn_type}"')

    if not conn_spec['enabled']:
        raise LavaError('Connection not enabled')

    with suppress(KeyError):
        conn_spec['port'] = int(conn_spec['port'])
    conn_spec['my_name'] = conn_spec.get('my_name', f'lava-{realm}')

    # ----------------------------------------
    # Work out subtype
    try:
        conn_spec.setdefault('subtype', _SMB_CONN_DEFAULT_SUBTYPE)
    except KeyError:
        LOG.critical(
            'No default subtype for SMB connections',
            extra={'event_type': 'worker'},
        )
        raise LavaError('No default subtype for SMB connections')


# ------------------------------------------------------------------------------
def get_smb_conn_spec(
    conn_id: str, realm: str, aws_session: boto3.Session = None
) -> dict[str, Any]:
    """
    Get an SMB connection spec from DynamoDB.

    The spec is validated and any other components (e.g. password from SSM or
    AWS Secrets Manager) are filled in.

    :param conn_id:         Connection ID
    :param realm:           Realm
    :param aws_session:     A boto3 Session().

    :return:                The inflated, validated conn spec.

    """

    conn_spec = get_connection_spec(conn_id, realm, aws_session=aws_session)
    expand_smb_conn_spec(conn_spec, realm, aws_session=aws_session)

    return conn_spec


# ------------------------------------------------------------------------------
def _get_smb_connection(conn_spec: dict[str, Any]) -> LavaSMBConnection:
    """
    Get a Python connection to the specified SMB share.

    The subtype specifies the python package used (pysmb/smbprotocol).

    :param conn_spec:       A connection spec. This must be pre-validated.
    :return:                An active SMB connection.

    """

    conn_id = conn_spec['conn_id']
    conn_type = conn_spec['type'].lower()
    subtype = conn_spec['subtype']

    # ----------------------------------------
    # Get the handler subtype
    try:
        connect_handler = _SMB_CONN_HANDLERS[subtype]
    except KeyError:
        raise LavaError(f'Unsupported connection type - {conn_type} = {subtype}')

    # ----------------------------------------
    # Try to connect
    LOG.debug(f'Connecting to SMB server {conn_id} ({conn_type}) = {subtype}')
    conn = connect_handler(conn_spec)
    LOG.debug(f'Connected to SMB server {conn_id}')

    return conn


# ------------------------------------------------------------------------------
def get_smb_connection(
    conn_id: str, realm: str, aws_session: boto3.Session = None
) -> LavaSMBConnection:
    """
    Get a Python connection to the specified SMB fileshare.

    The subtype specifies the python package used (pysmb/smbprotocol).

    :param conn_id:         Connection ID
    :param realm:           Realm
    :param aws_session:     A boto3 Session().

    :return:                An active SMB connection.

    """

    conn_spec = get_smb_conn_spec(conn_id, realm, aws_session=aws_session)
    return _get_smb_connection(conn_spec)


# ..............................................................................
# end region SMB connectors
# ..............................................................................

# ..............................................................................
# region CLI connectors
# ..............................................................................

_CLI_CONN_HANDLERS = {}


# ------------------------------------------------------------------------------
def cli_connector(*args: str):
    """
    Register connection handlers for external CLI based connectors.

    Usage:

    ```python
    @cli_connector('conn_type1', ...)
    a_func(conn_spec)
    ```

    :param args:        A list of connection types that the decorated function
                        handles.
    """

    def decorate(func):
        """
        Register the handler function.

        :param func:    Function to register.
        :return:        Unmodified function.

        """
        for db_type in args:
            if db_type in _CLI_CONN_HANDLERS:
                raise Exception(f'{db_type} is already registered')
            _CLI_CONN_HANDLERS[db_type] = func
        return func

    return decorate


# ------------------------------------------------------------------------------
def get_cli_connection(
    conn_id: str, realm: str, workdir: str, aws_session: boto3.Session = None
) -> str:
    """
    Get a CLI connection to a resource.

    This will produce an command line executable that will establish a
    connection to the target resource, managing authentication as required. The
    usage of the executable is dependent on the nature of the resource. The
    caller needs to know how to call it.

    :param conn_id:         Connection ID
    :param realm:           Realm
    :param workdir:         Working directory. This is required in case the
                            connection handler needs to store stuff like
                            credential files.
    :param aws_session:     A boto3 Session().

    :return:                Name of an executable that implements the connection.

    """

    # ----------------------------------------
    # Get the connection spec and make sure its ok

    conn_spec = get_connection_spec(conn_id, realm, aws_session=aws_session)

    if not conn_spec['enabled']:
        raise LavaError(f'Connection {conn_id}: Not enabled')

    # ----------------------------------------
    # Use a type specific handler to connect

    conn_type = conn_spec['type'].lower()

    try:
        connect_handler = _CLI_CONN_HANDLERS[conn_type]
    except KeyError:
        raise LavaError(f'{conn_id}: Unsupported connection type - {conn_type}')

    LOG.debug(f'Getting CLI connector to {conn_id} ({conn_type})')
    try:
        conn = connect_handler(conn_spec, workdir, aws_session=aws_session)
    except Exception as e:
        raise LavaError(f'Connection {conn_id}: {e}')

    LOG.debug(f'CLI Connector to {conn_id} is {conn}')

    return conn


# ..............................................................................
# endregion CLI connectors
# ..............................................................................
