# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""Data provider loader."""
from datetime import datetime
from functools import partial
from itertools import tee
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union

import pandas as pd
from tqdm.auto import tqdm

from ..._version import VERSION
from ...common.pkg_config import get_config
from ...common.utility import export, valid_pyname
from ...nbwidgets import QueryTime
from .. import drivers
from ..drivers.driver_base import DriverBase, DriverProps
from .param_extractor import extract_query_params
from .query_container import QueryContainer
from .query_defns import DataEnvironment
from .query_provider_connections_mixin import QueryProviderConnectionsMixin
from .query_provider_utils_mixin import QueryProviderUtilsMixin
from .query_source import QuerySource
from .query_store import QueryStore

__version__ = VERSION
__author__ = "Ian Hellen"

_HELP_FLAGS = ("help", "?")
_DEBUG_FLAGS = ("print", "debug_query", "print_query")
_COMPATIBLE_DRIVER_MAPPINGS = {
    "mssentinel": ["m365d"],
    "mde": ["m365d"],
    "mssentinel_new": ["mssentinel"],
    "kusto_new": ["kusto"],
}


# These are mixin classes that do not have an __init__ method
# pylint: disable=super-init-not-called
@export
class QueryProvider(QueryProviderConnectionsMixin, QueryProviderUtilsMixin):
    """
    Container for query store and query execution provider.

    Instances of this class hold the query set and execution
    methods for a specific data environment.

    """

    def __init__(  # noqa: MC0001
        self,
        data_environment: Union[str, DataEnvironment],
        driver: Optional[DriverBase] = None,
        query_paths: Optional[List[str]] = None,
        **kwargs,
    ):
        """
        Query provider interface to queries.

        Parameters
        ----------
        data_environment : Union[str, DataEnvironment]
            Name or Enum of environment for the QueryProvider
        driver : DriverBase, optional
            Override the builtin driver (query execution class)
            and use your own driver (must inherit from
            `DriverBase`)
        query_paths : List[str]
            Additional paths to look for query definitions.
        kwargs :
            Other arguments are passed to the data provider driver.

        Notes
        -----
        Additional keyword arguments are passed to the data provider driver.
        The driver may support additional keyword arguments, use
        the QueryProvider.driver_help() method to get a list of these
        parameters.

        See Also
        --------
        DataProviderBase : base class for data query providers.

        """
        # import at runtime to prevent circular import
        # pylint: disable=import-outside-toplevel, cyclic-import
        from ...init.pivot_init.pivot_data_queries import add_data_queries_to_entities

        # pylint: enable=import-outside-toplevel
        setattr(self.__class__, "_add_pivots", add_data_queries_to_entities)

        data_environment, self.environment_name = self._check_environment(
            data_environment
        )

        self._driver_kwargs = kwargs.copy()
        if driver is None:
            self.driver_class = drivers.import_driver(data_environment)
            if issubclass(self.driver_class, DriverBase):
                driver = self.driver_class(data_environment=data_environment, **kwargs)
            else:
                raise LookupError(
                    "Could not find suitable data provider for", f" {data_environment}"
                )
        else:
            self.driver_class = driver.__class__
        # allow the driver to override the data environment used for selecting queries
        self.environment_name = (
            driver.get_driver_property(DriverProps.EFFECTIVE_ENV)
            or self.environment_name
        )

        self._additional_connections: Dict[str, DriverBase] = {}
        self._query_provider = driver
        # replace the connect method docstring with that from
        # the driver's connect method
        self.__class__.connect.__doc__ = driver.connect.__doc__
        self.all_queries = QueryContainer()

        # Add any query files
        data_env_queries: Dict[str, QueryStore] = {}
        if driver.use_query_paths:
            data_env_queries.update(
                self._read_queries_from_paths(query_paths=query_paths)
            )
        self.query_store = data_env_queries.get(
            self.environment_name, QueryStore(self.environment_name)
        )
        self._add_query_functions()
        self._query_time = QueryTime(units="day")

    def _check_environment(
        self, data_environment
    ) -> Tuple[Union[str, DataEnvironment], str]:
        """Check environment against known names."""
        if isinstance(data_environment, str):
            data_env = DataEnvironment.parse(data_environment)
            if data_env != DataEnvironment.Unknown:
                data_environment = data_env
                environment_name = data_environment.name
            elif data_environment in drivers.CUSTOM_PROVIDERS:
                environment_name = data_environment
            else:
                raise TypeError(f"Unknown data environment {data_environment}")
        elif isinstance(data_environment, DataEnvironment):
            environment_name = data_environment.name
        else:
            raise TypeError(
                f"Unknown data environment type {data_environment} ({type(data_environment)})"
            )
        return data_environment, environment_name

    def __getattr__(self, name):
        """Return the value of the named property 'name'."""
        if "." in name:
            parent_name, child_name = name.split(".", maxsplit=1)
            parent = getattr(self, parent_name, None)
            if parent:
                return getattr(parent, child_name)
        raise AttributeError(f"{name} is not a valid attribute.")

    @property
    def environment(self) -> str:
        """Return the environment name."""
        return self.environment_name

    def connect(self, connection_str: Optional[str] = None, **kwargs):
        """
        Connect to data source.

        Parameters
        ----------
        connection_str : str
            Connection string for the data source

        """
        self._query_provider.connect(connection_str=connection_str, **kwargs)

        # If the driver has any attributes to expose via the provider
        # add those here.
        for attr_name, attr in self._query_provider.public_attribs.items():
            setattr(self, attr_name, attr)

        refresh_query_funcs = False
        # if the driver supports dynamic filtering of queries,
        # filter the query store based on connect-time parameters
        if self._query_provider.filter_queries_on_connect:
            self.query_store.apply_query_filter(self._query_provider.query_usable)
            refresh_query_funcs = True
        # Add any built-in or dynamically retrieved queries from driver
        if self._query_provider.has_driver_queries:
            driver_queries = self._query_provider.driver_queries
            self._add_driver_queries(queries=driver_queries)
            refresh_query_funcs = True

        if refresh_query_funcs:
            self._add_query_functions()

        # Since we're now connected, add Pivot functions
        self._add_pivots(lambda: self._query_time.timespan)

    def exec_query(self, query: str, **kwargs) -> Union[pd.DataFrame, Any]:
        """
        Execute simple query string.

        Parameters
        ----------
        query : str
            [description]
        use_connections : Union[str, List[str]]

        Other Parameters
        ----------------
        query_options : Dict[str, Any]
            Additional options passed to query driver.
        kwargs : Dict[str, Any]
            Additional options passed to query driver.

        Returns
        -------
        Union[pd.DataFrame, Any]
            Query results - a DataFrame if successful
            or a KqlResult if unsuccessful.

        """
        query_options = kwargs.pop("query_options", {}) or kwargs
        query_source = kwargs.pop("query_source", None)
        result = self._query_provider.query(
            query, query_source=query_source, **query_options
        )
        if not self._additional_connections:
            return result
        return self._exec_additional_connections(query, result, **kwargs)

    @property
    def query_time(self):
        """Return the default QueryTime control for queries."""
        return self._query_time

    def _execute_query(self, *args, **kwargs) -> Union[pd.DataFrame, Any]:
        if not self._query_provider.loaded:
            raise ValueError("Provider is not loaded.")
        if (
            not self._query_provider.connected
            and not _help_flag(*args)
            and not _debug_flag(*args, **kwargs)
        ):
            raise ValueError(
                "No connection to a data source.",
                "Please call connect(connection_str) and retry.",
            )
        query_name = kwargs.pop("query_name")
        family = kwargs.pop("query_path")

        query_source = self.query_store.get_query(
            query_path=family, query_name=query_name
        )
        if _help_flag(*args):
            query_source.help()
            return None

        params, missing = extract_query_params(query_source, *args, **kwargs)
        query_options = {
            "default_time_params": self._check_for_time_params(params, missing)
        }
        if missing:
            query_source.help()
            raise ValueError(f"No values found for these parameters: {missing}")

        split_by = kwargs.pop("split_query_by", None)
        if split_by:
            split_result = self._exec_split_query(
                split_by=split_by,
                query_source=query_source,
                query_params=params,
                args=args,
                **kwargs,
            )
            if split_result is not None:
                return split_result
            # if split queries could not be created, fall back to default
        query_str = query_source.create_query(
            formatters=self._query_provider.formatters, **params
        )
        # This looks for any of the "print query" debug args in args or kwargs
        if _debug_flag(*args, **kwargs):
            return query_str

        # Handle any query options passed and run the query
        query_options.update(_get_query_options(params, kwargs))
        return self.exec_query(query_str, query_source=query_source, **query_options)

    def _check_for_time_params(self, params, missing) -> bool:
        """Fall back on builtin query time if no time parameters were supplied."""
        defaults_added = False
        if "start" in missing:
            missing.remove("start")
            params["start"] = self._query_time.start
            defaults_added = True
        if "end" in missing:
            missing.remove("end")
            params["end"] = self._query_time.end
            defaults_added = True
        return defaults_added

    def _get_query_folder_for_env(self, root_path: str, environment: str) -> List[Path]:
        """Return query folder for current environment."""
        data_envs = [environment]
        if environment.casefold() in _COMPATIBLE_DRIVER_MAPPINGS:
            data_envs += _COMPATIBLE_DRIVER_MAPPINGS[environment.casefold()]
        return [
            _resolve_package_path(root_path).joinpath(data_env.casefold())
            for data_env in data_envs
        ]

    def _read_queries_from_paths(self, query_paths) -> Dict[str, QueryStore]:
        """Fetch queries from YAML files in specified paths."""
        settings: Dict[str, Any] = get_config("QueryDefinitions", {})
        all_query_paths: List[Union[Path, str]] = []
        for def_qry_path in settings.get("Default"):  # type: ignore
            # only read queries from environment folder
            builtin_qry_paths = self._get_query_folder_for_env(
                def_qry_path, self.environment_name
            )
            all_query_paths.extend(
                str(qry_path) for qry_path in builtin_qry_paths if qry_path.is_dir()
            )

        if settings.get("Custom") is not None:
            for custom_path in settings.get("Custom"):  # type: ignore
                custom_qry_path = _resolve_path(custom_path)
                if custom_qry_path:
                    all_query_paths.append(custom_qry_path)
        if query_paths:
            for param_path in query_paths:
                param_qry_path = _resolve_path(param_path)
                if param_qry_path:
                    all_query_paths.append(param_qry_path)
        if all_query_paths:
            return QueryStore.import_files(
                source_path=all_query_paths,
                recursive=True,
                driver_query_filter=self._query_provider.query_attach_spec,
            )
        # if no queries - just return an empty store
        return {self.environment_name: QueryStore(self.environment_name)}

    def _add_query_functions(self):
        """Add queries to the module as callable methods."""
        for qual_query_name in self.list_queries():
            query_path = qual_query_name.split(".")
            query_name = query_path[-1]
            current_node = self
            for container_name in query_path[:-1]:
                container_name = valid_pyname(container_name)
                if hasattr(current_node, container_name):
                    current_node = getattr(current_node, container_name)
                else:
                    new_node = QueryContainer()
                    setattr(current_node, container_name, new_node)
                    current_node = new_node

            query_cont_name = ".".join(query_path[:-1])

            # Create the partial function
            query_func = partial(
                self._execute_query, query_path=query_cont_name, query_name=query_name
            )
            query_func.__doc__ = self.query_store.get_query(
                query_path=query_cont_name, query_name=query_name
            ).create_doc_string()

            query_name = valid_pyname(query_name)
            setattr(current_node, query_name, query_func)
            setattr(self.all_queries, query_name, query_func)

    def _add_driver_queries(self, queries: Iterable[Dict[str, str]]):
        """Add driver queries to the query store."""
        for query in queries:
            self.query_store.add_query(
                name=query["name"],
                query=query["query"],
                query_paths=query.get(
                    "query_paths", query.get("query_container", "default")
                ),
                description=query["description"],
            )
        # For now, just add all of the functions again (with any connect-time acquired
        # queries) - we could be more efficient than this but unless there are 1000s of
        # queries it should not be noticeable.
        self._add_query_functions()

    def _exec_split_query(
        self,
        split_by: str,
        query_source: QuerySource,
        query_params: Dict[str, Any],
        args,
        **kwargs,
    ) -> Union[pd.DataFrame, str, None]:
        start = query_params.pop("start", None)
        end = query_params.pop("end", None)
        if not (start or end):
            print(
                "Cannot split a query that does not have 'start' and 'end' parameters"
            )
            return None
        try:
            split_delta = pd.Timedelta(split_by)
        except ValueError:
            split_delta = pd.Timedelta("1D")

        ranges = _calc_split_ranges(start, end, split_delta)

        split_queries = [
            query_source.create_query(
                formatters=self._query_provider.formatters,
                start=q_start,
                end=q_end,
                **query_params,
            )
            for q_start, q_end in ranges
        ]
        # This looks for any of the "print query" debug args in args or kwargs
        if _debug_flag(*args, **kwargs):
            return "\n\n".join(split_queries)

        # Retrieve any query options passed (other than query params)
        # and send to query function.
        query_options = _get_query_options(query_params, kwargs)
        query_dfs = [
            self.exec_query(query_str, query_source=query_source, **query_options)
            for query_str in tqdm(split_queries, unit="sub-queries", desc="Running")
        ]

        return pd.concat(query_dfs)


