#!/usr/bin/env python3
"""Recursive inference engine for stateless recursive math-solving models."""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


TAG_RE = re.compile(r"</?(problem|action|parallel|path)>", re.IGNORECASE)
NUMBER_RE = re.compile(r"-?\d+(?:,\d{3})*(?:\.\d+)?|-?\d+/\d+")
FINAL_ANSWER_LINE_RE = re.compile(
    r"(?im)^[ \t]*Final[ \t]+answer[ \t]*:[ \t]*(?P<answer>[^\n]+)$"
)
DEFAULT_PREFIX_PROMPT_FILE = Path("prompts") / "recursive_sft_prefix_prompt.txt"
FORCED_ROOT_OUTPUT_PREFIXES = {
    "SOLVE_DIRECTLY": "<action>\nSOLVE_DIRECTLY\n",
    "DECOMPOSE": "<action>\nDECOMPOSE\n</action>\n<parallel>\n",
}


@dataclass
class ModelCall:
    call_id: int
    depth: int
    role: str
    problem: str
    input_text: str
    output_text: str
    input_tokens: int
    output_tokens: int
    action_type: str | None = None
    parse_error: str | None = None
    forced_output_prefix: str = ""


@dataclass
class RolloutResult:
    problem: str
    final_answer: str | None
    full_trace: str
    calls: list[ModelCall] = field(default_factory=list)
    parse_errors: list[str] = field(default_factory=list)
    critical_path_tokens: int = 0
    stopped_reason: str = "completed"

    @property
    def parsed_ok(self) -> bool:
        return not self.parse_errors

    @property
    def model_call_count(self) -> int:
        return len(self.calls)


@dataclass
class NodeResult:
    final_answer: str | None
    trace: str
    critical_path_tokens: int
    parse_errors: list[str]
    stopped_reason: str = "completed"


def load_prefix_template(path: str | Path | None = None) -> str:
    prompt_path = Path(path) if path is not None else DEFAULT_PREFIX_PROMPT_FILE
    return prompt_path.read_text(encoding="utf-8")


def format_model_input(problem: str, prefix_template: str) -> str:
    return prefix_template.strip().replace("{PROBLEM}", problem.strip())


def token_count(tokenizer: Any, text: str) -> int:
    return len(tokenizer.encode(text, add_special_tokens=False))


def extract_blocks(text: str, tag: str) -> list[str]:
    return [
        match.group(1).strip()
        for match in re.finditer(rf"<{tag}>(.*?)</{tag}>", text, flags=re.DOTALL | re.IGNORECASE)
    ]


def normalize_answer(value: Any) -> str | None:
    if value is None:
        return None
    text = str(value).strip()
    if not text:
        return None
    if "####" in text:
        text = text.split("####")[-1].strip()
    boxed = re.findall(r"\\boxed\{([^{}]+)\}", text)
    if boxed:
        text = boxed[-1]
    final = re.findall(r"(?i)Final\s+answer\s*:\s*([^\n<]+)", text)
    if final:
        text = final[-1]
    text = text.replace("$", "").replace(",", "")
    text = re.sub(r"\\(?:text|mathrm)\{([^{}]+)\}", r"\1", text)
    text = re.sub(r"\s+", "", text)
    text = text.rstrip(".")
    return text or None


def extract_final_answer(text: str) -> str | None:
    patterns = [
        r"(?i)Final\s+answer\s*:\s*([^\n<]+)",
        r"(?i)answer\s+is\s*([^\n<]+)",
        r"\\boxed\{([^{}]+)\}",
    ]
    for pattern in patterns:
        matches = re.findall(pattern, text)
        if matches:
            return normalize_answer(matches[-1])
    matches = NUMBER_RE.findall(text)
    if matches:
        return normalize_answer(matches[-1])
    return None


def answers_match(pred: str | None, gt: str | None) -> bool:
    pred_norm = normalize_answer(pred)
    gt_norm = normalize_answer(gt)
    if pred_norm is None or gt_norm is None:
        return False

    if pred_norm == gt_norm:
        return True

    try:
        from math_verify import LatexExtractionConfig, parse, verify

        pred_parsed = parse(pred_norm, extraction_config=[LatexExtractionConfig()])
        gt_parsed = parse(gt_norm, extraction_config=[LatexExtractionConfig()])
        if pred_parsed and gt_parsed and verify(gt_parsed, pred_parsed):
            return True
    except Exception:
        pass

    return pred_norm.lower() == gt_norm.lower()


