# Licensed to Elasticsearch B.V under one or more agreements.
# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
# See the LICENSE file in the project root for more information

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
from eland import SortOrder


if TYPE_CHECKING:
    import pandas as pd  # type: ignore


class PostProcessingAction(ABC):
    def __init__(self, action_type: str) -> None:
        """
        Abstract class for postprocessing actions

        Parameters
        ----------
            action_type: str
                The action type (e.g. sort_index, head etc.)
        """
        self._action_type = action_type

    @property
    def type(self) -> str:
        return self._action_type

    @abstractmethod
    def resolve_action(self, df: "pd.DataFrame") -> "pd.DataFrame":
        pass

    @abstractmethod
    def __repr__(self) -> str:
        pass


class SortIndexAction(PostProcessingAction):
    def __init__(self) -> None:
        super().__init__("sort_index")

    def resolve_action(self, df: "pd.DataFrame") -> "pd.DataFrame":
        return df.sort_index()

    def __repr__(self) -> str:
        return f"('{self.type}')"


class HeadAction(PostProcessingAction):
    def __init__(self, count: int) -> None:
        super().__init__("head")

        self._count = count

    def resolve_action(self, df: "pd.DataFrame") -> "pd.DataFrame":
        return df.head(self._count)

    def __repr__(self) -> str:
        return f"('{self.type}': ('count': {self._count}))"


class TailAction(PostProcessingAction):
    def __init__(self, count: int) -> None:
        super().__init__("tail")

        self._count = count

    def resolve_action(self, df: "pd.DataFrame") -> "pd.DataFrame":
        return df.tail(self._count)

    def __repr__(self) -> str:
        return f"('{self.type}': ('count': {self._count}))"


class SortFieldAction(PostProcessingAction):
    def __init__(self, sort_params_string: str) -> None:
        super().__init__("sort_field")

        if sort_params_string is None:
            raise ValueError("Expected valid string")

        # Split string
        sort_field, _, sort_order = sort_params_string.partition(":")
        if not sort_field or sort_order not in ("asc", "desc"):
            raise ValueError(
                f"Expected ES sort params string (e.g. _doc:desc). Got '{sort_params_string}'"
            )

        self._sort_field = sort_field
        self._sort_order = SortOrder.from_string(sort_order)

    def resolve_action(self, df: "pd.DataFrame") -> "pd.DataFrame":
        return df.sort_values(self._sort_field, self._sort_order == SortOrder.ASC)

    def __repr__(self) -> str:
        return f"('{self.type}': ('sort_field': '{self._sort_field}', 'sort_order': {self._sort_order}))"