def _calc_split_ranges(start: datetime, end: datetime, split_delta: pd.Timedelta):
    """Return a list of time ranges split by `split_delta`."""
    # Use pandas date_range and split the result into 2 iterables
    s_ranges, e_ranges = tee(pd.date_range(start, end, freq=split_delta))
    next(e_ranges, None)  # skip to the next item in the 2nd iterable
    # Zip them together to get a list of (start, end) tuples of ranges
    # Note: we subtract 1 nanosecond from the 'end' value of each range so
    # to avoid getting duplicated records at the boundaries of the ranges.
    # Some providers don't have nanosecond granularity so we might
    # get duplicates in these cases
    ranges = [
        (s_time, e_time - pd.Timedelta("1ns"))
        for s_time, e_time in zip(s_ranges, e_ranges)
    ]

    # Since the generated time ranges are based on deltas from 'start'
    # we need to adjust the end time on the final range.
    # If the difference between the calculated last range end and
    # the query 'end' that the user requested is small (< 10% of a delta),
    # we just replace the last "end" time with our query end time.
    if (ranges[-1][1] - end) < (split_delta / 10):
        ranges[-1] = ranges[-1][0], end
    else:
        # otherwise append a new range starting after the last range
        # in ranges and ending in 'end"
        # note - we need to add back our subtracted 1 nanosecond
        ranges.append((ranges[-1][0] + pd.Timedelta("1ns"), end))
    return ranges


