"""Project utilities."""
# pylint: disable=cyclic-import
import importlib
import os
import threading
import traceback
import warnings
from pathlib import Path
from typing import Any, Dict, List, Union

import importlib_resources
import pyspark.sql.functions as F
from pyspark.sql import DataFrame as SparkDataFrame
from pyspark.sql import Window

from kedro.framework.context import KedroContext
from kedro.framework.hooks import _create_hook_manager
from kedro.framework.project import IMPORT_ERROR_MESSAGE, _create_pipeline, settings
from kedro.framework.startup import _get_project_metadata
from kedro.pipeline import Pipeline, pipeline


def _create_kedro_context():
    """
    Initializes the Kedro context.
    """

    project_path = Path(Path.cwd()).resolve()
    metadata = _get_project_metadata(project_path)
    context_class = settings.CONTEXT_CLASS
    env = os.getenv("KEDRO_ENV")
    hook_manager = _create_hook_manager()

    config_loader_class = settings.CONFIG_LOADER_CLASS
    config_loader = config_loader_class(
        conf_source=str(metadata.project_path / settings.CONF_SOURCE),
        env=env,
        **settings.CONFIG_LOADER_ARGS,
    )
    context = context_class(
        package_name=metadata.package_name,
        project_path=metadata.project_path,
        config_loader=config_loader,
        env=env,
        extra_params={},
        hook_manager=hook_manager,
    )
    return context


def get_kedro_context() -> Union[KedroContext, None]:
    """
    Returns the current Kedro context and tries to create one if not possible.
    """
    context = CurrentKedroContext().context
    if context is None:
        try:
            return _create_kedro_context()
        except Exception:  # pylint: disable=broad-except
            return None

    return context


def get_credentials() -> Dict[str, Any]:
    """
    Returns the credentials from the Kedro context.
    """
    if get_kedro_context() is None:
        raise ValueError("Kedro context is not set.")
    # pylint: disable=protected-access
    return get_kedro_context()._get_config_credentials()  # noqa


def get_credential(key: str) -> Union[str, Dict[str, Any], None]:
    """
    Returns the value of a credential from the Kedro context.

    Args:
        key: The key of the credential.

    Returns:
        The value of the credential.
    """
    return get_credentials().get(key, None)


def get_catalog() -> Dict[str, Any]:
    """
    Returns the catalog from the Kedro context.
    """
    if get_kedro_context() is None:
        raise ValueError("Kedro context is not set.")
    # pylint: disable=protected-access
    return get_kedro_context()._get_catalog()  # noqa


class SingletonMeta(type):
    """
    Implementation of a Singleton class using the metaclass method
    """

    _instances = {}
    _lock = threading.Lock()

    def __call__(cls, *args, **kwargs):
        """
        Possible changes to the value of the `__init__` argument do not affect
        the returned instance.
        """
        if cls not in cls._instances:
            with cls._lock:
                if cls not in cls._instances:
                    instance = super(SingletonMeta, cls).__call__(*args, **kwargs)
                    cls._instances[cls] = instance
        return cls._instances[cls]


class CurrentKedroContext(
    metaclass=SingletonMeta
):  # pylint: disable=too-few-public-methods
    """
    A singleton class that stores the current Kedro context.
    """

    def __init__(self, context: KedroContext = None):
        """
        Virtually private constructor.

        Args:
            context: The current Kedro context.
        """
        self._context: KedroContext = context

    @property
    def context(self) -> KedroContext:
        """Returns the current Kedro context."""
        return self._context

    @context.setter
    def context(self, context: KedroContext):
        """Sets the current Kedro context."""
        self._context = context


