"""Synthetic German lookup/calculator tasks with an authoritative tool environment.

Only reset, lookup and calculator are public environment methods: TRL discovers
public methods as native tools. Scoring and replay helpers deliberately live at
module scope. Never pass a row's private ``task`` metadata to the model prompt.
"""

from __future__ import annotations

import ast
import copy
from fractions import Fraction
import hashlib
import json
import random
import re
from typing import Any


MAX_TOOL_CALLS = 4
MAX_ASSISTANT_TURNS = 4
TEMPLATES = {
    "train": ["Ermittle den Gesamtpreis der Bestellung {record_id}.",
              "Berechne den Warenwert des Auftrags {record_id}."],
    "dev": ["Welcher Rechnungsbetrag ergibt sich für Vorgang {record_id}?",
            "Bestimme die Gesamtkosten des Datensatzes {record_id}."],
    "test": ["Wie hoch ist der Gesamtwert zum Registerschlüssel {record_id}?",
             "Nenne den zu zahlenden Betrag für den Beleg {record_id}."],
}
INSTRUCTIONS = (
    "\n\nWerkzeugprotokoll:\n"
    "Die nötigen Zahlen sind verborgen. Du kennst sie erst nach einer echten lookup-Antwort. "
    "Erfinde weder Zahlen noch Werkzeugantworten und berechne vorher kein Ergebnis.\n"
    "Jede Antwort darf genau einen Werkzeugaufruf ODER die abschließende Antwort enthalten. "
    "Schließe deinen Denkblock mit </think>, BEVOR du <tool_call> ausgibst. "
    "Werkzeugaufrufe innerhalb von <think>...</think> werden nicht ausgeführt.\n"
    "1. Rufe ausschließlich lookup mit dem vollständigen record_id auf, einschließlich beleg-. "
    "Beende danach deine Antwort direkt nach </tool_call> und WARTE auf die echte Werkzeugantwort.\n"
    "2. Erst nach der lookup-Antwort kennst du quantity (Stückzahl) und unit_price (Preis je Stück). "
    "Rufe jetzt in einer neuen Antwort genau einmal calculator auf. expression muss die beiden "
    "tatsächlich empfangenen Ganzzahlen als Zahlenliterale mit * multiplizieren; "
    "verwende keine Variablennamen oder erfundenen Werte. "
    "Schließe auch hier zuerst </think>, gib danach den Aufruf aus und WARTE auf das Rechnerergebnis.\n"
    "3. Erst nach der echten calculator-Antwort darfst du abschließen. "
    "Schreibe nach deinem geschlossenen Denkblock <answer>GANZZAHL</answer>, "
    "ersetze GANZZAHL durch das empfangene Rechnerergebnis und schreibe danach nichts mehr. "
    "Antworte auf Deutsch. Beide echten Werkzeuge sind erforderlich.\n\n"
    "Format deiner ERSTEN Antwort für diesen Auftrag (nur diesen einen Aufruf senden, dann warten):\n"
)
TOOLS = [
    {"type": "function", "function": {
        "name": "lookup", "description": "Lies quantity und unit_price für den angegebenen Datensatz.",
        "parameters": {"type": "object", "properties": {"record_id": {"type": "string"}},
                       "required": ["record_id"], "additionalProperties": False}}},
    {"type": "function", "function": {
        "name": "calculator", "description": "Berechne den Zahlenliteral-Ausdruck nach dem Nachschlagen.",
        "parameters": {"type": "object", "properties": {"expression": {"type": "string"}},
                       "required": ["expression"], "additionalProperties": False}}},
]


def _instructions(record_id: str) -> str:
    lookup_call = json.dumps({"name": "lookup", "arguments": {"record_id": record_id}},
                             ensure_ascii=False)
    return (INSTRUCTIONS + "<think>\nIch muss die echten Datensatzwerte nachschlagen.\n</think>\n\n"
            + "<tool_call>\n" + lookup_call + "\n</tool_call>\n\n"
            + "Dies ist nur der erste Schritt. Erzeuge noch keine calculator-Anfrage, "
            "keine Werkzeugantwort und kein <answer>, bevor die echte lookup-Antwort vorliegt.")


