"""Unit tests for individual functions in reconciliation.py."""

import pytest
from collections_extended import RangeMap
from pymultirole_plugins.v1.schema import Annotation, Document, Term

from pyprocessors_reconciliation.reconciliation import (
    ReconciliationParameters,
    ReconciliationProcessor,
    ReconciliationType,
    WrappedTerm,
    by_label,
    by_lexicon,
    by_lexicon_or_label,
    group_annotations,
    has_knowledge,
    is_kill_or_white_label,
    is_suspicious,
    is_whitelist,
    labels_are_compatible,
    left_longest_match,
    mark_whitelisted,
    natural_order,
    one_match,
    person_varnames,
)

# ── helpers ───────────────────────────────────────────────────────────────────


def ann(start, end, label="ORG", text=None, terms=None):
    return Annotation(start=start, end=end, labelName=label, text=text, terms=terms)


def term(identifier, lexicon, preferred=None, properties=None):
    return Term(identifier=identifier, lexicon=lexicon, preferredForm=preferred or identifier, properties=properties)


def make_range_map(*annotations):
    rm = RangeMap()
    for a in annotations:
        rm[a.start : a.end] = a
    return rm


# ── WrappedTerm ───────────────────────────────────────────────────────────────


class TestWrappedTerm:
    def test_equal_same_id_and_status(self):
        t = term("id1", "lex", properties={"status": "linked"})
        assert WrappedTerm(t) == WrappedTerm(t)

    def test_not_equal_different_status(self):
        t1 = term("id1", "lex", properties={"status": "linked"})
        t2 = term("id1", "lex", properties={"status": "white"})
        assert WrappedTerm(t1) != WrappedTerm(t2)

    def test_equal_no_properties(self):
        t = term("id1", "lex")
        assert WrappedTerm(t) == WrappedTerm(t)

    def test_hash_consistent_for_equal(self):
        t = term("id1", "lex")
        assert hash(WrappedTerm(t)) == hash(WrappedTerm(t))

    def test_set_deduplication(self):
        t = term("id1", "lex")
        s = {WrappedTerm(t), WrappedTerm(t)}
        assert len(s) == 1

    def test_status_defaults_to_empty_without_properties(self):
        t = term("id1", "lex")
        wt = WrappedTerm(t)
        assert wt.status == ""


# ── has_knowledge ─────────────────────────────────────────────────────────────


class TestHasKnowledge:
    def test_with_terms(self):
        assert has_knowledge(ann(0, 5, terms=[term("id1", "lex")]))

    def test_none_terms(self):
        assert not has_knowledge(ann(0, 5, terms=None))

    def test_empty_terms(self):
        assert not has_knowledge(ann(0, 5, terms=[]))


# ── is_suspicious ─────────────────────────────────────────────────────────────


class TestIsSuspicious:
    def test_all_lowercase(self):
        assert is_suspicious(ann(0, 7, text="company")) is True

    def test_uppercase_first_word(self):
        assert is_suspicious(ann(0, 7, text="Company")) is False

    def test_multi_word_all_lower(self):
        assert is_suspicious(ann(0, 10, text="some thing")) is True

    def test_multi_word_one_upper(self):
        assert is_suspicious(ann(0, 11, text="some Thing")) is False

    def test_no_text_returns_false(self):
        assert is_suspicious(ann(0, 5)) is False

    def test_number_string(self):
        assert is_suspicious(ann(0, 3, text="123")) is True


# ── is_whitelist ──────────────────────────────────────────────────────────────


class TestIsWhitelist:
    def test_no_terms(self):
        assert is_whitelist(ann(0, 5)) is False

    def test_term_with_lowercase_w_status(self):
        a = ann(0, 5, terms=[term("id1", "lex", properties={"status": "w"})])
        assert is_whitelist(a) is True

    def test_term_with_uppercase_W_status(self):
        a = ann(0, 5, terms=[term("id1", "lex", properties={"status": "W"})])
        assert is_whitelist(a) is True

    def test_term_with_non_w_status(self):
        a = ann(0, 5, terms=[term("id1", "lex", properties={"status": "linked"})])
        assert is_whitelist(a) is False

    def test_term_with_no_properties(self):
        assert is_whitelist(ann(0, 5, terms=[term("id1", "lex")])) is False


# ── person_varnames ───────────────────────────────────────────────────────────


