#  IBM Confidential
#  PID 5900-BAF
#  Copyright StreamSets Inc., an IBM Company 2025

"""Module containing Flow Model."""

import itertools
import requests
from abc import ABC, abstractmethod
from collections.abc import Callable
from ibm_watsonx_data_integration.common.models import BaseModel, CollectionModel
from ibm_watsonx_data_integration.common.utils import SeekableList
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeVar, cast, overload
from typing_extensions import override

if TYPE_CHECKING:
    from ibm_watsonx_data_integration.services.datastage.models import (
        BatchFlow,
    )
    from ibm_watsonx_data_integration.services.streamsets.models import (
        Environment,
        StreamingFlow,
    )
    from ibm_watsonx_data_integration.cpd_models.project_model import Project

FlowT = TypeVar("FlowT", bound="Flow")


class Flow(BaseModel, ABC):
    """Represents a generic Flow object."""

    _flow_registry: ClassVar[dict[str, type["Flow"]]] = {}

    @classmethod
    def register(cls, flow_type: str) -> Callable[[type[FlowT]], type[FlowT]]:
        """Class decorator to register a new Flow subclass.

        Args:
            flow_type: The key under which the subclass will be stored in `flow_registry`.

        Returns:
            A decorator that registers the subclass and returns it unchanged.
        """

        def inner(subclass: type[FlowT]) -> type[FlowT]:
            cls._flow_registry[flow_type] = subclass
            return subclass

        return inner

    @abstractmethod
    def _delete(self) -> requests.Response: ...

    @staticmethod
    @abstractmethod
    def _create(
        project: "Project", name: str, environment: "Environment | None", description: str, flow_type: str
    ) -> "Flow": ...

    @abstractmethod
    def _update(self) -> requests.Response: ...

    @abstractmethod
    def _duplicate(self, name: str, description: str | None) -> "Flow": ...


class PayloadExtender(ABC):
    """Interface class for flows with custom payload logic.

    This interface should be implemented by any flow that requires custom logic for generating a job's payload.
    It ensures a consistent contract for flows that need to override the default payload creation behavior.
    """

    def extend(self, payload: dict[str, Any], flow: Flow) -> dict[str, Any]:
        """Entry point for starting pre-check and extending create job payload."""
        self.run_pre_start_operations(flow)
        return self.prepare_payload(payload, flow)

    @abstractmethod
    def run_pre_start_operations(self, flow: Flow) -> None:
        """Here we should perform all flow prerequisite operations like validation, compilation etc.

        When pre-checks fail we should raise ``FlowPreStartOperationError`` exception
        to carry out this information to user.

        Raises:
            FlowPreStartOperationError: If any pre-check operation failed.
        """
        raise NotImplementedError

    @abstractmethod
    def prepare_payload(self, payload: dict[str, Any], flow: Flow) -> dict[str, Any]:
        """Here we should modify and return payload for job creation."""
        raise NotImplementedError


class DefaultFlowPayloadExtender(PayloadExtender):
    """Default payload extender which setup only `asset_ref`."""

    @override
    def run_pre_start_operations(self, flow: Flow) -> None:
        pass

    @override
    def prepare_payload(self, payload: dict[str, Any], flow: Flow) -> dict[str, Any]:
        payload["asset_ref"] = flow.flow_id
        return payload


class Flows(CollectionModel[Flow]):
    """Collection of BatchFlow and StreamingFlow objects."""

    def __init__(self, project: "Project") -> None:
        """The __init__ of the Flows class.

        Args:
            project: The Project object.
        """
        from ibm_watsonx_data_integration.services.datastage.models import BatchFlows
        from ibm_watsonx_data_integration.services.streamsets.models import StreamingFlows

        super().__init__(platform=project._platform)  # noqa
        self._streaming_flows = StreamingFlows(project)
        self._batch_flows = BatchFlows(project)
        self._project = project

    def _paginate(self, **kwargs: dict) -> Flow:
        for item in itertools.chain(self._streaming_flows.get_all(), self._batch_flows.get_all()):
            yield item

    def __len__(self) -> int:
        """Total amount of flows."""
        return len(self._streaming_flows) + len(self._batch_flows)

    @overload
    def get_all(self, flow_type: Literal["streaming"], **kwargs) -> "SeekableList[StreamingFlow]": ...  # noqa: ANN003

    @overload
    def get_all(self, flow_type: Literal["batch"], **kwargs) -> "SeekableList[BatchFlow]": ...  # noqa: ANN003

    @overload
    def get_all(self, flow_type: None = None, **kwargs) -> SeekableList[Flow]: ...  # noqa: ANN003

    def get_all(
        self,
        flow_type: str | None = None,
        **kwargs,  # noqa: ANN003
    ) -> "SeekableList[StreamingFlow] | SeekableList[BatchFlow] | SeekableList[Flow]":
        """Used to get multiple (all) results from flows api.

        Args:
            flow_type: The type of flow to be returned
            **kwargs: Optional other arguments to be passed to filter the results.

        Returns:
            A :py:obj:`list` of inherited instances of
                :py:class:`streamsets.sdk.sch_models.BaseModel`.

        Raises:
            TypeError: If flow_type is not 'streaming', 'batch', or None.
        """
        if flow_type is not None and flow_type not in ["streaming", "batch"]:
            raise TypeError(f"flow_type must be 'streaming', 'batch', got '{flow_type}'")

        if flow_type == "streaming":
            return cast(SeekableList[Flow], self._streaming_flows.get_all(**kwargs))
        elif flow_type == "batch":
            return cast(SeekableList[Flow], self._batch_flows.get_all(**kwargs))
        else:
            return SeekableList(
                itertools.chain(self._streaming_flows.get_all(**kwargs), self._batch_flows.get_all(**kwargs))
            )

    @overload
    def get(self, flow_type: Literal["streaming"], **kwargs) -> "StreamingFlow": ...  # noqa: ANN003

    @overload
    def get(self, flow_type: Literal["batch"], **kwargs) -> "BatchFlow": ...  # noqa: ANN003

    @overload
    def get(self, flow_type: None = None, **kwargs) -> Flow: ...  # noqa: ANN003

    def get(self, flow_type: str | None = None, **kwargs) -> Flow:  # noqa: ANN003
        """Used to get an instant result from the api.

        Args:
            flow_type: The type of flow to be returned
            **kwargs: Optional arguments to be passed to filter the results.

        Returns:
            An inherited instance of :py:class:`streamsets.sdk.sch_models.BaseModel`.

        Raises:
            TypeError: If flow_type is not 'streaming', 'batch', or None.
            ValueError: If instance is not in the list.
        """
        if flow_type is not None and flow_type not in ["streaming", "batch"]:
            raise TypeError(f"flow_type must be 'streaming', 'batch', got '{flow_type}'")

        if flow_type == "streaming":
            return self._streaming_flows.get(**kwargs)
        elif flow_type == "batch":
            return self._batch_flows.get(**kwargs)
        else:
            try:
                return self._streaming_flows.get(**kwargs)
            except Exception:
                return self._batch_flows.get(**kwargs)