def generate_tasks(split: str, count: int, seed: int = 42) -> list[dict]:
    """Generate reproducible split-specific task IDs, seeds, and wording templates."""
    if split not in TEMPLATES or type(count) is not int or count < 1:
        raise ValueError("Use train/dev/test and a positive integer task count")
    if type(seed) is not int or seed < 0:
        raise ValueError("seed must be a nonnegative integer")
    rows = []
    for index in range(count):
        identity = hashlib.sha256(f"german-tools-v1:{split}:{seed}:{index}".encode()).hexdigest()
        episode_seed = int(identity[:16], 16) & ((1 << 63) - 1)
        rng = random.Random(episode_seed)
        quantity, price = rng.randint(2, 9), rng.randint(2, 12)
        while quantity == price:
            price = rng.randint(2, 12)
        record_id = f"beleg-{identity[16:28]}"
        task_id = f"german-tools-v1-{split}-{seed}-{index:05d}"
        template_index = index % len(TEMPLATES[split])
        task = {"schema_version": 1, "task_id": task_id, "split": split,
                "record_id": record_id, "private_facts": {"quantity": quantity, "unit_price": price},
                "template_id": f"{split}-{template_index}", "episode_seed": episode_seed}
        rows.append({"id": task_id,
                     "prompt": TEMPLATES[split][template_index].format(record_id=record_id) + _instructions(record_id),
                     "task": task})
    return rows


def safe_calculate(expression: str) -> str:
    """Evaluate bounded literal arithmetic with AST interpretation, never eval."""
    return _format_number(_calculate(expression)[0])


def _format_number(value: Fraction) -> str:
    return str(value.numerator) if value.denominator == 1 else f"{value.numerator}/{value.denominator}"


def _calculate(expression: str) -> tuple[Fraction, ast.expr]:
    if not isinstance(expression, str) or not expression.strip() or len(expression) > 128:
        raise ValueError("expression must be a nonempty string of at most 128 characters")
    try:
        tree = ast.parse(expression.strip(), mode="eval")
    except (SyntaxError, ValueError) as exc:
        raise ValueError("invalid arithmetic syntax") from exc
    if sum(1 for _ in ast.walk(tree)) > 31:
        raise ValueError("arithmetic expression is too complex")

    def visit(node: ast.expr, depth: int = 0) -> Fraction:
        if depth > 8:
            raise ValueError("arithmetic nesting limit exceeded")
        if isinstance(node, ast.Constant) and type(node.value) is int and abs(node.value) <= 10_000:
            result = Fraction(node.value)
        elif isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.UAdd, ast.USub)):
            result = visit(node.operand, depth + 1) * (-1 if isinstance(node.op, ast.USub) else 1)
        elif isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Sub, ast.Mult, ast.Div)):
            left, right = visit(node.left, depth + 1), visit(node.right, depth + 1)
            if isinstance(node.op, ast.Add):
                result = left + right
            elif isinstance(node.op, ast.Sub):
                result = left - right
            elif isinstance(node.op, ast.Mult):
                result = left * right
            elif right:
                result = left / right
            else:
                raise ValueError("division by zero")
        else:
            raise ValueError("only bounded integers, parentheses, +, -, * and / are allowed")
        if abs(result) > 1_000_000 or result.denominator > 1_000_000:
            raise ValueError("arithmetic result limit exceeded")
        return result

    return visit(tree.body), tree.body


