import json
import tempfile
import unittest
from pathlib import Path

from llm.training_data import load_emotion_ids, load_examples, sha256_file


class TrainingDataTests(unittest.TestCase):
    def test_loads_valid_examples_and_honors_row_limit(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "data.jsonl"
            path.write_text(
                "\n".join(
                    json.dumps(row)
                    for row in (
                        {"user": "Hello", "emotion": "happy", "reply": "Hi!"},
                        {"user": "Wait", "emotion": "calm", "reply": "I can wait."},
                    )
                )
                + "\n"
            )
            rows = load_examples(path, {"happy", "calm"}, max_rows=1)
            self.assertEqual(rows, [{"user": "Hello", "emotion": "happy", "reply": "Hi!"}])
            self.assertEqual(len(sha256_file(path)), 64)

    def test_rejects_unknown_emotion(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "data.jsonl"
            path.write_text('{"user":"Hello","emotion":"unsafe","reply":"Hi"}\n')
            with self.assertRaisesRegex(ValueError, "unknown emotion"):
                load_examples(path, {"happy"})

    def test_loads_unique_emotion_ids(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "emotions.json"
            path.write_text('[{"id":"happy"},{"id":"calm"}]')
            self.assertEqual(load_emotion_ids(path), ["happy", "calm"])
            path.write_text('[{"id":"happy"},{"id":"happy"}]')
            with self.assertRaisesRegex(ValueError, "must be unique"):
                load_emotion_ids(path)


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