def find_matching_tag(text: str, tag: str, open_start: int, open_end: int) -> tuple[int, int] | None:
    depth = 1
    for match in TAG_RE.finditer(text, open_end):
        token = match.group(0).lower()
        name = match.group(1).lower()
        if name != tag:
            continue
        if token.startswith("</"):
            depth -= 1
            if depth == 0:
                return match.start(), match.end()
        else:
            depth += 1
    return None


def parse_model_output(output: str) -> tuple[str, list[str]]:
    action_match = re.search(r"<action>(.*?)</action>", output, flags=re.DOTALL | re.IGNORECASE)
    if not action_match:
        raise ValueError("missing <action>...</action> block")

    action_body = action_match.group(1).strip()
    first_line = action_body.splitlines()[0].strip().upper() if action_body else ""
    if first_line.startswith("SOLVE_DIRECTLY"):
        return "SOLVE_DIRECTLY", []
    if first_line.startswith("AGGREGATE"):
        return "AGGREGATE", []
    if not first_line.startswith("DECOMPOSE"):
        raise ValueError(f"unknown action type in first action line: {first_line!r}")

    next_nonspace = re.search(r"\S", output[action_match.end() :])
    if not next_nonspace:
        raise ValueError("DECOMPOSE action is not followed by <parallel>")

    parallel_start = action_match.end() + next_nonspace.start()
    if not output[parallel_start:].lower().startswith("<parallel>"):
        raise ValueError("DECOMPOSE action is not immediately followed by <parallel>")

    parallel_open_end = parallel_start + len("<parallel>")
    close = find_matching_tag(output, "parallel", parallel_start, parallel_open_end)
    if close is None:
        raise ValueError("unclosed <parallel> block")
    parallel_body = output[parallel_open_end : close[0]]
    paths = extract_blocks(f"<parallel>{parallel_body}</parallel>", "path")
    if len(paths) < 2:
        raise ValueError(f"DECOMPOSE parallel block has {len(paths)} path(s), expected at least 2")
    if any(not path for path in paths):
        raise ValueError("DECOMPOSE parallel block contains an empty path")
    return "DECOMPOSE", paths


def validate_action_protocol(output: str, action_type: str, role: str) -> None:
    """Enforce role-specific actions and explicit answers for terminal calls."""
    normalized_role = role.strip().lower()
    if normalized_role == "root":
        allowed_actions = {"DECOMPOSE", "SOLVE_DIRECTLY"}
    elif normalized_role == "aggregate":
        allowed_actions = {"AGGREGATE"}
    elif normalized_role == "child" or normalized_role.startswith("path_"):
        allowed_actions = {"DECOMPOSE", "SOLVE_DIRECTLY"}
    else:
        raise ValueError(f"unknown call role: {role!r}")

    if action_type not in allowed_actions:
        expected = " or ".join(sorted(allowed_actions))
        raise ValueError(
            f"{normalized_role} call returned {action_type}, expected {expected}"
        )

    if action_type not in {"SOLVE_DIRECTLY", "AGGREGATE"}:
        return

    action_match = re.search(
        r"<action>(.*?)</action>", output, flags=re.DOTALL | re.IGNORECASE
    )
    action_body = action_match.group(1) if action_match else ""
    final_answer = FINAL_ANSWER_LINE_RE.search(action_body)
    if final_answer is None or not final_answer.group("answer").strip():
        raise ValueError(f"{action_type} action is missing a non-empty Final answer: line")


def extract_final_answer_display(text: str) -> str | None:
    patterns = [
        r"(?i)Final\s+answer\s*:\s*([^\n<]+)",
        r"\\boxed\{([^{}]+)\}",
    ]
    for pattern in patterns:
        matches = re.findall(pattern, text)
        if matches:
            return matches[-1].strip().rstrip(".")
    return None


def child_result_summary(result: NodeResult) -> str:
    display = extract_final_answer_display(result.trace)
    if display is not None:
        return display
    if result.final_answer is not None:
        return str(result.final_answer)
    fallback = extract_final_answer(result.trace)
    if fallback is not None:
        return str(fallback)
    return "PARSE_ERROR" if result.parse_errors else "UNKNOWN"


def build_aggregation_problem(parent_problem: str, paths: list[str], child_results: list[NodeResult]) -> str:
    lines = [
        f"The original problem is: {parent_problem.strip()}",
        "The paths spawned were:",
    ]
    for idx, (path, result) in enumerate(zip(paths, child_results), start=1):
        summary = child_result_summary(result)
        lines.append(f"{idx}. {path.strip()} (Result: {summary}).")
    lines.append("Aggregate these results to find the answer to the original problem.")
    return "\n".join(lines)


