import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch

from hessian_recipes.german_data import align_german, prepare


def wrapper(language, question, answer):
    return {"prompt": "Below is an instruction that describes a task. Write a response that appropriately "
            f"completes the request in {language}. Please answer in {language}.\n\n"
            f"### Instruction:\n{question}\n\n### Response:", "chosen": answer, "reject": "I don't know."}


class GermanDataTests(unittest.TestCase):
    def setUp(self):
        self.source = {"question": "What is 6 times 7?", "answer": "Multiply 6 * 7 = 42.\n#### 42"}
        self.english = wrapper("English", self.source["question"], self.source["answer"])
        self.german = wrapper("German", "Was ist 6 mal 7?", "Wir multiplizieren 6 mit 7 und erhalten 42.")

    def test_exact_source_anchor_and_matching_numeric_target(self):
        rows, coverage = align_german([self.english, self.german], [self.source], [])
        self.assertEqual(coverage["accepted"], 1)
        self.assertEqual(rows[0]["source_id"], "gsm8k/train/00000")
        self.assertEqual(rows[0]["completion"], "<think>Wir multiplizieren 6 mit 7 und erhalten 42.</think>\n\n<answer>42</answer>")
        self.assertNotIn("42", rows[0]["prompt"])

    def test_rejects_mismatched_rationale_even_when_question_matches(self):
        self.english["chosen"] = "A different answer #### 42"
        rows, coverage = align_german([self.english, self.german], [self.source], [])
        self.assertEqual(rows, [])
        self.assertEqual(coverage["rejected"]["no_exact_unique_english_train_anchor"], 1)

    def test_rejects_test_source_and_does_not_guess_or_fabricate_answer(self):
        rows, _ = align_german([self.english, self.german], [self.source], [self.source])
        self.assertEqual(rows, [])
        self.german["chosen"] = "Das Ergebnis ist 41."
        rows, coverage = align_german([self.english, self.german], [self.source], [])
        self.assertEqual(rows, [])
        self.assertEqual(coverage["rejected"]["final_number_disagrees_with_canonical_gold"], 1)

    def test_duplicate_ids_and_ambiguous_numeric_formats_rejected(self):
        rows, coverage = align_german([self.english, self.german, self.german], [self.source], [])
        self.assertEqual(len(rows), 1)
        self.assertEqual(coverage["rejected"]["duplicate_german_source_id"], 1)
        self.german["chosen"] = "Das Ergebnis ist 42,0."
        rows, coverage = align_german([self.english, self.german], [self.source], [])
        self.assertEqual(rows, [])
        self.assertEqual(coverage["rejected"]["ambiguous_or_missing_final_number"], 1)

    def test_malformed_row_resets_anchor(self):
        rows, _ = align_german([self.english, {"prompt": "broken"}, self.german], [self.source], [])
        self.assertEqual(rows, [])

    def test_untranslated_source_and_audited_reasoning_errors_rejected(self):
        untranslated = wrapper("German", self.source["question"], "Wir erhalten 42.")
        rows, coverage = align_german([self.english, untranslated], [self.source], [])
        self.assertEqual(rows, [])
        self.assertEqual(coverage["rejected"]["untranslated_source_question"], 1)
        with patch("hessian_recipes.german_data.AUDIT_EXCLUSIONS", {"gsm8k/train/00000": "reviewed error"}):
            rows, coverage = align_german([self.english, self.german], [self.source], [])
        self.assertEqual(rows, [])
        self.assertEqual(coverage["rejected"]["audited_incorrect_intermediate_reasoning"], 1)

    def test_prepare_splits_source_ids_and_hashes_written_files(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            sources = root / "sources"
            sources.mkdir()
            canonical, translated = [], []
            for index in range(4):
                row = {"question": f"What is {index} plus 1?", "answer": f"We add 1.\n#### {index + 1}"}
                canonical.append(row)
                translated.extend([wrapper("English", row["question"], row["answer"]),
                                   wrapper("German", f"Was ist {index} plus 1?", f"Wir addieren und erhalten {index + 1}.")])
            for name, rows in {"gsm8k-train.jsonl": canonical, "gsm8k-test.jsonl": [], "mathoctopus.json": translated}.items():
                (sources / name).write_text("\n".join(json.dumps(row) for row in rows))
            external = [{"id": "external", "source_id": "mgsm/de/test/000", "prompt": "Andere Frage?", "answer": "9", "scorer": "final_numeric", "category": "German math"}]
            with patch("hessian_recipes.german_data.fetch_sources", return_value={}), patch("hessian_recipes.german_data.mgsm_examples", return_value=external):
                report = prepare(root / "data", 2, 1, sources, download=False)
            train = report["files"]["train.jsonl"]["source_ids"]
            dev = report["files"]["dev.jsonl"]["source_ids"]
            self.assertFalse(set(train) & set(dev))
            self.assertEqual(len(train), 2)
            self.assertEqual(len(report["files"]["train.jsonl"]["sha256"]), 64)
            with self.assertRaisesRegex(ValueError, "new or empty"):
                prepare(root / "data", 2, 1, sources, download=False)