import json
import unittest

from textbook.formalize import extract_json_block, validate

VALID_REPLY = """Here is the result:
```json
{
  "definitions": [
    {"label": "Definition 1.6.3", "concept": "prime number", "mathlib_ref": "Nat.Prime",
     "lean_code": null, "notes": ""}
  ],
  "theorems": [
    {"id": "goodman-1-6-6", "source_label": "Theorem 1.6.6", "kind": "theorem",
     "statement": "\\u2200 n : \\u2115, \\u2203 p, p > n \\u2227 p.Prime",
     "prelude": "import Mathlib", "natural_language": "There are infinitely many primes.",
     "requires": ["Nat.Prime"], "include_in_bench": true, "exclude_reason": null}
  ]
}
```
"""


class FormalizeParsingTests(unittest.TestCase):
    def test_extract_json_block_from_fenced_reply(self):
        data = extract_json_block(VALID_REPLY)
        self.assertEqual(set(data), {"definitions", "theorems"})

    def test_extract_json_block_without_fence_falls_back_to_whole_text(self):
        payload = '{"theorems": [], "definitions": []}'
        self.assertEqual(extract_json_block(payload), json.loads(payload))

    def test_validate_accepts_well_formed_response(self):
        data = extract_json_block(VALID_REPLY)
        theorems, definitions = validate(data)
        self.assertEqual(len(theorems), 1)
        self.assertEqual(len(definitions), 1)

    def test_validate_rejects_missing_field(self):
        data = {"theorems": [{"id": "x", "source_label": "Theorem 1", "statement": "True"}],
                 "definitions": []}
        with self.assertRaises(ValueError):
            validate(data)

    def test_validate_rejects_duplicate_ids(self):
        entry = {"id": "dup", "source_label": "Theorem 1", "kind": "theorem", "statement": "True",
                  "prelude": "import Mathlib", "include_in_bench": True}
        data = {"theorems": [entry, dict(entry)], "definitions": []}
        with self.assertRaises(ValueError):
            validate(data)

    def test_validate_rejects_non_object_response(self):
        with self.assertRaises(ValueError):
            validate([])


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