import os
import subprocess
import sys
from io import BytesIO

import pytest
from confit import Config
from confit.errors import ConfitValidationError
from confit.registry import validate_arguments
from spacy.tokens import Doc

import edsnlp
import edsnlp.pipes as eds
from edsnlp import Pipeline, registry
from edsnlp.core.registries import CurriedFactory
from edsnlp.pipes.base import BaseComponent

try:
    import torch.nn
except ImportError:
    torch = None


class CustomClass:
    pass

    def __call__(self, doc: Doc) -> Doc:
        return doc


def test_add_pipe_factory():
    model = edsnlp.blank("eds")
    model.add_pipe("eds.normalizer", name="normalizer")
    assert "normalizer" in model.pipe_names
    assert model.has_pipe("normalizer")

    model.add_pipe("eds.sentences", name="sentences")
    assert "sentences" in model.pipe_names
    assert model.has_pipe("sentences")

    with pytest.raises(ValueError):
        model.get_pipe("missing-pipe")


def test_add_pipe_component():
    model = edsnlp.blank("eds")
    model.add_pipe(eds.normalizer(nlp=model), name="normalizer")
    assert "normalizer" in model.pipe_names
    assert model.has_pipe("normalizer")
    assert model.pipes.normalizer is model.get_pipe("normalizer")

    model.add_pipe(eds.sentences(nlp=model), name="sentences")
    assert "sentences" in model.pipe_names
    assert model.has_pipe("sentences")
    assert model.pipes.sentences is model.get_pipe("sentences")

    with pytest.raises(ValueError):
        model.add_pipe(
            eds.sentences(nlp=model, name="sentences"),
            config={"punct_chars": ".?!"},
        )

    with pytest.raises(ValueError):
        model.add_pipe(CustomClass())


def test_sequence(frozen_ml_nlp: Pipeline):
    assert len(frozen_ml_nlp.pipeline) == 3
    assert list(frozen_ml_nlp.pipeline) == [
        ("sentences", frozen_ml_nlp.get_pipe("sentences")),
        ("transformer", frozen_ml_nlp.get_pipe("transformer")),
        ("ner", frozen_ml_nlp.get_pipe("ner")),
    ]
    assert list(frozen_ml_nlp.torch_components()) == [
        ("transformer", frozen_ml_nlp.get_pipe("transformer")),
        ("ner", frozen_ml_nlp.get_pipe("ner")),
    ]


def test_disk_serialization(tmp_path, ml_nlp):
    nlp = ml_nlp

    assert nlp.get_pipe("transformer").stride == 96
    ner = nlp.get_pipe("ner")
    ner.update_labels(["PERSON", "GIFT"])

    os.makedirs(tmp_path / "model", exist_ok=True)
    # by default, vocab is excluded
    nlp.to_disk(tmp_path / "model", exclude=set())

    assert (tmp_path / "model" / "config.cfg").exists()
    assert (tmp_path / "model" / "ner" / "parameters.safetensors").exists()
    assert (tmp_path / "model" / "transformer" / "parameters.safetensors").exists()
    # fmt: off
    assert (
            (tmp_path / "model" / "transformer" / "pytorch_model.bin").exists() or
            (tmp_path / "model" / "transformer" / "model.safetensors").exists()
    )
    # fmt: on

    assert (tmp_path / "model" / "config.cfg").read_text() == (
        config_str.replace("components = ${components}\n", "").replace(
            "prajjwal1/bert-tiny", "./transformer"
        )
    )

    nlp = edsnlp.load(
        tmp_path / "model",
        overrides={"components": {"transformer": {"stride": 64}}},
    )
    assert nlp.get_pipe("ner").labels == ["PERSON", "GIFT"]
    assert nlp.get_pipe("transformer").stride == 64


config_str = """\
[nlp]
lang = "eds"
pipeline = ["sentences", "transformer", "ner"]
components = ${components}

[nlp.tokenizer]
@tokenizers = "eds.tokenizer"

[components]

[components.sentences]
@factory = "eds.sentences"

[components.transformer]
@factory = "eds.transformer"
model = "prajjwal1/bert-tiny"
window = 128
stride = 96

[components.ner]
@factory = "eds.ner_crf"
embedding = ${components.transformer}
mode = "independent"
target_span_getter = ["ents", "ner-preds"]
labels = ["PERSON", "GIFT"]
infer_span_setter = false
window = 40
stride = 20

[components.ner.span_setter]
ents = true

"""


@pytest.mark.skipif(torch is None, reason="torch not installed")
def test_validate_config():
    @validate_arguments
    def function(model: Pipeline):
        assert len(model.pipe_names) == 3

    function(Config.from_str(config_str).resolve(registry=registry)["nlp"])


