import json
import tempfile
import unittest
from pathlib import Path

import make_batches
import merge_labels


def rec(i, vol=1, ch=1, speaker=None, conf=None):
    return {"id": i, "volume": vol, "chapter": ch, "text": f"對白{i}", "before": "", "after": "", "prev_par": "",
            "next_par": "", "candidates": [], "speaker": speaker, "confidence": conf}


class Batches(unittest.TestCase):
    def test_neighbors_stay_inside_the_chapter_and_hide_rule_speaker(self):
        recs = [rec("a1"), rec("a2", speaker="asuna", conf="high"), rec("a3"), rec("b1", ch=2)]
        nb = make_batches.with_neighbors(recs)
        self.assertEqual([x["id"] for x in nb["a2"]["prev"]], ["a1"])
        self.assertEqual([x["id"] for x in nb["a2"]["next"]], ["a3"])
        self.assertEqual(nb["b1"], {"prev": [], "next": []})
        self.assertNotIn("rule_speaker", nb["a2"]["prev"][0])

    def test_pilot_sample_is_stratified_and_deterministic(self):
        recs = [rec(f"u{i}") for i in range(50)] + [rec(f"m{i}", speaker="kirito", conf="medium") for i in range(50)] \
            + [rec(f"h{i}", speaker="asuna", conf="high") for i in range(50)]
        a = make_batches.pilot_sample(recs, 20, 1)
        b = make_batches.pilot_sample(recs, 20, 1)
        self.assertEqual([r["id"] for r in a], [r["id"] for r in b])
        kinds = [r["id"][0] for r in a]
        self.assertEqual((kinds.count("u"), kinds.count("m"), kinds.count("h")), (10, 5, 5))


class Check(unittest.TestCase):
    good = {"id": "x", "speaker": "asuna", "confidence": "high", "function": "tease", "emotion": "anger", "evidence": "前段"}

    def test_valid(self):
        self.assertEqual(merge_labels.check(self.good), [])

    def test_unknown_function_rejected(self):
        self.assertTrue(merge_labels.check({**self.good, "function": "flirt"}))

    def test_unknown_emotion_rejected(self):
        self.assertTrue(merge_labels.check({**self.good, "emotion": "rage"}))

    def test_extra_or_missing_fields_rejected(self):
        self.assertTrue(merge_labels.check({**self.good, "note": "hi"}))
        self.assertTrue(merge_labels.check({k: v for k, v in self.good.items() if k != "evidence"}))

    def test_bad_confidence_rejected(self):
        self.assertTrue(merge_labels.check({**self.good, "confidence": "sure"}))

    def test_none_speaker_is_valid_but_never_final(self):
        self.assertEqual(merge_labels.check({**self.good, "speaker": "none", "function": "other", "emotion": "neutral"}), [])

    def test_whitelist_matches_prompt(self):
        prompt = (Path(__file__).with_name("label_prompt.md")).read_text(encoding="utf-8")
        self.assertIn("tags.json", prompt)
        self.assertIn("other", merge_labels.TAGS["functions"])


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