from __future__ import annotations

import unittest

from br_sovereign_llm.corpus import Split, SplitConfig, assign_split, split_bucket


class SplitAssignmentTests(unittest.TestCase):
    def test_known_key_has_stable_bucket(self) -> None:
        self.assertEqual(split_bucket("group:lei-9610"), 6756)
        self.assertEqual(assign_split("group:lei-9610"), Split.TRAIN)

    def test_same_leakage_group_always_has_same_split(self) -> None:
        assignments = {assign_split("content-sha256:abc") for _ in range(20)}

        self.assertEqual(assignments, {assign_split("content-sha256:abc")})

    def test_all_partitions_are_reachable(self) -> None:
        config = SplitConfig(
            train_permyriad=8000,
            validation_permyriad=1000,
            test_permyriad=1000,
            seed="coverage-test",
        )
        observed = {
            assign_split(f"group-{index}", config=config) for index in range(500)
        }

        self.assertEqual(observed, {Split.TRAIN, Split.VALIDATION, Split.TEST})

    def test_seed_is_part_of_assignment(self) -> None:
        first = SplitConfig(
            train_permyriad=5000,
            validation_permyriad=2500,
            test_permyriad=2500,
            seed="first-seed",
        )
        second = SplitConfig(
            train_permyriad=5000,
            validation_permyriad=2500,
            test_permyriad=2500,
            seed="second-seed",
        )

        self.assertTrue(
            any(
                assign_split(f"group-{index}", config=first)
                != assign_split(f"group-{index}", config=second)
                for index in range(100)
            )
        )

    def test_split_proportions_must_sum_to_permyriad(self) -> None:
        with self.assertRaises(ValueError):
            SplitConfig(
                train_permyriad=9800,
                validation_permyriad=100,
                test_permyriad=99,
            )

    def test_boolean_proportion_is_rejected(self) -> None:
        with self.assertRaises(ValueError):
            SplitConfig(
                train_permyriad=True,
                validation_permyriad=9999,
                test_permyriad=0,
            )


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