import filecmp

import re

from itertools import islice

from os import listdir

from os.path import join

from pathlib import Path

from random import choice, randint, random

from string import ascii_letters, ascii_lowercase

import pytest

import edsnlp

from edsnlp.connectors.brat import BratConnector

from edsnlp.core import PipelineProtocol

def random_word():

    n = randint(1, 20)

    return "".join(

        [choice(ascii_letters)] + [choice(ascii_lowercase) for _ in range(n)]

    )

def random_text():

    n = randint(30, 60)

    return " ".join([random_word() for _ in range(n)])

def random_brat_file(text):

    brat = []

    for match in re.finditer(r"\w+", text):

        if random() > 0.8:

            line = (

                f"T{len(brat) + 1}\tTEST {match.start()} {match.end()}\t{match.group()}"

            )

            brat.append(line)

    return "\n".join(brat) + "\n"

@pytest.fixture

def brat_folder(tmpdir):

    for i in range(100):

        text = random_text()

        brat = random_brat_file(text)

        with open(join(tmpdir, f"{i}.txt"), "w") as f:

            f.write(text)

        with open(join(tmpdir, f"{i}.ann"), "w") as f:

            if i == 0:

                f.write("\n")

            else:

                f.write(brat)

    return tmpdir

# OLD BratConnector tests, deprecated

@pytest.fixture

def brat1(brat_folder) -> BratConnector:

    return BratConnector(brat_folder)

@pytest.fixture

def brat2(tmpdir) -> BratConnector:

    return BratConnector(tmpdir)

@pytest.fixture

def brat_importer():

    brat_dir = Path(__file__).parent.parent.resolve() / "resources" / "brat_data"

    return BratConnector(str(brat_dir), bool_attributes=["bool flag 0"])

@pytest.fixture

def brat_exporter(tmpdir):

    return BratConnector(tmpdir, attributes=["etat", "assertion", "bool flag 0"])

def test_empty_brat(brat2: BratConnector, blank_nlp: PipelineProtocol):

    with pytest.raises(AssertionError):

        brat2.brat2docs(blank_nlp)

def test_brat2brat(

    brat1: BratConnector, brat2: BratConnector, blank_nlp: PipelineProtocol

):

    docs = brat1.brat2docs(blank_nlp)

    brat2.docs2brat(docs)

    files = listdir(brat1.directory)

    assert files

    for file in files:

        assert file in listdir(brat2.directory)

        assert filecmp.cmp(join(brat1.directory, file), join(brat2.directory, file))

def test_docs2brat(nlp, brat2):

    text = (

        "Le patient est admis pour des douleurs dans le bras droit, "

        "mais n'a pas de problème de locomotion. "

        "Historique d'AVC dans la famille. pourrait être un cas de rhume.\n"

        "NBNbWbWbNbWbNBNbNbWbWbNBNbWbNbNbWbNBNbWbNbNBWbWbNbNbN"

        "BWbNbWbNbWBNbNbWbNbNBNbWbWbNbWBNbNbWbNBNbWbWbNb\n"

        "Pourrait être un cas de rhume.\n"

        "Motif :\n"

        "Douleurs dans le bras droit."

    )

    doc1 = nlp(text)

    doc1.ents = doc1.spans["pollutions"]

    doc2 = nlp(text)

    doc2.ents = doc2.spans["section_titles"]

    docs = [doc1, doc2]

    for i, doc in enumerate(docs):

        doc._.note_id = i + 1

    brat2.docs2brat(docs)

def assert_doc_read(doc):

    assert doc._.note_id == "subfolder/doc-1"

    attrs = ("etat", "assertion", "bool flag 0")

    spans_and_attributes = {

        "__ents__": sorted(

            [

                (e.start, e.end, e.text, tuple(getattr(e._, key) for key in attrs))

                for e in doc.ents

            ]

        ),

        **{

            name: sorted(

                [

                    (e.start, e.end, e.text, tuple(getattr(e._, key) for key in attrs))

                    for e in doc.spans[name]

                ]

            )

            for name in doc.spans

        },

    }

    assert spans_and_attributes == {

        "__ents__": [

            (6, 7, "douleurs", (None, None, False)),

            (7, 11, "dans le bras droit", (None, None, False)),

            (17, 21, "problème \nde locomotion", (None, "absent", True)),

            (25, 26, "AVC", ("passé", "non-associé", False)),

            (35, 36, "rhume", ("présent", "hypothétique", False)),

            (45, 46, "rhume", ("présent", "hypothétique", False)),

            (51, 52, "Douleurs", (None, None, False)),

            (52, 56, "dans le bras droit", (None, None, False)),

            (68, 69, "anomalie", (None, "absent", False)),

        ],

        "anatomie": [

            (9, 11, "bras droit", (None, None, False)),

            (54, 56, "bras droit", (None, None, False)),

        ],

        "localisation": [

            (7, 11, "dans le bras droit", (None, None, False)),

            (52, 56, "dans le bras droit", (None, None, False)),

        ],

        "pathologie": [

            (17, 21, "problème \nde locomotion", (None, "absent", True)),

            (25, 26, "AVC", ("passé", "non-associé", False)),

            (35, 36, "rhume", ("présent", "hypothétique", False)),

            (45, 46, "rhume", ("présent", "hypothétique", False)),

        ],

        "sosy": [

            (6, 7, "douleurs", (None, None, False)),

            (51, 52, "Douleurs", (None, None, False)),

            (68, 69, "anomalie", (None, "absent", False)),

        ],

        "test label 0": [(68, 69, "anomalie", (None, "absent", False))],

    }

