"""Strategy executors

`SparkExecutor` simply calls a partition processor where each partition will
then be processed with the `LocalExecutor`
"""

# mypy: disable-error-code="no-redef"

from __future__ import annotations

from collections import defaultdict
from collections.abc import Iterator
from functools import partial
from typing import TYPE_CHECKING
from typing import Protocol
from typing import Type
from typing import TypeAlias
from typing import TypeVar
from typing import cast
from typing import final

from networkx.utils.union_find import UnionFind
from pyspark.sql import Row
from pyspark.sql import SparkSession

from liken._constants import CANONICAL_ID
from liken._dataframe import Frame
from liken._dataframe import LocalDF
from liken._dataframe import SparkDF
from liken._strats_library import BaseStrategy
from liken._strats_library import BinaryDedupers
from liken._strats_manager import SEQUENTIAL_API_DEFAULT_KEY
from liken._strats_manager import Rules
from liken._strats_manager import StratsDict
from liken._types import Columns
from liken._types import Keep


if TYPE_CHECKING:
    from liken.dedupe import Dedupe


# TYPES:


SingleComponents: TypeAlias = dict[int, list[int]]
MultiComponents: TypeAlias = dict[tuple[int, ...], list[int]]
F = TypeVar("F", bound=Frame)


# EXECUTORS:


class Executor(Protocol[F]):
    def execute(
        self,
        df: F,
        /,
        *,
        columns: Columns | None,
        strats: StratsDict | Rules,
        keep: Keep,
        drop_duplicates: bool,
        drop_canonical_id: bool,
        id: str | None,
    ) -> F: ...


@final
class LocalExecutor(Executor):
    def execute(
        self,
        df: LocalDF,
        /,
        *,
        columns: Columns | None,
        strats: StratsDict | Rules,
        keep: Keep,
        drop_duplicates: bool,
        drop_canonical_id: bool,
        id: str | None = None,
    ) -> LocalDF:
        """Process a local dataframe according to the strategy collection

        Processing is defined according to whether the collections of
        strategies is:
            - Rules: in which case "and" combinations are allowed
            - StratsDict: in which case handles Sequential and Dict API

        For Rules, predication is implemented if an and combination contains at
        least one Binary Deduper. In that case the binary dedupers are proxy
        WHERE filters that propagate a set of dataframe indice positions to the
        next deduper (most likely a threshold deduper, but optionally binary
        also.)
        """

        del id  # Unused: here for interface symmetry with SparkExecutor

        call_strat = partial(
            self._call_strat,
            drop_duplicates=drop_duplicates,
            keep=keep,
        )

        if isinstance(strats, StratsDict):
            if not columns:
                for col, iter_strats in strats.items():
                    for strat in iter_strats:
                        uf, n = self._build_uf(strat, df, col)
                        components: SingleComponents = self._get_components(uf, n)
                        df = call_strat(strat, components)
            else:
                # For sequential API calls e.g.`.canonicalize("address")`
                for strat in strats[SEQUENTIAL_API_DEFAULT_KEY]:
                    uf, n = self._build_uf(strat, df, columns)
                    components: SingleComponents = self._get_components(uf, n)
                    df = call_strat(strat, components)

        if isinstance(strats, Rules):
            for stage in strats:
                has_any_binary: bool = stage.has_any_binary_strat

                # predication only if at least one binary strategy
                if has_any_binary:
                    indices = set()

                    for col, strat in stage.and_strats:
                        uf, n = self._build_uf(strat, df, col, predicate=indices)

                        components = defaultdict(list)
                        idx: list = sorted(indices)
                        for i in range(n):
                            if not indices:
                                components[uf[i]].append(i)
                            else:
                                components[idx[uf[i]]].append(idx[i])

                        if isinstance(strat, BinaryDedupers):
                            for c in components.values():
                                if len(c) > 1:
                                    indices = indices.union(set(c))

                else:
                    ufs = []

                    for col, strat in stage.and_strats:
                        uf, n = self._build_uf(strat, df, col)
                        ufs.append(uf)
                    components: MultiComponents = self._get_multi_components(ufs, n)

                df = call_strat(strat, components)

        if drop_canonical_id:
            return df.drop_col(CANONICAL_ID)
        return df

    @staticmethod
    def _build_uf(
        strat: BaseStrategy, df: LocalDF, columns: Columns, predicate: set = set()
    ) -> tuple[UnionFind[int], int]:
        return strat.set_frame(df).build_union_find(columns, predicate=predicate)

    @staticmethod
    def _get_components(
        uf: UnionFind[int],
        n: int,
    ) -> SingleComponents:
        components = defaultdict(list)
        for i in range(n):
            components[uf[i]].append(i)
        return components

    @staticmethod
    def _get_multi_components(
        ufs: list[UnionFind[int]],
        n: int,
    ) -> MultiComponents:
        components = defaultdict(list)
        for i in range(n):
            signature = tuple(uf[i] for uf in ufs)
            components[signature].append(i)
        return components

    @staticmethod
    def _call_strat(
        strat: BaseStrategy,
        components: SingleComponents | MultiComponents,
        drop_duplicates: bool,
        keep: Keep,
    ) -> LocalDF:
        return strat.canonicalizer(
            components=components,
            drop_duplicates=drop_duplicates,
            keep=keep,
        )