class GermanAgentEnv:
    """Per-rollout private state for TRL's native multi-turn environment_factory."""

    def __init__(self):
        self.trace: list[dict] = []
        self._task: dict | None = None
        self._lookup_success = False
        self._calculation_success = False
        self._workflow_valid = True
        self._failure_reason: str | None = None
        self._replay_turn = 0
        self._terminal: dict | None = None
        self._tool_limit_hit = False

    def reset(self, task: dict | None = None, **row: Any) -> None:
        """Reset private task state; the returned value never exposes hidden facts."""
        if not isinstance(task, dict) or task.get("schema_version") != 1:
            raise ValueError("reset requires generated task metadata with schema_version=1")
        facts = task.get("private_facts")
        if not isinstance(facts, dict) or set(facts) != {"quantity", "unit_price"}:
            raise ValueError("task.private_facts requires quantity and unit_price")
        if any(type(value) is not int or not 2 <= value <= 100 for value in facts.values()):
            raise ValueError("task facts must be bounded positive integers")
        if not isinstance(task.get("record_id"), str) or not task["record_id"]:
            raise ValueError("task requires a record_id")
        self.__init__()
        self._task = copy.deepcopy(task)

    def _record(self, tool: str, arguments: dict, observation: dict, ok: bool) -> str:
        text = json.dumps(observation, ensure_ascii=False, sort_keys=True)
        self.trace.append({"index": len(self.trace), "tool": tool,
                           "arguments": copy.deepcopy(arguments), "observation": text,
                           "ok": ok, "turn": self._replay_turn or None})
        if not ok:
            self._workflow_valid = False
            self._failure_reason = observation.get("error", "tool_error")
        return text

    def _guard(self, tool: str, arguments: dict) -> str | None:
        if self._task is None:
            raise ValueError("reset the environment before using tools")
        if self._terminal is not None:
            return self._record(tool, arguments, {"error": "episode_already_finished"}, False)
        if len(self.trace) >= MAX_TOOL_CALLS:
            if self._tool_limit_hit:
                return json.dumps({"error": "tool_call_limit"})
            self._tool_limit_hit = True
            return self._record(tool, arguments, {"error": "tool_call_limit"}, False)
        return None

    def lookup(self, record_id: str) -> str:
        """Look up hidden quantity and unit_price for a record.

        Args:
            record_id: The exact record identifier supplied in the German task.
        """
        arguments = {"record_id": record_id}
        if error := self._guard("lookup", arguments):
            return error
        if record_id != self._task["record_id"]:
            return self._record("lookup", arguments, {"error": "unknown_record"}, False)
        self._lookup_success = True
        return self._record("lookup", arguments, {"record_id": record_id, **self._task["private_facts"]}, True)

    def calculator(self, expression: str) -> str:
        """Evaluate bounded literal arithmetic after looking up the hidden values.

        Args:
            expression: The product of the two retrieved integer values, as literals.
        """
        arguments = {"expression": expression}
        if error := self._guard("calculator", arguments):
            return error
        if not self._lookup_success:
            return self._record("calculator", arguments, {"error": "lookup_required"}, False)
        try:
            value, tree = _calculate(expression)
        except ValueError as exc:
            return self._record("calculator", arguments, {"error": "invalid_expression", "message": str(exc)}, False)
        facts = self._task["private_facts"]
        operands = [tree.left.value, tree.right.value] if (
            isinstance(tree, ast.BinOp) and isinstance(tree.op, ast.Mult)
            and isinstance(tree.left, ast.Constant) and isinstance(tree.right, ast.Constant)
            and type(tree.left.value) is int and type(tree.right.value) is int
        ) else []
        semantic_match = sorted(operands) == sorted(facts.values())
        if not semantic_match:
            return self._record("calculator", arguments,
                                {"error": "must_multiply_retrieved_values", "result": _format_number(value)}, False)
        self._calculation_success = True
        return self._record("calculator", arguments, {"result": _format_number(value)}, True)


def _visible_content(text: str) -> str | None:
    if not isinstance(text, str):
        return None
    if text.rfind("<think>") > text.rfind("</think>"):
        return None
    return text.rsplit("</think>", 1)[-1].strip()


def score_episode(env: GermanAgentEnv, completion: str | list[dict]) -> dict:
    """Reward outcome AND legal observed workflow; explanations/language are unscored."""
    if env._task is None:
        raise ValueError("reset environment before scoring")
    assistant_turns = None
    distinct_tool_turns = True
    transcript_valid = True
    if isinstance(completion, list):
        transcript_valid = bool(
            completion and isinstance(completion[-1], dict)
            and completion[-1].get("role") == "assistant"
            and not completion[-1].get("tool_calls")
        )
        assistants = [message for message in completion if isinstance(message, dict) and message.get("role") == "assistant"]
        assistant_turns = len(assistants)
        text = (assistants[-1].get("content") or "") if assistants else ""
        lookup_turns, calculator_turns = [], []
        declared_actions = []
        for index, message in enumerate(assistants):
            for call in message.get("tool_calls") or []:
                try:
                    name = call["function"]["name"]
                    arguments = call["function"]["arguments"]
                    arguments = json.loads(arguments) if isinstance(arguments, str) else arguments
                    required = {"record_id"} if name == "lookup" else {"expression"}
                    if not isinstance(arguments, dict) or set(arguments) != required or any(
                        not isinstance(value, str) for value in arguments.values()
                    ):
                        raise ValueError("invalid native arguments")
                    declared_actions.append((name, arguments))
                except (KeyError, TypeError, ValueError):
                    transcript_valid = False
                    continue
                if name == "lookup":
                    lookup_turns.append(index)
                elif name == "calculator":
                    calculator_turns.append(index)
                else:
                    transcript_valid = False
        distinct_tool_turns = bool(lookup_turns and calculator_turns and min(lookup_turns) < min(calculator_turns))
        transcript_valid = transcript_valid and assistant_turns <= MAX_ASSISTANT_TURNS
        transcript_valid = transcript_valid and declared_actions == [
            (row["tool"], row["arguments"]) for row in env.trace
        ]
    else:
        text = completion
        if env._replay_turn:
            assistant_turns = env._replay_turn
            transcript_valid = assistant_turns <= MAX_ASSISTANT_TURNS
            lookups = [row["turn"] for row in env.trace if row["tool"] == "lookup" and row["ok"]]
            calculations = [row["turn"] for row in env.trace if row["tool"] == "calculator" and row["ok"]]
            distinct_tool_turns = bool(lookups and calculations and min(lookups) < min(calculations))
        else:
            # Direct tool calls without a transcript carry no model-turn evidence.
            transcript_valid = False
    visible = _visible_content(text)
    match = re.search(r"<answer>(-?(?:0|[1-9][0-9]*))</answer>\Z", visible) if visible is not None else None
    valid_format = match is not None and visible.count("<answer>") == visible.count("</answer>") == 1
    answer = match.group(1) if valid_format else None
    workflow = (env._workflow_valid and env._lookup_success and env._calculation_success
                and len(env.trace) <= MAX_TOOL_CALLS and transcript_valid and distinct_tool_turns)
    expected = str(env._task["private_facts"]["quantity"] * env._task["private_facts"]["unit_price"])
    success = workflow and valid_format and answer == expected
    reason = "success" if success else (
        env._failure_reason or ("invalid_workflow" if not workflow else
                                "invalid_final_answer" if not valid_format else "incorrect_final_answer"))
    return {"reward": float(success), "success": bool(success), "reason": reason,
            "tool_calls": len(env.trace), "lookup_success": env._lookup_success,
            "calculation_success": env._calculation_success, "workflow_valid": workflow,
            "final_answer": answer, "assistant_turns": assistant_turns, "trace": copy.deepcopy(env.trace)}