def assert_doc_write(exported_ann_text):

    assert exported_ann_text == (

        "T1 sosy 30 38  douleurs\n"

        "A1 etat T1 test\n"

        "T2 localisation 39 57  dans le bras droit\n"

        "T3 anatomie 47 57  bras droit\n"

        "T4 pathologie 75 83;85 98  problème de locomotion\n"

        "A2 assertion T4 absent\n"

        "A3 bool flag 0 T4\n"

        "T5 pathologie 114 117  AVC\n"

        "A4 etat T5 passé\n"

        "A5 assertion T5 non-associé\n"

        "T6 pathologie 159 164  rhume\n"

        "A6 etat T6 présent\n"

        "A7 assertion T6 hypothétique\n"

        "T7 pathologie 291 296  rhume\n"

        "A8 etat T7 présent\n"

        "A9 assertion T7 hypothétique\n"

        "T8 sosy 306 314    Douleurs\n"

        "T9 localisation 315 333    dans le bras droit\n"

        "T10    anatomie 323 333    bras droit\n"

        "T11    sosy 378 386    anomalie\n"

        "A10    assertion T11 absent\n"

        "T12    test label 0 378 386    anomalie\n"

        "A11    assertion T12 absent\n"

    )

def test_brat(

    brat_importer: BratConnector,

    brat_exporter: BratConnector,

    blank_nlp: PipelineProtocol,

):

    doc = brat_importer.brat2docs(blank_nlp)[0]

    assert_doc_read(doc)

    doc.ents[0]._.etat = "test"

    brat_exporter.docs2brat([doc])

    with open(brat_exporter.directory / "subfolder" / "doc-1.ann") as f:

        exported_ann_text = f.read()

    assert_doc_write(exported_ann_text)

# New `edsnlp.data.read_standoff` and `edsnlp.data.write_standoff` tests

def test_read_to_standoff(blank_nlp, tmpdir):

    input_dir = Path(__file__).parent.parent.resolve() / "resources" / "brat_data"

    output_dir = Path(tmpdir)

    doc = list(

        edsnlp.data.read_standoff(

            input_dir,

            bool_attributes=["bool flag 0"],

            notes_as_span_attribute="cui",

        )

    )[0]

    assert_doc_read(doc)

    doc.ents[0]._.etat = "test"

    doc.ents[0]._.cui = "C0030193"

    edsnlp.data.write_standoff(

        [doc],

        output_dir,

        span_attributes=["etat", "assertion", "bool flag 0"],

        span_getter=[

            "ents",

            "sosy",

            "localisation",

            "anatomie",

            "pathologie",

            "test label 0",

        ],

    )

    with open(output_dir / "subfolder" / "doc-1.ann") as f:

        exported_ann_text = f.read()

    assert_doc_write(exported_ann_text)

@pytest.mark.parametrize("num_cpu_workers", [0, 2])

def test_read_shuffle_loop(num_cpu_workers: int):

    notes = (

        edsnlp.data.read_standoff(

            Path(__file__).parent.parent.resolve() / "resources" / "brat_data",

            shuffle="dataset",

            keep_txt_only_docs=True,

            seed=42,

            loop=True,

        )

        .map(lambda x: x._.note_id)

        .set_processing(num_cpu_workers=num_cpu_workers)

    )

    notes = list(islice(notes, 6))

    assert notes == [

        "subfolder/doc-2",

        "subfolder/doc-1",

        "subfolder/doc-3",

        "subfolder/doc-3",

        "subfolder/doc-2",

        "subfolder/doc-1",

    ]