class RecursiveEngine:
    def __init__(
        self,
        *,
        model: Any,
        tokenizer: Any,
        prefix_template: str,
        max_new_tokens: int = 768,
        temperature: float = 0.0,
        top_p: float = 1.0,
        top_k: int | None = None,
        max_depth: int = 3,
        max_calls: int = 10,
        force_root_action: str | None = None,
    ) -> None:
        self.model = model
        self.tokenizer = tokenizer
        self.prefix_template = prefix_template
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.max_depth = max_depth
        self.max_calls = max_calls
        normalized_force = force_root_action.upper() if force_root_action else None
        if normalized_force not in {None, *FORCED_ROOT_OUTPUT_PREFIXES}:
            raise ValueError(
                "force_root_action must be SOLVE_DIRECTLY, DECOMPOSE, or None"
            )
        self.force_root_action = normalized_force
        self.calls: list[ModelCall] = []

    def generate_once(self, input_text: str, *, output_prefix: str = "") -> str:
        import torch

        generation_input = input_text + output_prefix
        inputs = self.tokenizer(generation_input, return_tensors="pt")
        device = next(self.model.parameters()).device
        inputs = {key: value.to(device) for key, value in inputs.items()}
        generation_kwargs = {
            "max_new_tokens": self.max_new_tokens,
            "do_sample": self.temperature > 0,
            "top_p": self.top_p,
            "use_cache": True,
            "pad_token_id": self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
            "eos_token_id": self.tokenizer.eos_token_id,
        }
        if self.temperature > 0:
            generation_kwargs["temperature"] = self.temperature
            if self.top_k is not None:
                generation_kwargs["top_k"] = self.top_k
        with torch.no_grad():
            generated = self.model.generate(
                **inputs,
                **generation_kwargs,
            )
        new_tokens = generated[0, inputs["input_ids"].shape[1] :]
        continuation = self.tokenizer.decode(new_tokens, skip_special_tokens=False).strip()
        return output_prefix + continuation

    def rollout(self, problem: str) -> RolloutResult:
        self.calls = []
        root = self._solve_node(problem=problem, depth=0, role="root")
        return RolloutResult(
            problem=problem,
            final_answer=root.final_answer,
            full_trace=root.trace,
            calls=self.calls,
            parse_errors=root.parse_errors,
            critical_path_tokens=root.critical_path_tokens,
            stopped_reason=root.stopped_reason,
        )

    def _solve_node(self, *, problem: str, depth: int, role: str) -> NodeResult:
        if len(self.calls) >= self.max_calls:
            return NodeResult(
                final_answer=None,
                trace="",
                critical_path_tokens=0,
                parse_errors=[f"max model calls exceeded ({self.max_calls})"],
                stopped_reason="max_calls",
            )
        if depth > self.max_depth:
            return NodeResult(
                final_answer=None,
                trace="",
                critical_path_tokens=0,
                parse_errors=[f"max recursion depth exceeded ({self.max_depth})"],
                stopped_reason="max_depth",
            )

        input_text = format_model_input(problem, self.prefix_template)
        forced_output_prefix = ""
        if depth == 0 and role == "root" and self.force_root_action is not None:
            forced_output_prefix = FORCED_ROOT_OUTPUT_PREFIXES[self.force_root_action]
        output_text = self.generate_once(
            input_text,
            output_prefix=forced_output_prefix,
        )
        generated_output = (
            output_text[len(forced_output_prefix) :]
            if forced_output_prefix and output_text.startswith(forced_output_prefix)
            else output_text
        )
        call = ModelCall(
            call_id=len(self.calls) + 1,
            depth=depth,
            role=role,
            problem=problem,
            input_text=input_text,
            output_text=output_text,
            input_tokens=token_count(self.tokenizer, input_text + forced_output_prefix),
            output_tokens=token_count(self.tokenizer, generated_output),
            forced_output_prefix=forced_output_prefix,
        )
        self.calls.append(call)
        own_tokens = call.input_tokens + call.output_tokens

        try:
            action_type, paths = parse_model_output(output_text)
            validate_action_protocol(output_text, action_type, role)
            call.action_type = action_type
        except ValueError as exc:
            call.parse_error = str(exc)
            return NodeResult(
                final_answer=None,
                trace=output_text,
                critical_path_tokens=own_tokens,
                parse_errors=[f"call_{call.call_id}: {exc}"],
                stopped_reason="parse_error",
            )

        if action_type in {"SOLVE_DIRECTLY", "AGGREGATE"}:
            return NodeResult(
                final_answer=extract_final_answer(output_text),
                trace=output_text,
                critical_path_tokens=own_tokens,
                parse_errors=[],
            )

        child_results = [
            self._solve_node(problem=path, depth=depth + 1, role=f"path_{idx}")
            for idx, path in enumerate(paths, start=1)
        ]
        aggregate_problem = build_aggregation_problem(problem, paths, child_results)
        aggregate_result = self._solve_node(
            problem=aggregate_problem,
            depth=depth + 1,
            role="aggregate",
        )
        child_trace = "\n\n".join(
            f"<problem>\n{path.strip()}\n</problem>\n{result.trace.strip()}"
            for path, result in zip(paths, child_results)
        )
        full_trace = "\n\n".join(
            part
            for part in [
                output_text.strip(),
                child_trace.strip(),
                f"<problem>\n{aggregate_problem.strip()}\n</problem>\n{aggregate_result.trace.strip()}",
            ]
            if part
        )
        child_critical = max((result.critical_path_tokens for result in child_results), default=0)
        parse_errors: list[str] = []
        stopped_reason = "completed"
        for result in child_results + [aggregate_result]:
            parse_errors.extend(result.parse_errors)
            if result.stopped_reason != "completed":
                stopped_reason = result.stopped_reason

        return NodeResult(
            final_answer=aggregate_result.final_answer,
            trace=full_trace,
            critical_path_tokens=own_tokens + child_critical + aggregate_result.critical_path_tokens,
            parse_errors=parse_errors,
            stopped_reason=stopped_reason,
        )


