from spacy.tokens import Doc
from spacy.tokens.span import Span

import edsnlp
import edsnlp.pipes as eds


def test_split_line_jump():
    txt = """This is a test. Another test.

A third test!"""
    nlp = edsnlp.blank("eds")
    nlp.add_pipe("eds.matcher", config={"terms": {"test": "test"}})
    doc = nlp(txt)
    Span.set_extension("test_dict", default={}, force=True)
    Doc.set_extension("global_attr", default=None, force=True)
    Span.set_extension("ent_attr", default=None, force=True)
    doc._.global_attr = "global"
    doc.ents[0]._.test_dict = {"key": doc.ents[1]}
    doc.ents[0]._.ent_attr = "ent-0"
    doc.ents[1]._.test_dict = {"key": doc.ents[0]}
    doc.ents[2]._.test_dict = {"key": doc.ents[0]}
    doc.ents[2]._.ent_attr = "ent-2"
    doc.spans["section"] = [doc[:]]
    doc.spans["section"][0]._.ent_attr = "section"
    subdocs = list(eds.split(regex="\n\n")(doc))

    assert len(subdocs) == 2
    assert subdocs[0].text == "This is a test. Another test.\n\n"
    assert subdocs[0]._.global_attr == "global"
    assert subdocs[0].ents[0]._.test_dict["key"] == subdocs[0].ents[1]
    assert subdocs[0].ents[0]._.ent_attr == "ent-0"
    assert subdocs[0].ents[1]._.test_dict["key"] == subdocs[0].ents[0]
    assert subdocs[0].spans["section"][0].text == "This is a test. Another test.\n\n"
    assert subdocs[0].spans["section"][0]._.ent_attr == "section"

    assert subdocs[1].text == "A third test!"
    assert subdocs[1]._.global_attr == "global"
    assert subdocs[1].ents[0]._.test_dict == {}
    assert subdocs[1].ents[0]._.ent_attr == "ent-2"
    assert subdocs[1].spans["section"][0].text == "A third test!"
    assert subdocs[1].spans["section"][0]._.ent_attr == "section"


def test_filter():
    txt = """This is a test. Another test."""
    doc = edsnlp.blank("eds")(txt)
    subdocs = list(
        eds.split(
            regex=r"[.!?]\s+()[A-Z]",
            filter_expr='"Another" in doc.text',
        )(doc)
    )
    assert len(subdocs) == 1
    assert subdocs[0].text == "Another test."


def test_max_length():
    txt = (
        "Le patient mange des pates depuis le début du confinement, "
        "il est donc un peu ballonné, mais pense revenir à un régime plus "
        "équilibré en mangeant des légumes et des fruits."
    )
    doc = edsnlp.blank("eds")(txt)
    texts = [d.text for d in eds.split(max_length=12)(doc)]
    assert texts == [
        "Le patient mange des pates depuis le début du confinement, il ",
        "est donc un peu ballonné, mais pense revenir à un régime ",
        "plus équilibré en mangeant des légumes et des fruits.",
    ]