"""Tests for rbsp.core.synonyms -- UnionFind and SynonymMap."""

from __future__ import annotations

import os
import sys
import tempfile
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from rbsp.core.synonyms import SynonymMap, UnionFind


# ---------------------------------------------------------------------------
# UnionFind
# ---------------------------------------------------------------------------

class TestUnionFindBasic(unittest.TestCase):
    """Basic UnionFind operations."""

    def test_find_self(self):
        uf = UnionFind()
        root = uf.find("a")
        self.assertEqual(root, "a")

    def test_contains_after_find(self):
        uf = UnionFind()
        self.assertNotIn("a", uf)
        uf.find("a")
        self.assertIn("a", uf)

    def test_contains_not_present(self):
        uf = UnionFind()
        self.assertNotIn("missing", uf)


class TestUnionFindUnion(unittest.TestCase):
    """Union and find with groups."""

    def test_union_two_elements(self):
        uf = UnionFind()
        uf.union("a", "b")
        self.assertEqual(uf.find("a"), uf.find("b"))

    def test_union_chain(self):
        uf = UnionFind()
        uf.union("a", "b")
        uf.union("b", "c")
        self.assertEqual(uf.find("a"), uf.find("c"))

    def test_union_same_element(self):
        uf = UnionFind()
        uf.union("a", "a")
        self.assertEqual(uf.find("a"), "a")

    def test_union_already_same_group(self):
        uf = UnionFind()
        uf.union("a", "b")
        uf.union("a", "b")
        self.assertEqual(uf.find("a"), uf.find("b"))


class TestUnionFindGroup(unittest.TestCase):
    """Group retrieval."""

    def test_single_element_group(self):
        uf = UnionFind()
        group = uf.group("x")
        self.assertEqual(group, {"x"})

    def test_merged_group(self):
        uf = UnionFind()
        uf.union("a", "b")
        uf.union("b", "c")
        group = uf.group("a")
        self.assertEqual(group, {"a", "b", "c"})

    def test_separate_groups(self):
        uf = UnionFind()
        uf.union("a", "b")
        uf.union("c", "d")
        self.assertEqual(uf.group("a"), {"a", "b"})
        self.assertEqual(uf.group("c"), {"c", "d"})


class TestUnionFindCanonical(unittest.TestCase):
    """Canonical form for groups."""

    def test_single_element_canonical(self):
        uf = UnionFind()
        uf._ensure("x")
        self.assertEqual(uf.canonical("x"), "x")

    def test_canonical_preserved_after_union(self):
        uf = UnionFind()
        uf._ensure("first")
        uf.union("first", "second")
        # The root's canonical is preserved -- it should be "first"
        root = uf.find("first")
        self.assertEqual(uf.canonical("first"), uf._canonical[root])

    def test_canonical_consistent_in_group(self):
        uf = UnionFind()
        uf.union("a", "b")
        uf.union("a", "c")
        # All members should share the same canonical
        ca = uf.canonical("a")
        cb = uf.canonical("b")
        cc = uf.canonical("c")
        self.assertEqual(ca, cb)
        self.assertEqual(cb, cc)


# ---------------------------------------------------------------------------
# SynonymMap
# ---------------------------------------------------------------------------

class TestSynonymMapAddEquivalent(unittest.TestCase):
    """add_equivalent creates an equivalence class."""

    def test_basic_equivalence(self):
        sm = SynonymMap()
        sm.add_equivalent(["usa", "united states", "america"])
        expanded = sm.expand("usa")
        self.assertEqual(expanded, {"usa", "united states", "america"})

    def test_canonical_is_first_term(self):
        sm = SynonymMap()
        sm.add_equivalent(["usa", "united states", "america"])
        replaced = sm.apply_replacements(["united states"])
        self.assertEqual(replaced, ["usa"])

    def test_single_term_is_noop(self):
        sm = SynonymMap()
        sm.add_equivalent(["only_one"])
        # Should not crash; only_one expands to just itself
        expanded = sm.expand("only_one")
        self.assertEqual(expanded, {"only_one"})

    def test_len_after_add_equivalent(self):
        sm = SynonymMap()
        sm.add_equivalent(["a", "b", "c"])
        self.assertEqual(len(sm), 1)

    def test_multiple_equivalence_groups(self):
        sm = SynonymMap()
        sm.add_equivalent(["big", "large"])
        sm.add_equivalent(["small", "tiny"])
        self.assertEqual(sm.expand("big"), {"big", "large"})
        self.assertEqual(sm.expand("small"), {"small", "tiny"})
        self.assertEqual(len(sm), 2)


class TestSynonymMapAddRule(unittest.TestCase):
    """add_rule with bidirectional and directional modes."""

    def test_bidirectional_rule(self):
        sm = SynonymMap()
        sm.add_rule(["color"], ["colour"], bidirectional=True)
        expanded = sm.expand("color")
        self.assertIn("colour", expanded)
        expanded2 = sm.expand("colour")
        self.assertIn("color", expanded2)

    def test_directional_rule(self):
        sm = SynonymMap()
        sm.add_rule(["big"], ["large", "huge"], bidirectional=False)
        expanded = sm.expand("big")
        self.assertIn("large", expanded)
        self.assertIn("huge", expanded)
        # "large" should NOT expand to "big" directionally
        expanded_large = sm.expand("large")
        self.assertNotIn("big", expanded_large)

    def test_rule_count(self):
        sm = SynonymMap()
        sm.add_rule(["a"], ["b"])
        sm.add_rule(["c"], ["d"], bidirectional=False)
        self.assertEqual(len(sm), 2)