def _resolve_package_path(config_path: str) -> Path:
    """Resolve path relative to current package."""
    return (
        Path(config_path)
        if Path(config_path).is_absolute()
        else Path(__file__).resolve().parent.parent.joinpath(config_path)
    )


def _resolve_path(config_path: str) -> Optional[str]:
    """Resolve path."""
    if not Path(config_path).is_absolute():
        config_path = str(Path(config_path).expanduser().resolve())
    if not Path(config_path).is_dir():
        print(f"Warning: Custom query definitions path {config_path} not found")
        return None
    return config_path


def _help_flag(*args) -> bool:
    """Return True if help parameter passed."""
    return any(help_flag for help_flag in _HELP_FLAGS if help_flag in args)


def _debug_flag(*args, **kwargs) -> bool:
    """Return True if debug/print args passed."""
    return any(db_arg for db_arg in _DEBUG_FLAGS if db_arg in args) or any(
        db_arg for db_arg in _DEBUG_FLAGS if kwargs.get(db_arg, False)
    )


def _get_query_options(
    params: Dict[str, Any], kwargs: Dict[str, Any]
) -> Dict[str, Any]:
    # sourcery skip: inline-immediately-returned-variable, use-or-for-fallback
    """Return any kwargs not already in params."""
    query_options = kwargs.pop("query_options", {})
    if not query_options:
        # Any kwargs left over we send to the query provider driver
        query_options = {key: val for key, val in kwargs.items() if key not in params}
    query_options["time_span"] = {
        "start": params.get("start"),
        "end": params.get("end"),
    }
    return query_options
