import unittest

from app.ml.dataset import build_dataset, label_name, LABEL_NAMES


class DatasetBuilderTests(unittest.TestCase):
    def test_labels_are_the_documented_four(self):
        self.assertEqual(LABEL_NAMES, ["NORMAL", "DISPOSABLE", "ROLE", "SPAM_TRAP"])
        self.assertEqual(label_name(0), "NORMAL")
        self.assertEqual(label_name(1), "DISPOSABLE")
        self.assertEqual(label_name(2), "ROLE")
        self.assertEqual(label_name(3), "SPAM_TRAP")

    def test_blocklist_domains_become_disposable(self):
        rows = build_dataset(
            blocklist={"tempmail.example", "10minutemail.example"},
            free={"gmail.com"},
            role_prefixes={"admin"},
            normal_domains={"acme-corp.test"},
            normal_locals={"john.smith", "maria.garcia"},
            per_class_limit=50,
        )
        by_label = {}
        for r in rows:
            by_label.setdefault(r["label"], []).append(r["text"])
        self.assertIn(1, by_label, "blocklist domains must yield DISPOSABLE rows")
        self.assertTrue(any("tempmail.example" in t for t in by_label[1]))

    def test_role_prefixes_become_role_label(self):
        rows = build_dataset(
            blocklist=set(),
            free=set(),
            role_prefixes={"admin", "info"},
            normal_domains={"acme-corp.test"},
            normal_locals={"john.smith"},
            per_class_limit=50,
        )
        role_rows = [r for r in rows if r["label"] == 2]
        self.assertTrue(role_rows, "role prefixes must yield ROLE rows")
        self.assertTrue(any(r["text"].startswith("admin") for r in role_rows))

    def test_normal_class_never_contains_blocklist_domains(self):
        rows = build_dataset(
            blocklist={"tempmail.example"},
            free=set(),
            role_prefixes=set(),
            normal_domains={"acme-corp.test", "northwind.test"},
            normal_locals={"john.smith", "maria.garcia", "kenji.tanaka"},
            per_class_limit=50,
        )
        normal_texts = [r["text"] for r in rows if r["label"] == 0]
        self.assertTrue(normal_texts, "NORMAL class must be populated")
        self.assertTrue(all("tempmail.example" not in t for t in normal_texts))

    def test_per_class_limit_is_respected(self):
        rows = build_dataset(
            blocklist={f"d{i}.example" for i in range(100)},
            free=set(),
            role_prefixes=set(),
            normal_domains={"acme.test"},
            normal_locals={"john.smith"},
            per_class_limit=25,
        )
        counts = {}
        for r in rows:
            counts[r["label"]] = counts.get(r["label"], 0) + 1
        self.assertTrue(all(c <= 25 for c in counts.values()), counts)

    def test_text_format_is_local_sep_domain(self):
        rows = build_dataset(
            blocklist=set(),
            free=set(),
            role_prefixes={"admin"},
            normal_domains={"acme.test"},
            normal_locals={"john.smith"},
            per_class_limit=5,
        )
        self.assertTrue(rows)
        for r in rows:
            self.assertIn(" [SEP] ", r["text"])

    def test_deterministic_with_same_seed(self):
        kwargs = dict(
            blocklist={"tempmail.example"},
            free={"gmail.com"},
            role_prefixes={"admin"},
            normal_domains={"acme.test"},
            normal_locals={"john.smith"},
            per_class_limit=20,
            seed=7,
        )
        a = build_dataset(**kwargs)
        b = build_dataset(**kwargs)
        self.assertEqual(a, b, "same seed must produce identical rows")


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