class TestSynonymMapExpand(unittest.TestCase):
    """expand returns all synonyms for a term."""

    def test_unknown_term_returns_self(self):
        sm = SynonymMap()
        expanded = sm.expand("unknown")
        self.assertEqual(expanded, {"unknown"})

    def test_combined_directional_and_equivalence(self):
        sm = SynonymMap()
        sm.add_equivalent(["fast", "quick"])
        sm.add_rule(["fast"], ["speedy"], bidirectional=False)
        expanded = sm.expand("fast")
        self.assertIn("fast", expanded)
        self.assertIn("quick", expanded)
        self.assertIn("speedy", expanded)


class TestSynonymMapExpandQuery(unittest.TestCase):
    """expand_query returns per-token alternatives."""

    def test_basic_expansion(self):
        sm = SynonymMap()
        sm.add_equivalent(["big", "large"])
        result = sm.expand_query(["very", "big"])
        # "very" has no synonyms -> ["very"]
        self.assertEqual(result[0], ["very"])
        # "big" has "large" as synonym -> ["big", "large"]
        self.assertEqual(result[1][0], "big")  # original first
        self.assertIn("large", result[1])

    def test_no_synonyms(self):
        sm = SynonymMap()
        result = sm.expand_query(["hello", "world"])
        self.assertEqual(result, [["hello"], ["world"]])


class TestSynonymMapApplyReplacements(unittest.TestCase):
    """apply_replacements replaces tokens with canonical form."""

    def test_basic_replacement(self):
        sm = SynonymMap()
        sm.add_equivalent(["usa", "united states", "america"])
        result = sm.apply_replacements(["visit", "america"])
        self.assertEqual(result, ["visit", "usa"])

    def test_unknown_token_unchanged(self):
        sm = SynonymMap()
        sm.add_equivalent(["big", "large"])
        result = sm.apply_replacements(["unknown"])
        self.assertEqual(result, ["unknown"])

    def test_all_synonyms_map_to_canonical(self):
        sm = SynonymMap()
        sm.add_equivalent(["color", "colour"])
        result = sm.apply_replacements(["colour", "color"])
        self.assertEqual(result[0], result[1])
        self.assertEqual(result[0], "color")


class TestSynonymMapPersistence(unittest.TestCase):
    """save_to_file and load_from_file round-trip."""

    def test_round_trip_equivalence(self):
        sm = SynonymMap()
        sm.add_equivalent(["usa", "united states", "america"])
        sm.add_equivalent(["big", "large", "huge"])

        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".txt", delete=False, encoding="utf-8"
        ) as f:
            path = f.name

        try:
            sm.save_to_file(path)

            sm2 = SynonymMap()
            sm2.load_from_file(path)

            # Both should expand the same way
            self.assertEqual(sm.expand("usa"), sm2.expand("usa"))
            self.assertEqual(sm.expand("big"), sm2.expand("big"))
        finally:
            os.unlink(path)

    def test_round_trip_directional(self):
        sm = SynonymMap()
        sm.add_rule(["fast"], ["speedy", "quick"], bidirectional=False)

        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".txt", delete=False, encoding="utf-8"
        ) as f:
            path = f.name

        try:
            sm.save_to_file(path)

            sm2 = SynonymMap()
            sm2.load_from_file(path)

            expanded = sm2.expand("fast")
            self.assertIn("speedy", expanded)
            self.assertIn("quick", expanded)
        finally:
            os.unlink(path)

    def test_load_ignores_comments_and_blanks(self):
        content = "# This is a comment\n\nusa, united states\n"
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".txt", delete=False, encoding="utf-8"
        ) as f:
            f.write(content)
            path = f.name

        try:
            sm = SynonymMap()
            sm.load_from_file(path)
            expanded = sm.expand("usa")
            self.assertEqual(expanded, {"usa", "united states"})
        finally:
            os.unlink(path)

    def test_load_directional_format(self):
        content = "=> big, large, huge\n"
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".txt", delete=False, encoding="utf-8"
        ) as f:
            f.write(content)
            path = f.name

        try:
            sm = SynonymMap()
            sm.load_from_file(path)
            expanded = sm.expand("big")
            self.assertIn("large", expanded)
            self.assertIn("huge", expanded)
        finally:
            os.unlink(path)


class TestSynonymMapLen(unittest.TestCase):
    """__len__ counts rules."""

    def test_empty(self):
        sm = SynonymMap()
        self.assertEqual(len(sm), 0)

    def test_after_rules(self):
        sm = SynonymMap()
        sm.add_equivalent(["a", "b"])
        sm.add_rule(["c"], ["d"], bidirectional=False)
        self.assertEqual(len(sm), 2)

    def test_repr(self):
        sm = SynonymMap()
        sm.add_equivalent(["a", "b"])
        r = repr(sm)
        self.assertIn("SynonymMap", r)
        self.assertIn("rules=1", r)


if __name__ == "__main__":
    unittest.main()
