#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8

"""
Create connectors with `meerschaum.connectors.get_connector()`.
For ease of use, you can also import from the root `meerschaum` module:
```
>>> from meerschaum import get_connector
>>> conn = get_connector()
```
"""

from __future__ import annotations

import meerschaum as mrsm
from meerschaum.utils.typing import Any, SuccessTuple, Union, Optional, List, Dict
from meerschaum.utils.threading import Lock, RLock
from meerschaum.utils.warnings import error, warn

from meerschaum.connectors.Connector import Connector, InvalidAttributesError
from meerschaum.connectors.sql.SQLConnector import SQLConnector
from meerschaum.connectors.api.APIConnector import APIConnector
from meerschaum.connectors.sql._create_engine import flavor_configs as sql_flavor_configs

__all__ = ("Connector", "SQLConnector", "APIConnector", "get_connector", "is_connected")

### store connectors partitioned by
### type, label for reuse
connectors: Dict[str, Dict[str, Connector]] = {
    'api'   : {},
    'sql'   : {},
    'plugin': {},
}
instance_types: List[str] = ['sql', 'api']
_locks: Dict[str, RLock] = {
    'connectors'               : RLock(),
    'types'                    : RLock(),
    'custom_types'             : RLock(),
    '_loaded_plugin_connectors': RLock(),
    'instance_types'           : RLock(),
}
attributes: Dict[str, Dict[str, Any]] = {
    'api': {
        'required': [
            'host',
            'username',
            'password'
        ],
        'default': {
            'protocol': 'http',
        },
    },
    'sql': {
        'flavors': sql_flavor_configs,
    },
}
### Fill this with objects only when connectors are first referenced.
types: Dict[str, Any] = {}
custom_types: set = set()
_loaded_plugin_connectors: bool = False


def get_connector(
        type: str = None,
        label: str = None,
        refresh: bool = False,
        debug: bool = False,
        **kw: Any
    ) -> Connector:
    """
    Return existing connector or create new connection and store for reuse.
    
    You can create new connectors if enough parameters are provided for the given type and flavor.
    

    Parameters
    ----------
    type: Optional[str], default None
        Connector type (sql, api, etc.).
        Defaults to the type of the configured `instance_connector`.

    label: Optional[str], default None
        Connector label (e.g. main). Defaults to `'main'`.

    refresh: bool, default False
        Refresh the Connector instance / construct new object. Defaults to `False`.

    kw: Any
        Other arguments to pass to the Connector constructor.
        If the Connector has already been constructed and new arguments are provided,
        `refresh` is set to `True` and the old Connector is replaced.

    Returns
    -------
    A new Meerschaum connector (e.g. `meerschaum.connectors.api.APIConnector`,
    `meerschaum.connectors.sql.SQLConnector`).
    
    Examples
    --------
    The following parameters would create a new
    `meerschaum.connectors.sql.SQLConnector` that isn't in the configuration file.

    ```
    >>> conn = get_connector(
    ...     type = 'sql',
    ...     label = 'newlabel',
    ...     flavor = 'sqlite',
    ...     database = '/file/path/to/database.db'
    ... )
    >>>
    ```

    """
    from meerschaum.connectors.parse import parse_instance_keys
    from meerschaum.config import get_config
    from meerschaum.config.static import STATIC_CONFIG
    from meerschaum.utils.warnings import warn
    global _loaded_plugin_connectors
    if isinstance(type, str) and not label and ':' in type:
        type, label = type.split(':', maxsplit=1)
    with _locks['_loaded_plugin_connectors']:
        if not _loaded_plugin_connectors:
            load_plugin_connectors()
            _loaded_plugin_connectors = True
    if type is None and label is None:
        default_instance_keys = get_config('meerschaum', 'instance', patch=True)
        ### recursive call to get_connector
        return parse_instance_keys(default_instance_keys)

    ### NOTE: the default instance connector may not be main.
    ### Only fall back to 'main' if the type is provided by the label is omitted.
    label = label if label is not None else STATIC_CONFIG['connectors']['default_label']

    ### type might actually be a label. Check if so and raise a warning.
    if type not in connectors:
        possibilities, poss_msg = [], ""
        for _type in get_config('meerschaum', 'connectors'):
            if type in get_config('meerschaum', 'connectors', _type):
                possibilities.append(f"{_type}:{type}")
        if len(possibilities) > 0:
            poss_msg = " Did you mean"
            for poss in possibilities[:-1]:
                poss_msg += f" '{poss}',"
            if poss_msg.endswith(','):
                poss_msg = poss_msg[:-1]
            if len(possibilities) > 1:
                poss_msg += " or"
            poss_msg += f" '{possibilities[-1]}'?"

        warn(f"Cannot create Connector of type '{type}'." + poss_msg, stack=False)
        return None

    if 'sql' not in types:
        from meerschaum.connectors.plugin import PluginConnector
        with _locks['types']:
            types.update({
                'api'   : APIConnector,
                'sql'   : SQLConnector,
                'plugin': PluginConnector,
            })
    
    ### determine if we need to call the constructor
    if not refresh:
        ### see if any user-supplied arguments differ from the existing instance
        if label in connectors[type]:
            warning_message = None
            for attribute, value in kw.items():
                if attribute not in connectors[type][label].meta:
                    import inspect
                    cls = connectors[type][label].__class__
                    cls_init_signature = inspect.signature(cls)
                    cls_init_params = cls_init_signature.parameters
                    if attribute not in cls_init_params:
                        warning_message = (
                            f"Received new attribute '{attribute}' not present in connector " +
                            f"{connectors[type][label]}.\n"
                        )
                elif connectors[type][label].__dict__[attribute] != value:
                    warning_message = (
                        f"Mismatched values for attribute '{attribute}' in connector "
                        + f"'{connectors[type][label]}'.\n" +
                        f"  - Keyword value: '{value}'\n" +
                        f"  - Existing value: '{connectors[type][label].__dict__[attribute]}'\n"
                    )
            if warning_message is not None:
                warning_message += (
                    "\nSetting `refresh` to True and recreating connector with type:"
                    + f" '{type}' and label '{label}'."
                )
                refresh = True
                warn(warning_message)
        else: ### connector doesn't yet exist
            refresh = True

    ### only create an object if refresh is True
    ### (can be manually specified, otherwise determined above)
    if refresh:
        with _locks['connectors']:
            try:
                ### will raise an error if configuration is incorrect / missing
                conn = types[type](label=label, **kw)
                connectors[type][label] = conn
            except InvalidAttributesError as ie:
                warn(
                    f"Incorrect attributes for connector '{type}:{label}'.\n"
                    + str(ie),
                    stack = False,
                )
                conn = None
            except Exception as e:
                from meerschaum.utils.formatting import get_console
                console = get_console()
                if console:
                    console.print_exception()
                warn(
                    f"Exception when creating connector '{type}:{label}'.\n" + str(e),
                    stack = False,
                )
                conn = None
        if conn is None:
            return None

    return connectors[type][label]