def load_model_and_tokenizer(model_path: str, dtype: str = "bf16", device_map: str | None = None) -> tuple[Any, Any]:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    dtype_map = {
        "bf16": torch.bfloat16,
        "fp16": torch.float16,
        "fp32": torch.float32,
        "auto": "auto",
    }
    torch_dtype = dtype_map.get(dtype, torch.bfloat16)
    kwargs: dict[str, Any] = {
        "torch_dtype": torch_dtype,
        "trust_remote_code": True,
    }
    if device_map:
        kwargs["device_map"] = device_map
    model = AutoModelForCausalLM.from_pretrained(model_path, **kwargs)
    if not device_map:
        model = model.to("cuda" if torch.cuda.is_available() else "cpu")
    # SFT checkpoints commonly persist use_cache=False from gradient checkpointing.
    # Generation should always use the KV cache, or long rollouts become prohibitively slow.
    model.config.use_cache = True
    model.eval()
    return model, tokenizer


def row_to_eval_record(
    item: dict[str, Any],
    *,
    idx: int,
    split: str,
    problem_field: str | None,
    answer_field: str | None,
) -> dict[str, Any]:
    p_field = problem_field or infer_field(item, ["problem", "question", "input"])
    a_field = answer_field or infer_field(item, ["answer", "target", "final_answer", "solution"])
    if p_field is None:
        raise ValueError(f"Could not infer problem field. Available fields: {list(item.keys())}")
    return {
        "id": str(item.get("id", item.get("unique_id", f"{split}_{idx:06d}"))),
        "problem": str(item[p_field]),
        "answer": str(item[a_field]) if a_field and item.get(a_field) is not None else None,
        "raw": dict(item),
    }


def load_eval_jsonl(
    *,
    path: str | Path,
    problem_field: str | None,
    answer_field: str | None,
    limit: int | None,
    start_index: int = 0,
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    path = Path(path)
    with path.open("r", encoding="utf-8") as handle:
        for absolute_idx, line in enumerate(handle):
            if absolute_idx < start_index:
                continue
            if limit is not None and len(rows) >= limit:
                break
            if not line.strip():
                continue
            item = json.loads(line)
            rows.append(
                row_to_eval_record(
                    item,
                    idx=absolute_idx,
                    split="jsonl",
                    problem_field=problem_field,
                    answer_field=answer_field,
                )
            )
    return rows


def load_eval_dataset(
    *,
    dataset: str,
    dataset_config: str | None,
    split: str,
    problem_field: str | None,
    answer_field: str | None,
    limit: int | None,
    start_index: int = 0,
) -> list[dict[str, Any]]:
    from datasets import load_dataset

    ds = load_dataset(dataset, dataset_config, split=split) if dataset_config else load_dataset(dataset, split=split)
    if start_index:
        ds = ds.select(range(start_index, len(ds)))
    if limit is not None:
        ds = ds.select(range(min(limit, len(ds))))

    rows: list[dict[str, Any]] = []
    for idx, item in enumerate(ds, start=start_index):
        rows.append(
            row_to_eval_record(
                dict(item),
                idx=idx,
                split=split,
                problem_field=problem_field,
                answer_field=answer_field,
            )
        )
    return rows


def infer_field(item: dict[str, Any], candidates: list[str]) -> str | None:
    for candidate in candidates:
        if candidate in item:
            return candidate
    return None


def write_jsonl(path: str | Path, rows: list[dict[str, Any]]) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")