def parse_response(text: str) -> list[dict]:
    """Strict native-Qwen text parsing for fixture/replay; native TRL parsing is preferred."""
    visible = _visible_content(text)
    if visible is None:
        raise ValueError("unfinished thinking block")
    if "<tool_call>" not in visible:
        return [{"final": visible}]
    blocks = re.findall(r"<tool_call>(.*?)</tool_call>", visible, re.DOTALL)
    remainder = re.sub(r"<tool_call>.*?</tool_call>", "", visible, flags=re.DOTALL)
    if not blocks or remainder.strip():
        raise ValueError("malformed native tool call")
    actions = []
    for block in blocks:
        call = json.loads(block)
        if not isinstance(call, dict) or set(call) != {"name", "arguments"}:
            raise ValueError("tool call must contain name and arguments")
        arguments = json.loads(call["arguments"]) if isinstance(call["arguments"], str) else call["arguments"]
        if not isinstance(arguments, dict):
            raise ValueError("tool arguments must be an object")
        if call["name"] not in {"lookup", "calculator"} or "tool" in arguments:
            raise ValueError("unsupported tool or reserved argument")
        actions.append({"tool": call["name"], **arguments})
    return actions


def step_action(env: GermanAgentEnv, action: dict | str) -> dict:
    """Dispatch an already parsed native tool action without exposing private metadata."""
    if isinstance(action, str):
        action = json.loads(action)
    if not isinstance(action, dict):
        raise ValueError("action must be an object")
    if env._terminal is not None:
        result = env._terminal
        return {"observation": "", "done": True,
                **{key: result[key] for key in ("reward", "success", "reason")}, "result": result}
    if set(action) == {"final"}:
        result = score_episode(env, action["final"])
        env._terminal = result
        return {"observation": "", "done": True, **{key: result[key] for key in ("reward", "success", "reason")}, "result": result}
    if action.get("tool") == "lookup" and set(action) == {"tool", "record_id"}:
        observation = env.lookup(action["record_id"])
    elif action.get("tool") == "calculator" and set(action) == {"tool", "expression"}:
        observation = env.calculator(action["expression"])
    else:
        env._workflow_valid = False
        env._failure_reason = "invalid_action"
        observation = json.dumps({"error": "invalid_action"})
    return {"observation": observation, "done": len(env.trace) > MAX_TOOL_CALLS,
            "reward": 0.0, "success": False, "reason": env._failure_reason}


def apply_response(env: GermanAgentEnv, text: str) -> dict:
    """Replay one model turn, retaining observation boundaries and terminal outcomes."""
    if env._terminal is not None:
        return {"observations": [], "terminal": True, "result": env._terminal}
    env._replay_turn += 1
    try:
        if env._replay_turn > MAX_ASSISTANT_TURNS:
            raise ValueError("assistant turn limit")
        actions = parse_response(text)
    except (ValueError, TypeError):
        env._workflow_valid = False
        env._failure_reason = "malformed_response_or_turn_limit"
        result = score_episode(env, "")
        env._terminal = result
        return {"observations": [], "terminal": True, "result": result}
    observations = []
    for action in actions:
        result = step_action(env, action)
        if result["done"]:
            terminal = result.get("result") or score_episode(env, "")
            env._terminal = terminal
            return {"observations": observations, "terminal": True, "result": terminal}
        observations.append({"tool": action.get("tool"),
                             "arguments": {key: value for key, value in action.items() if key != "tool"},
                             "observation": result["observation"]})
    return {"observations": observations, "terminal": False, "result": None}