"""Shared parsing/masking helpers for the GSM8K reasoning dataset translation."""
import re

RECORD_RE = re.compile(
    r"<question>(.*?)</question>\s*<thinking>(.*?)</thinking>\s*<answer>(.*?)</answer>", re.S
)
NUM_RE = re.compile(r"\d+(?:\.\d+)?")
PH_RE = re.compile(r"#(\d+)#")

SRC_PROMPT = "Translate the following English sentence into Arabic:\n{text} <ar>"


def parse(text):
    m = RECORD_RE.match(text.strip())
    if not m:
        return None
    return m.group(1).strip(), m.group(2).strip(), m.group(3).strip()


def mask_numbers(text):
    """'168 + 19 = 187' -> ('#0# + #1# = #2#', ['168', '19', '187'])"""
    nums = []

    def repl(m):
        nums.append(m.group(0))
        return f"#{len(nums) - 1}#"

    return NUM_RE.sub(repl, text), nums


def unmask_numbers(text, nums):
    """Restore. Returns (text, ok) — ok is False if any placeholder was lost or duplicated."""
    seen = []

    def repl(m):
        i = int(m.group(1))
        seen.append(i)
        return nums[i] if i < len(nums) else m.group(0)

    out = PH_RE.sub(repl, text)
    return out, sorted(seen) == list(range(len(nums)))


def build_record(question, thinking, answer):
    return f"<question>{question}</question> <thinking>{thinking}</thinking> <answer>{answer}</answer>"