import json
import tempfile
import unittest
from pathlib import Path

from textbook.build_library import build


def write_jsonl(path, records):
    path.write_text('\n'.join(json.dumps(r) for r in records) + '\n', encoding='utf-8')


class BuildLibraryTests(unittest.TestCase):
    def test_build_includes_only_solved_theorems_and_hoists_shared_imports(self):
        with tempfile.TemporaryDirectory() as directory:
            directory = Path(directory)
            dataset = directory / 'dataset.jsonl'
            results = directory / 'results.jsonl'
            write_jsonl(dataset, [
                {"id": "solved-one", "statement": "True", "prelude": "import Mathlib",
                 "source_label": "Theorem 1", "natural_language": "trivial"},
                {"id": "unsolved-one", "statement": "False", "prelude": "import Mathlib",
                 "source_label": "Theorem 2", "natural_language": "not actually true"},
            ])
            write_jsonl(results, [
                {"id": "solved-one", "success": True,
                 "attempts": [{"proof": "trivial", "success": True, "diagnostics": ""}]},
                {"id": "unsolved-one", "success": False, "attempts": []},
            ])

            text, solved, skipped = build(dataset, results)

            self.assertEqual(solved, ['solved-one'])
            self.assertEqual(skipped, ['unsolved-one'])
            self.assertIn('import Mathlib', text)
            self.assertEqual(text.count('import Mathlib'), 1)  # hoisted once, not per-theorem
            self.assertIn('theorem solved_one : True := by', text)
            self.assertNotIn('unsolved_one', text)

    def test_build_keeps_non_import_prelude_attached_per_theorem(self):
        with tempfile.TemporaryDirectory() as directory:
            directory = Path(directory)
            dataset = directory / 'dataset.jsonl'
            results = directory / 'results.jsonl'
            write_jsonl(dataset, [
                {"id": "custom-def", "statement": "myPred 1", "prelude": "def myPred (n : Nat) := True",
                 "source_label": "Theorem 3", "natural_language": ""},
            ])
            write_jsonl(results, [
                {"id": "custom-def", "success": True,
                 "attempts": [{"proof": "trivial", "success": True, "diagnostics": ""}]},
            ])

            text, solved, skipped = build(dataset, results)

            self.assertIn('def myPred (n : Nat) := True', text)
            self.assertEqual(solved, ['custom-def'])


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