class TestPersonVarnames:
    def test_non_person_label_returns_none(self):
        assert person_varnames(ann(0, 7, label="ORG", text="Company"), "Company") is None

    def test_single_word_name(self):
        a = ann(0, 4, label="person", text="Jean")
        assert person_varnames(a, "Jean") == ["Jean"]

    def test_two_word_name_contains_all_variants(self):
        a = ann(0, 11, label="person", text="Jean Dupont")
        result = person_varnames(a, "Jean Dupont")
        assert "Jean Dupont" in result
        assert "Dupont" in result
        assert "Jean" in result

    def test_three_word_name_variants(self):
        a = ann(0, 16, label="person", text="Jean Paul Dupont")
        result = person_varnames(a, "Jean Paul Dupont")
        assert "Jean Paul Dupont" in result
        assert "Paul Dupont" in result
        assert "Dupont" in result
        assert "Jean" in result
        assert "Jean Paul" in result

    def test_uses_text_slice_when_annotation_has_no_text(self):
        text = "Hello Jean Dupont end"
        a = Annotation(start=6, end=17, labelName="person")
        result = person_varnames(a, text)
        assert result is not None
        assert "Jean Dupont" in result


# ── key functions ─────────────────────────────────────────────────────────────


class TestKeyFunctions:
    def test_by_lexicon_with_terms(self):
        assert by_lexicon(ann(0, 5, terms=[term("id1", "wikidata")])) == "wikidata"

    def test_by_lexicon_no_terms(self):
        assert by_lexicon(ann(0, 5)) == ""

    def test_by_lexicon_or_label_with_terms(self):
        a = ann(0, 5, label="ORG", terms=[term("id1", "wikidata")])
        assert by_lexicon_or_label(a) == "wikidata"

    def test_by_lexicon_or_label_no_terms(self):
        assert by_lexicon_or_label(ann(0, 5, label="ORG")) == "ORG"

    def test_by_label(self):
        assert by_label(ann(0, 5, label="PER")) == "PER"

    def test_left_longest_match(self):
        assert left_longest_match(ann(2, 7)) == (5, -2)

    def test_natural_order(self):
        assert natural_order(ann(3, 8)) == (-3, 5)


# ── one_match ─────────────────────────────────────────────────────────────────


class TestOneMatch:
    def test_none_matches_returns_false_none(self):
        perfect, match = one_match(ann(0, 5), None)
        assert perfect is False
        assert match is None

    def test_perfect_match(self):
        a = ann(0, 5, label="ORG")
        b = ann(0, 5, label="ORG")
        perfect, match = one_match(a, make_range_map(b))
        assert perfect is True
        assert match is b

    def test_imperfect_match_larger_span(self):
        a = ann(2, 5, label="ORG")
        b = ann(0, 7, label="ORG")
        perfect, match = one_match(a, make_range_map(b))
        assert perfect is False
        assert match is b


# ── label compatibility ───────────────────────────────────────────────────────


class TestLabelCompatibility:
    @pytest.fixture
    def params(self):
        return ReconciliationParameters(kill_label="kill", white_label="white")

    def test_same_label_compatible(self, params):
        assert labels_are_compatible("ORG", "ORG", params) is True

    def test_kill_label_is_compatible(self, params):
        assert labels_are_compatible("kill", "ORG", params) is True

    def test_white_label_is_compatible(self, params):
        assert labels_are_compatible("ORG", "white", params) is True

    def test_unrelated_labels_incompatible(self, params):
        assert labels_are_compatible("ORG", "PER", params) is False

    def test_is_kill_label(self, params):
        assert is_kill_or_white_label("kill", params) is True

    def test_is_white_label(self, params):
        assert is_kill_or_white_label("white", params) is True

    def test_unrelated_label_is_neither(self, params):
        assert is_kill_or_white_label("ORG", params) is False


# ── ReconciliationParameters defaults ─────────────────────────────────────────


class TestReconciliationParametersDefaults:
    def test_defaults(self):
        p = ReconciliationParameters()
        assert p.type == ReconciliationType.linker
        assert p.kill_label is None
        assert p.white_label is None
        assert p.whitelisted_lexicons is None
        assert p.person_label is None
        assert p.remove_suspicious is True
        assert p.resolve_lastnames is False


# ── group_annotations ─────────────────────────────────────────────────────────


class TestGroupAnnotations:
    def test_non_overlapping_annotations_in_same_group(self):
        t = term("id1", "lex")
        a1 = ann(0, 5, terms=[t])
        a2 = ann(10, 15, terms=[t])
        groups = group_annotations([a1, a2], by_lexicon)
        assert len(groups["lex"]) == 2

    def test_same_span_different_terms_merged(self):
        t1 = term("id1", "lex")
        t2 = term("id2", "lex")
        a1 = ann(0, 5, terms=[t1])
        a2 = ann(0, 5, terms=[t2])
        groups = group_annotations([a1, a2], by_lexicon)
        merged = groups["lex"][0]
        assert len(merged.terms) == 2

    def test_same_span_duplicate_term_deduplicated(self):
        t = term("id1", "lex")
        a1 = ann(0, 5, terms=[t])
        a2 = ann(0, 5, terms=[t])
        groups = group_annotations([a1, a2], by_lexicon)
        merged = groups["lex"][0]
        assert len(merged.terms) == 1

    def test_empty_input(self):
        groups = group_annotations([], by_lexicon)
        assert len(groups) == 0

    def test_annotations_split_into_groups_by_lexicon(self):
        a1 = ann(0, 5, terms=[term("id1", "wiki")])
        a2 = ann(0, 5, terms=[term("id2", "kb")])
        groups = group_annotations([a1, a2], by_lexicon)
        assert "wiki" in groups
        assert "kb" in groups


