from functools import wraps
import traceback
from typing import Any, Callable, List, Protocol, Tuple, TypeVar

from kazu.data.data import Document, PROCESSING_EXCEPTION


class Step(Protocol):
    @classmethod
    def namespace(cls) -> str:
        """Metadata to name/describe the step, used in various places.

        Defaults to  ``cls.__name__``.
        """
        # Note: it would be nice for this to be a property, but you can't use @classmethod
        # and @property together in most versions of python (you could in 3.9 and 3.10), and
        # it being a classmethod is quite useful.
        return cls.__name__

    def __call__(self, docs: List[Document]) -> Tuple[List[Document], List[Document]]:
        """Process documents and respond with processed and failed documents.

        Note that many steps will be decorated by :func:`document_iterating_step` or
        :func:`document_batch_step` which will modify the 'original' ``__call__`` function
        signature to match the expected signature for a step, as the decorators handle the
        exception/failed documents logic for you.

        :param docs:
        :return: The first element is all the provided docs (now modified by the processing), the
            second is the docs that failed to (fully) process correctly.
        """
        raise NotImplementedError()


Self = TypeVar("Self")
"""A TypeVar for the type of the class whose method is decorated with :func:`document_iterating_step`
or :func:`document_batch_step`\\ ."""


def document_iterating_step(
    per_doc_callable: Callable[[Self, Document], Any]
) -> Callable[[Self, List[Document]], Tuple[List[Document], List[Document]]]:
    """Handle a list of :class:`~kazu.data.data.Document`\\ s and add error handling.

    Use this to decorate a method that processes a single :class:`~kazu.data.data.Document`\\ .
    The resulting method will then iterate over a list of
    :class:`~kazu.data.data.Document`\\ s, calling the decorated function for each
    :class:`~kazu.data.data.Document`\\ . Errors are handled automatically and added to the
    ``PROCESSING_EXCEPTION`` metadata of documents, with failed docs returned as the second element
    of the return value, as expected of the ``__call__`` method of a :class:`Step`\\ .

    Generally speaking, it will save effort and repetition to decorate a :class:`Step` with either
    :func:`document_iterating_step` or :func:`document_batch_step`\\, rather than implementing the
    error handling in the :class:`Step` itself.

    Normally, :func:`document_iterating_step` would be used in preference to
    :func:`document_batch_step`\\, unless the method involves computation which is more efficient
    when run in a batch, such as inference with a transformer-based Machine Learning model, or
    using spacy's `pipe <https://spacy.io/api/language/#pipe>`_ method.

    Note that this will only work for a method of a class, rather than a standalone function,
    as it expects to have to pass through 'self' as a parameter.

    :param per_doc_callable: A function that processes a single document, that you want to use as
        the :literal:`__call__` method of a :class:`Step`\\ . This must do its work by mutating the
        input document: the return value is ignored.
    """
    # note - this excludes __annotations__ as the returned function has a different type signature
    @wraps(per_doc_callable, assigned=("__module__", "__name__", "__qualname__", "__doc__"))
    def step_call(self: Self, docs: List[Document]) -> Tuple[List[Document], List[Document]]:
        failed_docs = []
        for doc in docs:
            try:
                per_doc_callable(self, doc)
            except Exception:
                doc.metadata[PROCESSING_EXCEPTION] = traceback.format_exc()
                failed_docs.append(doc)
        return docs, failed_docs

    return step_call


def document_batch_step(
    batch_doc_callable: Callable[[Self, List[Document]], Any]
) -> Callable[[Self, List[Document]], Tuple[List[Document], List[Document]]]:
    """Add error handling to a method that processes batches of :class:`~kazu.data.data.Document`\\ s.

    Use this to decorate a method that processes a batch of :class:`~kazu.data.data.Document`\\ s
    at a time. The resulting method will wrap a call to the decorated function with error handling
    which will added exceptions to the ``PROCESSING_EXCEPTION`` metadata of documents. Failed documents
    will be returned as the second element of the return value, as expected of the ``__call__``
    method of a :class:`Step`\\ .

    Generally speaking, it will save effort and repetition to decorate a :class:`Step` with either
    :func:`document_iterating_step` or :func:`document_batch_step`\\, rather than implementing the
    error handling in the :class:`Step` itself.

    Normally, :func:`document_iterating_step` would be used in preference to
    :func:`document_batch_step`\\, unless the method involves computation which is more efficient
    when run in a batch, such as inference with a transformer-based Machine Learning model, or
    using spacy's `pipe <https://spacy.io/api/language/#pipe>`_ method.

    Note that this will only work for a method of a class, rather than a standalone function,
    as it expects to have to pass through 'self' as a parameter.

    :param batch_doc_callable: A function that processes a batch of documents, that you want to use
        as the :literal:`__call__` method of a :class:`Step`. This must do its work by mutating the
        input documents: the return value is ignored.
    """
    # note - this excludes __annotations__ as the returned function has a different type signature
    @wraps(batch_doc_callable, assigned=("__module__", "__name__", "__qualname__", "__doc__"))
    def step_call(self: Self, docs: List[Document]) -> Tuple[List[Document], List[Document]]:
        failed_docs = []
        try:
            batch_doc_callable(self, docs)
        except Exception:
            affected_doc_ids = [doc.idx for doc in docs]
            for doc in docs:
                message = (
                    f"batch failed: affected ids: {affected_doc_ids}\n" + traceback.format_exc()
                )
                doc.metadata[PROCESSING_EXCEPTION] = message
                failed_docs.append(doc)
        return docs, failed_docs

    return step_call