def find_nested_pipelines(
    base_module: str, package_name: str = "virgo"
) -> Dict[str, Pipeline]:
    """Automatically find modular pipelines having a ``create_pipeline``
    function. By default, projects created using Kedro 0.18.3 and higher
    call this function to autoregister pipelines upon creation/addition.

    Projects that require more fine-grained control can still define the
    pipeline registry without calling this function. Alternatively, they
    can modify the mapping generated by the ``find_pipelines`` function.

    For more information on the pipeline registry and autodiscovery, see
    https://kedro.readthedocs.io/en/latest/nodes_and_pipelines/pipeline_registry.html

    Returns:
        A generated mapping from pipeline names to ``Pipeline`` objects.

    Warns:
        UserWarning: When a module does not expose a ``create_pipeline``
            function, the ``create_pipeline`` function does not return a
            ``Pipeline`` object, or if the module import fails up front.
    """
    pipeline_obj = None
    pipelines_dict = {"_default_": pipeline_obj or pipeline([])}

    # Handle the case that a project doesn't have a pipelines' directory.
    pipelines_package = None
    try:
        pipelines_package = importlib_resources.files(
            f"{package_name}.pipelines.{base_module}"
        )
    except ModuleNotFoundError as exc:
        if str(exc) == f"No module named '{package_name}.pipelines.{base_module}'":
            return pipelines_dict

    for pipeline_dir in pipelines_package.iterdir():
        if not pipeline_dir.is_dir():
            continue

        pipeline_name = pipeline_dir.name
        if pipeline_name == "_pycache_":
            continue

        pipeline_module_name = f"{package_name}.pipelines.{base_module}.{pipeline_name}"
        try:
            pipeline_module = importlib.import_module(pipeline_module_name)
        except:  # pylint: disable=bare-except  # noqa: E722
            warnings.warn(
                IMPORT_ERROR_MESSAGE.format(
                    module=pipeline_module_name, tb_exc=traceback.format_exc()
                )
            )
            continue

        pipeline_obj = _create_pipeline(pipeline_module)
        if pipeline_obj is not None:
            pipelines_dict[pipeline_name] = pipeline_obj
    return pipelines_dict


def fill_data_frame_columns(
    df: SparkDataFrame,  # pylint: disable=invalid-name
    columns_to_fill: List[str],
    method: str = "ffill",
    partition_columns: List[str] = None,
    order_columns: List[str] = None,
    value: Any = None,
) -> SparkDataFrame:
    """Fills the given spark DataFrame columns.

    Fills the given spark DataFrame columns using the specified filling method.

    Args:
        df: The Spark DataFrame containing the columns to be filled.
        columns_to_fill: The list of the names of the columns to be filled.
        method: The filling method.
            It can be one of:
            - ffill: forward filling
            - bfill: backwards filling
            - any: fill with value
        partition_columns: Partition columns used as the base for the filling.
            It can be empty.
        order_columns: Columns used for the ordering when filling. They are only used
            ffill or bfill are used.
        value: Value used to fill the specified columns. It can be a dictionary to
            fill every column with a different value.
    Returns:
        A spark DataFrame with the filled columns.
    Raises:
        RuntimeError: Incorrect function call.
    """
    # pylint: disable=too-many-arguments

    partition_columns = [] if partition_columns is None else partition_columns
    order_columns = [] if order_columns is None else order_columns

    if df is None or columns_to_fill is None:
        raise RuntimeError("One of the required arguments doesn't exist.")

    df_columns = df.columns
    columns_to_retain = [c for c in df_columns if c not in columns_to_fill]
    if value is not None:
        if isinstance(value, dict):
            filled_columns = [
                F.when(F.col(f"`{c}`").isNull(), value[c]).otherwise(F.col(f"`{c}`"))
                for c in columns_to_fill
            ]
        else:
            filled_columns = [
                F.when(F.col(f"`{c}`").isNull(), value).otherwise(F.col(f"`{c}`"))
                for c in columns_to_fill
            ]
        remain_columns = [F.col(f"`{c}`") for c in columns_to_retain]
        return df.select(*remain_columns, *filled_columns).select(
            [f"`{c}`" for c in df_columns]
        )

    if method == "ffill":
        if not order_columns:
            raise RuntimeError("Ordering is required to by ffill.")
        if not partition_columns:
            partition_columns = F.lit(0)
        window_ffill = (
            Window.partitionBy(partition_columns)
            .orderBy(*[F.asc(f"`{c}`") for c in order_columns])
            .rowsBetween(Window.unboundedPreceding, 0)
        )
        filled_columns = [
            F.last(f"`{c}`", ignorenulls=True).over(window_ffill).alias(c)
            for c in columns_to_fill
        ]
        remain_columns = [F.col(f"`{c}`") for c in columns_to_retain]
        return df.select(*remain_columns, *filled_columns).select(
            [f"`{c}`" for c in df_columns]
        )
    if method == "bfill":
        if not order_columns:
            raise RuntimeError("Ordering is required to by bfill.")
        if not partition_columns:
            partition_columns = F.lit(0)
        window_bfill = (
            Window.partitionBy(partition_columns)
            .orderBy(*[F.asc(f"`{c}`") for c in order_columns])
            .rowsBetween(0, Window.unboundedFollowing)
        )
        filled_columns = [
            F.first(f"`{c}`", ignorenulls=True).over(window_bfill).alias(c)
            for c in columns_to_fill
        ]
        remain_columns = [F.col(f"`{c}`") for c in columns_to_retain]
        return df.select(*remain_columns, *filled_columns).select(
            [f"`{c}`" for c in df_columns]
        )
    raise RuntimeError("Unknown method type.")