# ── mark_whitelisted ──────────────────────────────────────────────────────────


class TestMarkWhitelisted:
    def test_passthrough_when_no_config(self):
        t = term("id1", "lex")
        annotations = [ann(0, 5, terms=[t]), ann(10, 15)]
        result = mark_whitelisted(annotations, ReconciliationParameters())
        assert result == annotations

    def test_white_label_clears_terms(self):
        t = term("id1", "lex")
        a = Annotation(start=0, end=5, labelName="white", terms=[t])
        params = ReconciliationParameters(white_label="white")
        result = mark_whitelisted([a], params)
        assert len(result) == 1
        assert result[0].terms is None

    def test_whitelisted_lexicon_adds_term_free_duplicate(self):
        t = term("id1", "mylex")
        a = ann(0, 5, label="ORG", terms=[t])
        params = ReconciliationParameters(whitelisted_lexicons=["mylex"])
        result = mark_whitelisted([a], params)
        assert len(result) == 2
        assert result[0].terms is not None  # original kept
        assert result[1].terms is None  # copy without terms

    def test_annotation_without_terms_always_passed_through(self):
        a = ann(0, 5)
        params = ReconciliationParameters(white_label="white", whitelisted_lexicons=["mylex"])
        result = mark_whitelisted([a], params)
        assert len(result) == 1


# ── process: edge cases ───────────────────────────────────────────────────────


class TestProcessEdgeCases:
    @pytest.fixture
    def processor(self):
        return ReconciliationProcessor()

    def test_document_without_annotations_returned_unchanged(self, processor):
        doc = Document(text="Hello world")
        result = processor.process([doc], ReconciliationParameters())
        assert result[0].annotations is None

    def test_sentence_annotations_stripped_before_reconciliation(self, processor):
        model_ann = Annotation(start=0, end=5, labelName="ORG", text="Apple")
        sentence = Annotation(start=0, end=11, labelName="sentence", text="Apple world")
        doc = Document(text="Apple world", annotations=[model_ann, sentence])
        result = processor.process([doc], ReconciliationParameters())
        labels = [a.labelName for a in result[0].annotations]
        assert "sentence" not in labels
        assert "ORG" in labels

    def test_suspicious_annotation_removed_by_default(self, processor):
        a = Annotation(start=0, end=7, labelName="ORG", text="company")
        doc = Document(text="company", annotations=[a])
        result = processor.process([doc], ReconciliationParameters(remove_suspicious=True))
        assert result[0].annotations == []

    def test_suspicious_annotation_kept_when_flag_disabled(self, processor):
        a = Annotation(start=0, end=7, labelName="ORG", text="company")
        doc = Document(text="company", annotations=[a])
        result = processor.process([doc], ReconciliationParameters(remove_suspicious=False))
        assert len(result[0].annotations) == 1

    def test_kill_label_removes_matching_model_annotation(self, processor):
        model_ann = Annotation(start=0, end=5, labelName="ORG", text="Apple")
        killer = Annotation(start=0, end=5, labelName="kill", text="Apple", terms=[term("id1", "blocklist")])
        doc = Document(text="Apple falls", annotations=[model_ann, killer])
        result = processor.process([doc], ReconciliationParameters(kill_label="kill"))
        labels = [a.labelName for a in result[0].annotations]
        assert "ORG" not in labels

    def test_kb_annotation_enriches_matching_model_annotation(self, processor):
        kb = Annotation(start=0, end=5, labelName="ORG", text="Apple", terms=[term("Q312", "wikidata", "Apple Inc.")])
        model_ann = Annotation(start=0, end=5, labelName="ORG", text="Apple")
        doc = Document(text="Apple rules", annotations=[model_ann, kb])
        result = processor.process([doc], ReconciliationParameters())
        enriched = result[0].annotations[0]
        assert enriched.terms is not None
        assert any(t.lexicon == "wikidata" for t in enriched.terms)

    def test_multiple_documents_processed_independently(self, processor):
        doc1 = Document(text="Apple", annotations=[Annotation(start=0, end=5, labelName="ORG", text="Apple")])
        doc2 = Document(text="Hello world")
        result = processor.process([doc1, doc2], ReconciliationParameters(remove_suspicious=False))
        assert len(result[0].annotations) == 1
        assert result[1].annotations is None
