from typing import Any, Callable, Dict, Iterable, Optional

import spacy
from spacy import util
from spacy.errors import Errors
from spacy.language import FactoryMeta
from spacy.pipe_analysis import validate_attrs
from spacy.util import SimpleFrozenDict, SimpleFrozenList, registry


@classmethod
def factory(
    cls,
    name: str,
    *,
    default_config: Dict[str, Any] = SimpleFrozenDict(),
    assigns: Iterable[str] = SimpleFrozenList(),
    requires: Iterable[str] = SimpleFrozenList(),
    retokenizes: bool = False,
    default_score_weights: Dict[str, Optional[float]] = SimpleFrozenDict(),
    func: Optional[Callable] = None,
) -> Callable:
    """
    Patched from spaCy to allow back dots in factory
    names (https://github.com/aphp/edsnlp/pull/152)

    Register a new pipeline component factory. Can be used as a decorator
    on a function or classmethod, or called as a function with the factory
    provided as the func keyword argument. To create a component and add
    it to the pipeline, you can use nlp.add_pipe(name).

    name (str): The name of the component factory.
    default_config (Dict[str, Any]): Default configuration, describing the
        default values of the factory arguments.
    assigns (Iterable[str]): Doc/Token attributes assigned by this component,
        e.g. "token.ent_id". Used for pipeline analysis.
    requires (Iterable[str]): Doc/Token attributes required by this component,
        e.g. "token.ent_id". Used for pipeline analysis.
    retokenizes (bool): Whether the component changes the tokenization.
        Used for pipeline analysis.
    default_score_weights (Dict[str, Optional[float]]): The scores to report during
        training, and their default weight towards the final score used to
        select the best model. Weights should sum to 1.0 per component and
        will be combined and normalized for the whole pipeline. If None,
        the score won't be shown in the logs or be weighted.
    func (Optional[Callable]): Factory function if not used as a decorator.

    DOCS: https://spacy.io/api/language#factory
    """
    if not isinstance(name, str):
        raise ValueError(Errors.E963.format(decorator="factory"))
    if not isinstance(default_config, dict):
        raise ValueError(
            Errors.E962.format(
                style="default config", name=name, cfg_type=type(default_config)
            )
        )

    def add_factory(factory_func: Callable) -> Callable:
        internal_name = cls.get_factory_name(name)
        if internal_name in registry.factories:
            # We only check for the internal name here – it's okay if it's a
            # subclass and the base class has a factory of the same name. We
            # also only raise if the function is different to prevent raising
            # if module is reloaded.
            existing_func = registry.factories.get(internal_name)
            if not util.is_same_func(factory_func, existing_func):
                err = Errors.E004.format(
                    name=name, func=existing_func, new_func=factory_func
                )
                raise ValueError(err)

        arg_names = util.get_arg_names(factory_func)
        if "nlp" not in arg_names or "name" not in arg_names:
            raise ValueError(Errors.E964.format(name=name))
        # Officially register the factory so we can later call
        # registry.resolve and refer to it in the config as
        # @factories = "spacy.Language.xyz". We use the class name here so
        # different classes can have different factories.
        registry.factories.register(internal_name, func=factory_func)
        factory_meta = FactoryMeta(
            factory=name,
            default_config=default_config,
            assigns=validate_attrs(assigns),
            requires=validate_attrs(requires),
            scores=list(default_score_weights.keys()),
            default_score_weights=default_score_weights,
            retokenizes=retokenizes,
        )
        cls.set_factory_meta(name, factory_meta)
        # We're overwriting the class attr with a frozen dict to handle
        # backwards-compat (writing to Language.factories directly). This
        # wouldn't work with an instance property and just produce a
        # confusing error – here we can show a custom error
        cls.factories = SimpleFrozenDict(
            registry.factories.get_all(), error=Errors.E957
        )
        return factory_func

    if func is not None:  # Support non-decorator use cases
        return add_factory(func)
    return add_factory


spacy.Language.factory = factory
