import json
import unittest

from copienator.latex_json import (
    normalize_literal_newlines,
    repair_decoded_latex_controls,
    repair_latex_json,
)
from copienator.commands.gemini_for_enonce import ExamQuestions, ExamSolutions, ExamContext, GroupRubrics


class LatexJsonTests(unittest.TestCase):
    def test_literal_newlines_preserve_tex_commands_and_row_breaks(self):
        content = r"Paragraph\n\n Next $\neq \notin \newline$ \\ next \\\n row"
        expected = "Paragraph\n\n Next " + r"$\neq \notin \newline$ \\ next \\" + "\n row"
        self.assertEqual(normalize_literal_newlines(content), expected)
        self.assertEqual(normalize_literal_newlines(expected), expected)
        raw = json.dumps({'solutions': [{'label': '1)', 'solution_content': content}]})
        parsed = ExamSolutions.model_validate_json(repair_latex_json(raw))
        self.assertEqual(parsed.solutions[0].solution_content, expected)

    def test_mixed_correct_and_broken_escapes_in_all_extraction_fields(self):
        cases = [
            (ExamQuestions, 'questions', 'question_content', {'label': '1)'}),
            (ExamSolutions, 'solutions', 'solution_content', {'label': '1)'}),
            (ExamContext, 'contexts', 'context_content',
             {'target_question_label': '1)', 'last_question_label': '2)'}),
            (GroupRubrics, 'rubrics', 'rubric_content', {'label': '1)'})]
        raw_content = r'$\circ \rightarrow \frac{1}{2} \not= \theta \begin{x} \unknown \underline{x}$\nSuite : $\\sqrt{x}$'
        expected = r'$\circ \rightarrow \frac{1}{2} \not= \theta \begin{x} \unknown \underline{x}$' + '\n' + r'Suite : $\sqrt{x}$'
        for model, collection, field, labels in cases:
            with self.subTest(field=field):
                raw = json.dumps({collection: [{**labels, field: 'PLACEHOLDER'}]}).replace('PLACEHOLDER', raw_content)
                repaired = repair_latex_json(raw)
                parsed = model.model_validate_json(repaired)
                self.assertEqual(getattr(getattr(parsed, collection)[0], field), expected)
                self.assertEqual(repair_latex_json(repaired), repaired)

    def test_valid_json_is_unchanged(self):
        content = 'Énoncé "cité"\nNouvelle ligne\ntexte\nnext\ni \ne \nu \tindent\r\n' + r'$\not= \frac{1}{2} \\ \begin{array}$'
        raw = json.dumps({'question_content': content, 'label': r'\not'})
        self.assertEqual(repair_latex_json(raw), raw)
        self.assertEqual(repair_latex_json(r'{"rubric_content":"\u00e9 \/ \"ok\""}'),
                         r'{"rubric_content":"\u00e9 \/ \"ok\""}')

    def test_only_content_values_are_repaired(self):
        raw = r'{"label":"\not", "other":"\frac", "rubric_content":"\frac"}'
        self.assertEqual(repair_latex_json(raw),
                         r'{"label":"\not", "other":"\frac", "rubric_content":"\\frac"}')

    def test_feedback_text_can_be_repaired_before_json_decoding(self):
        raw = (
            r'{"text":"$\forall x, f \big(x\big) \neq x \to \text{false}$",'
            r'"other":"\forall"}'
        )
        repaired = repair_latex_json(
            raw,
            additional_content_fields={"text"},
        )
        parsed = json.loads(repaired)
        self.assertEqual(
            parsed["text"],
            r"$\forall x, f \big(x\big) \neq x \to \text{false}$",
        )
        self.assertEqual(parsed["other"], "\x0corall")

        valid_prose = json.dumps({"text": "First line\ntext outside math"})
        self.assertEqual(
            repair_latex_json(
                valid_prose,
                additional_content_fields={"text"},
            ),
            valid_prose,
        )

    def test_decoded_control_repair_is_limited_to_known_commands_in_math(self):
        corrupted = (
            "$"
            + "\x0corall x, f "
            + "\x08ig(x\x08ig) "
            + "\neq x "
            + "\to "
            + "\text{false}$"
        )
        self.assertEqual(
            repair_decoded_latex_controls(corrupted),
            r"$\forall x, f \big(x\big) \neq x \to \text{false}$",
        )
        ordinary_whitespace = "First paragraph\nnext line\tindented"
        self.assertEqual(
            repair_decoded_latex_controls(ordinary_whitespace),
            ordinary_whitespace,
        )

    def test_does_not_repair_structural_errors(self):
        with self.assertRaises(json.JSONDecodeError):
            json.loads(repair_latex_json(r'{"question_content":"\frac", broken}'))