def is_connected(keys: str, **kw) -> bool:
    """
    Check if the connector keys correspond to an active connection.
    If the connector has not been created, it will immediately return `False`.
    If the connector exists but cannot communicate with the source, return `False`.
    
    **NOTE:** Only works with instance connectors (`SQLConnectors` and `APIConnectors`).
    Keyword arguments are passed to `meerschaum.connectors.poll.retry_connect`.

    Parameters
    ----------
    keys:
        The keys to the connector (e.g. `'sql:main'`).
        
    Returns
    -------
    A `bool` corresponding to whether a successful connection may be made.

    """
    import warnings
    if ':' not in keys:
        warn(f"Invalid connector keys '{keys}'")

    try:
        typ, label = keys.split(':')
    except Exception as e:
        return False
    if typ not in instance_types:
        return False
    if not (label in connectors.get(typ, {})):
        return False

    from meerschaum.connectors.parse import parse_instance_keys
    conn = parse_instance_keys(keys)
    try:
        with warnings.catch_warnings():
            warnings.filterwarnings('ignore')
            return conn.test_connection(**kw)
    except Exception as e:
        return False


def make_connector(
        cls,
    ):
    """
    Register a class as a `Connector`.
    The `type` will be the lower case of the class name, without the suffix `connector`.

    Parameters
    ----------
    instance: bool, default False
        If `True`, make this connector type an instance connector.
        This requires implementing the various pipes functions and lots of testing.

    Examples
    --------
    >>> import meerschaum as mrsm
    >>> from meerschaum.connectors import make_connector, Connector
    >>> class FooConnector(Connector):
    ...     def __init__(self, label: str, **kw):
    ...         super().__init__('foo', label, **kw)
    ... 
    >>> make_connector(FooConnector)
    >>> mrsm.get_connector('foo', 'bar')
    foo:bar
    >>> 
    """
    import re
    typ = re.sub(r'connector$', '', cls.__name__.lower())
    with _locks['types']:
        types[typ] = cls
    with _locks['custom_types']:
        custom_types.add(typ)
    with _locks['connectors']:
        if typ not in connectors:
            connectors[typ] = {}
    if getattr(cls, 'IS_INSTANCE', False):
        with _locks['instance_types']:
            if typ not in instance_types:
                instance_types.append(typ)

    return cls


def load_plugin_connectors():
    """
    If a plugin makes use of the `make_connector` decorator,
    load its module.
    """
    from meerschaum.plugins import get_plugins, import_plugins
    to_import = []
    for plugin in get_plugins():
        with open(plugin.__file__, encoding='utf-8') as f:
            text = f.read()
        if 'make_connector' in text:
            to_import.append(plugin.name)
    if not to_import:
        return
    import_plugins(*to_import) 


def get_connector_plugin(
        connector: Connector,
    ) -> Union[str, None, mrsm.Plugin]:
    """
    Determine the plugin for a connector.
    This is useful for handling virtual environments for custom instance connectors.

    Parameters
    ----------
    connector: Connector
        The connector which may require a virtual environment.

    Returns
    -------
    A Plugin, 'mrsm', or None.
    """
    if not hasattr(connector, 'type'):
        return None
    plugin_name = (
        connector.__module__.replace('plugins.', '').split('.')[0]
        if connector.type in custom_types else (
            connector.label
            if connector.type == 'plugin'
            else 'mrsm'
        )
    )
    plugin = mrsm.Plugin(plugin_name)
    return plugin if plugin.is_installed() else None