def test_torch_module(frozen_ml_nlp: Pipeline):
    with frozen_ml_nlp.train(True):
        for name, component in frozen_ml_nlp.torch_components():
            assert component.training is True

    with frozen_ml_nlp.train(False):
        for name, component in frozen_ml_nlp.torch_components():
            assert component.training is False

    frozen_ml_nlp.to("cpu")


def test_cache(frozen_ml_nlp: Pipeline):
    from edsnlp.core.torch_component import _caches

    text = "Ceci est un exemple"
    frozen_ml_nlp(text)

    doc = frozen_ml_nlp.make_doc(text)
    with frozen_ml_nlp.cache():
        for name, pipe in frozen_ml_nlp.pipeline:
            # This is a hack to get around the ambiguity
            # between the __call__ method of Pytorch modules
            # and the __call__ methods of spacy components
            if hasattr(pipe, "batch_process"):
                doc = next(iter(pipe.batch_process([doc])))
            else:
                doc = pipe(doc)
        trf_forward_cache_entries = [
            key
            for key in _caches["default"]
            if isinstance(key, tuple) and key[0] == "forward"
        ]
        assert len(trf_forward_cache_entries) == 2

    assert len(_caches) == 0


def test_select_pipes(frozen_ml_nlp: Pipeline):
    text = "Ceci est un exemple"
    with frozen_ml_nlp.select_pipes(enable=["transformer", "ner"]):
        assert len(frozen_ml_nlp.disabled) == 1
        assert not frozen_ml_nlp(text).has_annotation("SENT_START")
    assert len(frozen_ml_nlp.disabled) == 0


@pytest.mark.skip(reason="Deprecated behavior")
def test_different_names():
    nlp = edsnlp.blank("eds")

    extractor = eds.sentences(nlp=nlp, name="custom_name")

    with pytest.raises(ValueError) as exc_info:
        nlp.add_pipe(extractor, name="sentences")

    assert (
        "The provided name 'sentences' does not "
        "match the name of the component 'custom_name'."
    ) in str(exc_info.value)


@pytest.mark.skipif(torch is None, reason="torch not installed")
def test_load_config(run_in_test_dir):
    nlp = edsnlp.load("training/qlf_config.yml")
    assert nlp.pipe_names == [
        "normalizer",
        "sentencizer",
        "covid",
        "qualifier",
    ]


fail_config = """
[nlp]
lang = "eds"
pipeline = ["transformer", "ner"]

[nlp.tokenizer]
@tokenizers = "eds.tokenizer"

[components]

[components.transformer]
@factory = "eds.transformer"
model = "prajjwal1/bert-tiny"
window = 128
stride = 96

[components.ner]
@factory = "eds.ner_crf"
embedding = ${components.transformer}
mode = "error-mode"
span_setter = "ents"
"""


@pytest.mark.skipif(torch is None, reason="torch not installed")
def test_config_validation_error():
    with pytest.raises(ConfitValidationError) as e:
        Pipeline.from_config(Config.from_str(fail_config))

    assert "1 validation error for" in str(e.value)
    assert "got 'error-mode'" in str(e.value)


@edsnlp.registry.factory.register("test_wrapper", spacy_compatible=False)
class WrapperComponent:
    def __init__(self, *, copy_list, copy_dict, sub):
        pass


fail_config_sub = """
nlp:
    lang: "eds"
    components:
        wrapper:
            "@factory": "test_wrapper"

            copy_list:
                - ${nlp.components.wrapper.sub}

            copy_dict:
                key: ${nlp.components.wrapper.sub}

            sub:
                "@factory": "eds.matcher"
                terms: 100.0  # clearly wrong

        matcher_copy: ${nlp.components.wrapper.sub}
"""


def test_config_sub_validation_error():
    with pytest.raises(ConfitValidationError):
        Pipeline.from_config(Config.from_yaml_str(fail_config_sub))

    fix = {"nlp": {"components": {"wrapper": {"sub": {"terms": {"pattern": ["ok"]}}}}}}
    Pipeline.from_config(Config.from_yaml_str(fail_config_sub).merge(fix))


def test_add_pipe_validation_error():
    model = edsnlp.blank("eds")
    with pytest.raises(ConfitValidationError) as e:
        model.add_pipe("eds.covid", name="extractor", config={"foo": "bar"})

    assert str(e.value) == (
        "1 validation error for "
        "edsnlp.pipes.ner.covid.factory.create_component()\n"
        "-> extractor.foo\n"
        "   unexpected keyword argument"
    )


def test_spacy_component():
    nlp = edsnlp.blank("fr")
    nlp.add_pipe("sentencizer")


