from abc import abstractmethod

from theflow.base import Composable


class BaseComponent(Composable):
    """Base class for component

    A component is a class that can be used to compose a pipeline. To use the
    component, you should implement the following methods:
        - run_raw: run on raw input
        - run_batch_raw: run on batch of raw input
        - run_document: run on document
        - run_batch_document: run on batch of documents
        - is_document: check if input is document
        - is_batch: check if input is batch
    """

    @abstractmethod
    def run_raw(self, *args, **kwargs):
        ...

    @abstractmethod
    def run_batch_raw(self, *args, **kwargs):
        ...

    @abstractmethod
    def run_document(self, *args, **kwargs):
        ...

    @abstractmethod
    def run_batch_document(self, *args, **kwargs):
        ...

    @abstractmethod
    def is_document(self, *args, **kwargs) -> bool:
        ...

    @abstractmethod
    def is_batch(self, *args, **kwargs) -> bool:
        ...

    def run(self, *args, **kwargs):
        """Run the component."""

        is_document = self.is_document(*args, **kwargs)
        is_batch = self.is_batch(*args, **kwargs)

        if is_document and is_batch:
            return self.run_batch_document(*args, **kwargs)
        elif is_document and not is_batch:
            return self.run_document(*args, **kwargs)
        elif not is_document and is_batch:
            return self.run_batch_raw(*args, **kwargs)
        else:
            return self.run_raw(*args, **kwargs)