import abc
import inspect
import warnings
from operator import attrgetter
from typing import (
    Iterable,
    List,
    Optional,
    Tuple,
)

from spacy.tokens import Doc, Span

from edsnlp.core import PipelineProtocol
from edsnlp.core.registries import CurriedFactory
from edsnlp.utils.span_getters import (
    SpanGetter,  # noqa: F401
    SpanGetterArg,  # noqa: F401
    SpanSetter,
    SpanSetterArg,
    get_spans,  # noqa: F401
    set_spans,
    validate_span_getter,  # noqa: F401
    validate_span_setter,
)


def value_getter(span: Span):
    key = span._._get_key("value")
    if key in span.doc.user_data:
        return span.doc.user_data[key]
    return span._.get(span.label_) if span._.has(span.label_) else None


class BaseComponentMeta(abc.ABCMeta):
    def __init__(cls, name, bases, dct):
        super().__init__(name, bases, dct)

        sig = inspect.signature(cls.__init__)
        sig = sig.replace(parameters=tuple(sig.parameters.values())[1:])
        cls.__signature__ = sig

    def __call__(cls, nlp=inspect.Signature.empty, *args, **kwargs):
        # If this component is missing the nlp argument, we curry it with the
        # provided arguments and return a CurriedFactory object.
        sig = inspect.signature(cls.__init__)
        try:
            bound = sig.bind_partial(None, nlp, *args, **kwargs)
            bound.arguments.pop("self", None)
            if (
                "nlp" in sig.parameters
                and sig.parameters["nlp"].default is sig.empty
                and bound.arguments.get("nlp", sig.empty) is sig.empty
            ):
                return CurriedFactory(cls, bound.arguments)
            if nlp is inspect.Signature.empty:
                bound.arguments.pop("nlp", None)
        except TypeError:  # pragma: no cover
            if nlp is inspect.Signature.empty:
                super().__call__(*args, **kwargs)
            else:
                super().__call__(nlp, *args, **kwargs)
        return super().__call__(**bound.arguments)


class BaseComponent(abc.ABC, metaclass=BaseComponentMeta):
    """
    The `BaseComponent` adds a `set_extensions` method,
    called at the creation of the object.

    It helps decouple the initialisation of the pipeline from
    the creation of extensions, and is particularly usefull when
    distributing EDSNLP on a cluster, since the serialisation mechanism
    imposes that the extensions be reset.
    """

    def __init__(
        self,
        nlp: Optional[PipelineProtocol] = None,
        name: Optional[str] = None,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.name = name
        self.set_extensions()

    def set_extensions(self):
        """
        Set `Doc`, `Span` and `Token` extensions.
        """
        if Span.has_extension("value"):
            if Span.get_extension("value")[2] is not value_getter:
                warnings.warn(
                    "A Span extension 'value' already exists with a different getter. "
                    "Keeping the existing extension, but some components of edsnlp may "
                    "not work as expected."
                )
            return
        Span.set_extension(
            "value",
            getter=value_getter,
        )

    def get_spans(self, doc: Doc):  # noqa: F811
        """
        Returns sorted spans of interest according to the
        possible value of `on_ents_only`.
        Includes `doc.ents` by default, and adds eventual SpanGroups.
        """
        ents = list(doc.ents) + list(doc.spans.get("discarded", []))

        on_ents_only = getattr(self, "on_ents_only", None)

        if isinstance(on_ents_only, str):
            on_ents_only = [on_ents_only]
        if isinstance(on_ents_only, (set, list)):
            for spankey in set(on_ents_only) & set(doc.spans.keys()):
                ents.extend(doc.spans.get(spankey, []))

        return sorted(list(set(ents)), key=(attrgetter("start", "end")))

    def _boundaries(
        self, doc: Doc, terminations: Optional[List[Span]] = None
    ) -> List[Tuple[int, int]]:
        """
        Create sub sentences based sentences and terminations found in text.

        Parameters
        ----------
        doc:
            spaCy Doc object
        terminations:
            List of tuples with (match_id, start, end)

        Returns
        -------
        boundaries:
            List of tuples with (start, end) of spans
        """

        if terminations is None:
            terminations = []

        sent_starts = [sent.start for sent in doc.sents]
        termination_starts = [t.start for t in terminations]

        starts = sent_starts + termination_starts + [len(doc)]

        # Remove duplicates
        starts = list(set(starts))

        # Sort starts
        starts.sort()

        boundaries = [(start, end) for start, end in zip(starts[:-1], starts[1:])]

        return boundaries

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.set_extensions()


class BaseNERComponent(BaseComponent, abc.ABC):
    span_setter: SpanSetter

    def __init__(
        self,
        nlp: PipelineProtocol = None,
        name: str = None,
        *args,
        span_setter: SpanSetterArg,
        **kwargs,
    ):
        super().__init__(nlp, name, *args, **kwargs)
        self.span_setter: SpanSetter = validate_span_setter(span_setter)  # type: ignore

    def set_spans(self, doc, matches):
        return set_spans(doc, matches, self.span_setter)


class BaseSpanAttributeClassifierComponent(BaseComponent, abc.ABC):
    span_getter: SpanGetter
    attributes: Iterable[str]

    def __init__(
        self,
        nlp: PipelineProtocol = None,
        name: str = None,
        *args,
        span_getter: SpanGetterArg,
        **kwargs,
    ):
        super().__init__(nlp, name, *args, **kwargs)
        self.span_getter: SpanGetter = validate_span_getter(span_getter)  # type: ignore

    # For backwards compatibility
    @property
    def qualifiers(self):  # pragma: no cover
        return self.attributes

    @qualifiers.setter
    def qualifiers(self, value):  # pragma: no cover
        self.attributes = value