import json
from pathlib import Path

import pytest

import edsnlp


def assert_doc_read(doc):
    assert doc._.note_id == "subfolder/doc-1"
    assert doc._.context_var == "test"

    attrs = ("etat", "assertion")
    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)),
            (7, 11, "dans le bras droit", (None, None)),
            (17, 21, "problème \nde locomotion", (None, "absent")),
            (25, 26, "AVC", ("passé", "non-associé")),
            (35, 36, "rhume", ("présent", "hypothétique")),
            (45, 46, "rhume", ("présent", "hypothétique")),
            (51, 52, "Douleurs", (None, None)),
            (52, 56, "dans le bras droit", (None, None)),
            (68, 69, "anomalie", (None, "absent")),
        ],
        "anatomie": [
            (9, 11, "bras droit", (None, None)),
            (54, 56, "bras droit", (None, None)),
        ],
        "localisation": [
            (7, 11, "dans le bras droit", (None, None)),
            (52, 56, "dans le bras droit", (None, None)),
        ],
        "pathologie": [
            (17, 21, "problème \nde locomotion", (None, "absent")),
            (25, 26, "AVC", ("passé", "non-associé")),
            (35, 36, "rhume", ("présent", "hypothétique")),
            (45, 46, "rhume", ("présent", "hypothétique")),
        ],
        "sosy": [
            (6, 7, "douleurs", (None, None)),
            (51, 52, "Douleurs", (None, None)),
            (68, 69, "anomalie", (None, "absent")),
        ],
    }


def assert_doc_write(exported_obj):
    assert exported_obj == {
        "entities": [
            {
                "assertion": None,
                "end_char": 38,
                "etat": "test",
                "lexical_variant": "douleurs",
                "note_nlp_id": 0,
                "note_nlp_source_value": "sosy",
                "start_char": 30,
            },
            {
                "assertion": None,
                "end_char": 57,
                "etat": None,
                "lexical_variant": "dans le bras droit",
                "note_nlp_id": 1,
                "note_nlp_source_value": "localisation",
                "start_char": 39,
            },
            {
                "assertion": None,
                "end_char": 57,
                "etat": None,
                "lexical_variant": "bras droit",
                "note_nlp_id": 2,
                "note_nlp_source_value": "anatomie",
                "start_char": 47,
            },
            {
                "assertion": "absent",
                "end_char": 98,
                "etat": None,
                "lexical_variant": "problème \nde locomotion",
                "note_nlp_id": 3,
                "note_nlp_source_value": "pathologie",
                "start_char": 75,
            },
            {
                "assertion": "non-associé",
                "end_char": 117,
                "etat": "passé",
                "lexical_variant": "AVC",
                "note_nlp_id": 4,
                "note_nlp_source_value": "pathologie",
                "start_char": 114,
            },
            {
                "assertion": "hypothétique",
                "end_char": 164,
                "etat": "présent",
                "lexical_variant": "rhume",
                "note_nlp_id": 5,
                "note_nlp_source_value": "pathologie",
                "start_char": 159,
            },
            {
                "assertion": "hypothétique",
                "end_char": 296,
                "etat": "présent",
                "lexical_variant": "rhume",
                "note_nlp_id": 6,
                "note_nlp_source_value": "pathologie",
                "start_char": 291,
            },
            {
                "assertion": None,
                "end_char": 314,
                "etat": None,
                "lexical_variant": "Douleurs",
                "note_nlp_id": 7,
                "note_nlp_source_value": "sosy",
                "start_char": 306,
            },
            {
                "assertion": None,
                "end_char": 333,
                "etat": None,
                "lexical_variant": "dans le bras droit",
                "note_nlp_id": 8,
                "note_nlp_source_value": "localisation",
                "start_char": 315,
            },
            {
                "assertion": None,
                "end_char": 333,
                "etat": None,
                "lexical_variant": "bras droit",
                "note_nlp_id": 9,
                "note_nlp_source_value": "anatomie",
                "start_char": 323,
            },
            {
                "assertion": "absent",
                "end_char": 386,
                "etat": None,
                "lexical_variant": "anomalie",
                "note_nlp_id": 10,
                "note_nlp_source_value": "sosy",
                "start_char": 378,
            },
        ],
        "note_id": "subfolder/doc-1",
        "context_var": "test",
        "note_text": "Le patient est admis pour des douleurs dans le bras droit, mais "
        "n'a pas de problème \n"
        "de locomotion. \n"
        "Historique d'AVC dans la famille. pourrait être un cas de "
        "rhume.\n"
        "NBNbWbWbNbWbNBNbNbWbWbNBNbWbNbNbWbNBNbWbNbNBWbWbNbNbNBWbNbWbNbWBNbNbWbNbNBNbWb"
        "WbNbWBNbNbWbNBNbWbWbNb\n"
        "Pourrait être un cas de rhume.\n"
        "Motif :\n"
        "Douleurs dans le bras droit.\n"
        "ANTÉCÉDENTS\n"
        "Le patient est déjà venu\n"
        "Pas d'anomalie détectée.\n",
    }


def test_read_in_worker(blank_nlp, tmpdir):
    input_dir = Path(__file__).parent.parent.resolve() / "resources" / "docs.jsonl"
    list(
        edsnlp.data.read_json(
            input_dir,
            converter="omop",
            span_attributes=["etat", "assertion"],
            read_in_worker=True,
        )
    )[0]


def test_read_to_json(blank_nlp, tmpdir):
    input_dir = Path(__file__).parent.parent.resolve() / "resources" / "docs.jsonl"
    output_dir = Path(tmpdir)
    doc = list(
        edsnlp.data.read_json(
            input_dir,
            converter="omop",
            span_attributes=["etat", "assertion"],
            doc_attributes=["context_var"],
        )
    )[0]
    assert_doc_read(doc)
    doc.ents[0]._.etat = "test"

    edsnlp.data.write_json(
        [doc],
        output_dir / "docs.jsonl",
        converter="omop",
        doc_attributes=["context_var"],
        span_attributes=["etat", "assertion"],
        span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
        lines=True,
    )

    with pytest.raises(FileExistsError):
        edsnlp.data.write_json(
            [doc],
            output_dir / "docs.jsonl",
            converter="omop",
            doc_attributes=["context_var"],
            span_attributes=["etat", "assertion"],
            span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
        )

    with open(output_dir / "docs.jsonl") as f:
        exported_obj = json.loads(f.readlines()[0])
    assert_doc_write(exported_obj)

    with pytest.raises(FileExistsError):
        edsnlp.data.write_json(
            [doc],
            output_dir,
            converter="omop",
            doc_attributes=["context_var"],
            span_attributes=["etat", "assertion"],
            span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
            lines=False,
        )

    edsnlp.data.write_json(
        [doc],
        output_dir,
        converter="omop",
        doc_attributes=["context_var"],
        span_attributes=["etat", "assertion"],
        span_getter=["ents", "sosy", "localisation", "anatomie", "pathologie"],
        lines=False,
        overwrite=True,
    )

    with open(output_dir / "subfolder" / "doc-1.json") as f:
        exported_obj = json.loads(f.read())
    assert_doc_write(exported_obj)