@final
class SparkExecutor(Executor):
    def __init__(self, spark_session: SparkSession, id: str | None = None):
        self._spark_session = spark_session
        self._id = id

    def execute(
        self,
        df: SparkDF,
        /,
        *,
        columns: Columns | None,
        strats: StratsDict | Rules,
        keep: Keep,
        drop_duplicates: bool,
        drop_canonical_id: bool,
        id: str | None = None,
    ) -> SparkDF:
        """Spark specific deduplication helper

        Maps dataframe partitions to be processed via the RDD API yielding low-
        level list[Rows], which are then post-processed back to a dataframe.

        Args:
            columns: The attribute to deduplicate.
            strats: the collection of strats
        Retuns:
            Instance's _df attribute is updated
        """

        # import in worker node
        from liken.dedupe import Dedupe

        # IMPORTANT: Use local variables, no references to Self
        process_partition = self._process_partition

        rdd = df.mapPartitions(
            lambda partition: process_partition(
                factory=Dedupe,
                partition=partition,
                strats=strats,
                id=id,
                columns=columns,
                drop_duplicates=drop_duplicates,
                keep=keep,
            )
        )

        schema = df._schema

        df = SparkDF(self._spark_session.createDataFrame(rdd, schema=schema), is_init=False)

        if drop_canonical_id:
            return df.drop_col(CANONICAL_ID)
        return df

    @staticmethod
    def _process_partition(
        *,
        factory: Type[Dedupe],
        partition: Iterator[Row],
        strats: StratsDict | Rules,
        id: str | None,
        columns: Columns | None,
        drop_duplicates: bool,
        keep: Keep = "first",
    ) -> Iterator[Row]:
        """process a spark dataframe partition i.e. a list[Row]

        This function is functionality mapped to a worker node. For clean
        separation from the driver, strats are re-instantiated and the main
        liken API is executed *per* worker node.

        Args:
            paritition_iter: a partition
            strats: the collection of strats
            id: the unique identified of the dataset a.k.a "business key"
            columns: the attribute on which to deduplicate

        Returns:
            A list[Row], deduplicated
        """
        # handle empty partitions
        rows: list[Row] = list(partition)
        if not rows:
            return iter([])

        # Core API reused per partition, per worker node
        lk = factory._from_rows(rows)  # type: ignore
        lk.apply(strats)  # type: ignore
        df = lk.canonicalize(
            columns,
            keep=keep,
            drop_duplicates=drop_duplicates,
            id=id,
        )

        return iter(cast(list[Row], df))