def test_rule_based_pipeline():
    nlp = edsnlp.blank("eds")
    nlp.add_pipe("eds.normalizer")
    nlp.add_pipe("eds.covid")

    assert nlp.pipe_names == ["normalizer", "covid"]
    assert nlp.get_pipe("normalizer") == nlp.pipeline[0][1]
    assert nlp.has_pipe("covid")

    with pytest.raises(ValueError) as exc_info:
        nlp.get_pipe("unknown")

    assert str(exc_info.value) == "Pipe 'unknown' not found in pipeline."

    doc = nlp.make_doc("Mon patient a le covid")

    new_doc = nlp(doc)

    assert len(doc.ents) == 1
    assert new_doc is doc

    assert nlp.get_pipe_meta("covid").assigns == ["doc.ents", "doc.spans"]


def test_torch_save(ml_nlp):
    import torch

    ml_nlp.get_pipe("ner").update_labels(["LOC", "PER"])
    buffer = BytesIO()
    torch.save(ml_nlp, buffer)
    buffer.seek(0)
    nlp = torch.load(buffer, weights_only=False)
    assert nlp.get_pipe("ner").labels == ["LOC", "PER"]
    assert len(list(nlp("Une phrase. Deux phrases.").sents)) == 2


def test_parameters(frozen_ml_nlp):
    assert len(list(frozen_ml_nlp.parameters())) == 42


def test_missing_factory(nlp):
    with pytest.raises(ValueError) as exc_info:
        nlp.add_pipe("__test_missing_pipe__")

    assert "__test_missing_pipe__" in str(exc_info.value)


@edsnlp.registry.factory("custom-curry-test")
class CustomComponent(BaseComponent):
    def __init__(self, nlp, name):
        self.nlp = nlp

    def __call__(self, doc):
        return doc


def test_curried_nlp_pipe():
    nlp = edsnlp.blank("eds")
    nlp.add_pipe(eds.sentences(name="my-sentences"))
    nlp.add_pipe(eds.normalizer())
    nlp.add_pipe(eds.sections(), name="sections")
    pipe = CustomComponent()

    assert isinstance(pipe, CurriedFactory)
    err = (
        f"This component CurriedFactory({pipe.factory}) has not been instantiated "
        f"yet, likely because it was missing an `nlp` pipeline argument. You should "
        f"either:\n"
        "- add it to a pipeline: `pipe = nlp.add_pipe(pipe)`\n"
        "- or fill its `nlp` argument: `pipe = factory(nlp=nlp, ...)`"
    )
    with pytest.raises(TypeError) as exc_info:
        pipe("Demo texte")
    assert str(exc_info.value) == err

    with pytest.raises(TypeError) as exc_info:
        pipe.forward("Demo texte")
    assert str(exc_info.value) == err

    nlp.add_pipe(pipe, name="custom")

    assert nlp.pipes.custom.nlp is nlp

    assert nlp.pipe_names == ["my-sentences", "normalizer", "sections", "custom"]


@pytest.mark.skipif(
    sys.version_info < (3, 8),
    reason="Can't run on GH CI with Python 3.7",
)
@pytest.mark.skipif(torch is None, reason="torch not installed")
def test_huggingface():
    nlp = edsnlp.load(
        "AP-HP/dummy-ner",
        auto_update=True,
        install_dependencies=True,
    )
    doc = nlp("On lui prescrit du paracetamol à 500mg.")
    assert doc.ents[0].text == "paracetamol"
    assert doc.ents[1].text == "500mg"

    # Try loading it twice for coverage
    edsnlp.load(
        "AP-HP/dummy-ner",
        auto_update=True,
        install_dependencies=True,
    )

    subprocess.run(["pip", "uninstall", "dummy-pip-package", "-y"], check=True)


@pytest.mark.skipif(
    sys.version_info < (3, 8),
    reason="Can't run on GH CI with Python 3.7",
)
def test_missing_huggingface():
    with pytest.raises(ValueError) as exc_info:
        edsnlp.load(
            "AP-HP/does-not-exist",
            auto_update=True,
        )

    assert "The load function expects either :" in str(exc_info.value)


def test_repr(frozen_ml_nlp):
    with frozen_ml_nlp.select_pipes(disable=["sentences"]):
        assert (
            repr(frozen_ml_nlp)
            == """\
Pipeline(lang=eds, pipes={
  "sentences": [disabled] eds.sentences,
  "transformer": eds.transformer,
  "ner": eds.ner_crf
})"""
        )


@edsnlp.registry.factory.register("test_nlp_less", spacy_compatible=False)
class NlpLessComponent:
    def __init__(self, nlp=None, name: str = "nlp_less", *, value: int):
        self.value = value
        self.name = name

    def __call__(self, doc):
        return doc


def test_nlp_less_component():
    component = NlpLessComponent(value=42)
    assert component.value == 42

    config = """
[component]
@factory = "test_nlp_less"
value = 42
"""
    component = Config.from_str(config).resolve(registry=registry)["component"]
    assert component.value == 42