# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""Data provider sub-package."""
import importlib
from typing import Union

from ..query_defns import DataEnvironment

# flake8: noqa: F403
from .driver_base import DriverBase

from ..._version import VERSION

__version__ = VERSION

_ENVIRONMENT_DRIVERS = {
    DataEnvironment.LogAnalytics: ("kql_driver", "KqlDriver"),
    DataEnvironment.AzureSecurityCenter: ("kql_driver", "KqlDriver"),
    DataEnvironment.SecurityGraph: ("security_graph_driver", "SecurityGraphDriver"),
    DataEnvironment.Kusto: ("kusto_driver", "KustoDriver"),
    DataEnvironment.MDATP: ("mdatp_driver", "MDATPDriver"),
    DataEnvironment.MDE: ("mdatp_driver", "MDATPDriver"),
    DataEnvironment.LocalData: ("local_data_driver", "LocalDataDriver"),
    DataEnvironment.Splunk: ("splunk_driver", "SplunkDriver"),
    DataEnvironment.Mordor: ("mordor_driver", "MordorDriver"),
    DataEnvironment.Sumologic: ("sumologic_driver", "SumologicDriver"),
    DataEnvironment.ResourceGraph: ("resource_graph_driver", "ResourceGraphDriver"),
    DataEnvironment.M365D: ("mdatp_driver", "MDATPDriver"),
}


def import_driver(data_environment: DataEnvironment) -> type:
    """Import driver class for a data environment."""
    mod_name, cls_name = _ENVIRONMENT_DRIVERS.get(data_environment, (None, None))

    if not (mod_name and cls_name):
        raise ValueError(
            f"No driver available for environment {data_environment.name}.",
            "Possible values are:",
            ", ".join(env.name for env in _ENVIRONMENT_DRIVERS),
        )

    imp_module = importlib.import_module(
        f"msticpy.data.drivers.{mod_name}", package="msticpy"
    )
    return getattr(imp_module, cls_name)
