#!/usr/bin/env python3
"""
評価ユーティリティ: boxed answer抽出、数学的等価判定、コード実行、フォーマットチェック
"""

import re
import subprocess
import tempfile
import math
from typing import Optional


# =============================================================================
# Boxed Answer Extraction
# =============================================================================

def extract_boxed_answer(text: str) -> Optional[str]:
    """
    \\boxed{...} から最終回答を抽出する。
    ネスト対応: \\boxed{\\frac{1}{2}} のような場合も正しく抽出。
    複数ある場合は最後のものを返す (最終回答は通常末尾)。
    全角括弧（）や日本語「箱に入れ」表記にも対応。
    """
    # 表記ゆれに対応: \boxed, \\boxed, \Boxed, \\Boxed + 全角括弧
    pattern = r'\\?\\?[Bb]oxed\s*[\{（]'
    matches = list(re.finditer(pattern, text))

    # フォールバック: 日本語「箱に入れ」パターン
    if not matches:
        jp_pattern = r'[\{（]([^}）]+)[}）]\s*(?:を箱に入れ|を.*?箱)'
        jp_match = re.search(jp_pattern, text)
        if jp_match:
            return jp_match.group(1).strip()
        return None

    # 最後のマッチから抽出 (最終回答)
    last_match = matches[-1]
    start = last_match.end()  # '{' または '（' の直後

    # ネスト対応のブレースマッチング（全角括弧対応）
    depth = 1
    i = start
    while i < len(text) and depth > 0:
        if text[i] in ('{', '（'):
            depth += 1
        elif text[i] in ('}', '）'):
            depth -= 1
        i += 1

    if depth == 0:
        return text[start:i - 1].strip()
    return None


def normalize_answer(answer: str) -> str:
    """回答文字列を正規化する (比較用)"""
    if answer is None:
        return ""
    s = answer.strip()
    # $記号を除去
    s = s.replace("$", "")
    # LaTeXコマンドの正規化
    s = s.replace("\\%", "%")
    s = s.replace("\\$", "$")
    # LaTeXコマンド周辺のスペースを正規化
    # "\\ frac {a} {b}" → "\frac{a}{b}"
    s = re.sub(r'\\+\s*frac\s*', r'\\frac', s)
    s = re.sub(r'(\\frac)\s*\{', r'\1{', s)
    s = re.sub(r'\}\s*\{', '}{', s)
    # 余分な空白を除去
    s = re.sub(r'\s+', ' ', s).strip()
    return s


def math_equivalent(pred: str, gold: str) -> bool:
    """
    2つの数学的表現が等価かどうかを判定する。
    sympyを使った数式パースを試み、失敗した場合は文字列比較にフォールバック。
    """
    pred_norm = normalize_answer(pred)
    gold_norm = normalize_answer(gold)

    if not pred_norm or not gold_norm:
        return False

    # 1. 完全一致
    if pred_norm == gold_norm:
        return True

    # 2. 数値比較 (小数・整数)
    try:
        pred_val = float(eval_simple_expr(pred_norm))
        gold_val = float(eval_simple_expr(gold_norm))
        if math.isclose(pred_val, gold_val, rel_tol=1e-6, abs_tol=1e-9):
            return True
    except (ValueError, TypeError, SyntaxError, ZeroDivisionError):
        pass

    # 3. sympy による数式等価判定
    try:
        import sympy
        from sympy.parsing.latex import parse_latex

        # LaTeX表記をsympyで解析
        try:
            pred_expr = parse_latex(pred_norm)
            gold_expr = parse_latex(gold_norm)
        except Exception:
            # LaTeXパースが失敗した場合、sympy.sympifyを試す
            pred_expr = sympy.sympify(pred_norm)
            gold_expr = sympy.sympify(gold_norm)

        # 数値的に等しいか
        diff = sympy.simplify(pred_expr - gold_expr)
        if diff == 0:
            return True

        # 数値に変換して比較
        pred_float = float(pred_expr.evalf())
        gold_float = float(gold_expr.evalf())
        if math.isclose(pred_float, gold_float, rel_tol=1e-6, abs_tol=1e-9):
            return True
    except Exception:
        pass

    # 4. 文字列の緩い比較 (空白、カンマ区切り等を無視)
    pred_clean = re.sub(r'[,\s\\{}]', '', pred_norm.lower())
    gold_clean = re.sub(r'[,\s\\{}]', '', gold_norm.lower())
    if pred_clean == gold_clean:
        return True

    return False


def eval_simple_expr(s: str) -> float:
    """簡単な数式文字列を評価する (分数、パーセント等に対応)"""
    s = s.strip()
    # パーセント
    if s.endswith('%'):
        return float(s[:-1])
    # LaTeX分数 \frac{a}{b}
    frac_match = re.match(r'\\frac\s*\{([^}]+)\}\s*\{([^}]+)\}', s)
    if frac_match:
        num = float(frac_match.group(1))
        den = float(frac_match.group(2))
        return num / den
    # 通常の分数 a/b
    if '/' in s and not any(c.isalpha() for c in s):
        parts = s.split('/')
        if len(parts) == 2:
            return float(parts[0]) / float(parts[1])
    return float(s)


# =============================================================================
# Code Execution
# =============================================================================

def extract_code_blocks(text: str) -> list[str]:
    """<llm-code>...</llm-code> タグからPythonコードブロックを抽出"""
    pattern = r'<llm-code>(.*?)</llm-code>'
    matches = re.findall(pattern, text, re.DOTALL)
    return [m.strip() for m in matches if m.strip()]


def _auto_print_last_expr(code: str) -> str:
    """
    コードの最終行が print() を含まない式の場合、自動で print() を付与する。
    Jupyter ノートブック形式のコード（最終行が式評価のみ）に対応。
    """
    lines = code.rstrip().split("\n")
    if not lines:
        return code

    last_line = lines[-1].strip()

    # 空行、コメント、代入文、制御文、print文はスキップ
    if (not last_line
            or last_line.startswith("#")
            or "=" in last_line and not last_line.startswith("=") and "==" not in last_line
            or last_line.startswith(("if ", "for ", "while ", "def ", "class ", "import ",
                                     "from ", "return ", "try:", "except", "with ", "raise "))
            or "print(" in last_line):
        return code

    # 最終行が式→ print() で囲む
    lines[-1] = f"print({last_line})"
    return "\n".join(lines)


def execute_code_safely(code: str, timeout: int = 10) -> dict:
    """
    Pythonコードをサブプロセスで安全に実行する。
    最終行が式の場合は自動で print() を付与する。
    Returns: {"success": bool, "stdout": str, "stderr": str}
    """
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
        # sympy は頻出なのでimportを追加
        code_with_print = _auto_print_last_expr(code)
        wrapped = "import sympy\nfrom sympy import *\n" + code_with_print
        f.write(wrapped)
        f.flush()
        try:
            result = subprocess.run(
                ['python3', f.name],
                capture_output=True,
                text=True,
                timeout=timeout,
            )
            return {
                "success": result.returncode == 0,
                "stdout": result.stdout.strip(),
                "stderr": result.stderr.strip(),
            }
        except subprocess.TimeoutExpired:
            return {
                "success": False,
                "stdout": "",
                "stderr": f"Timeout after {timeout}s",
            }
        except Exception as e:
            return {
                "success": False,
                "stdout": "",
                "stderr": str(e),
            }


def check_code_execution(text: str) -> dict:
    """
    テキスト中のコードブロックを実行し、結果をまとめる。
    Returns: {"has_code": bool, "num_blocks": int, "all_success": bool, "results": list}
    """
    blocks = extract_code_blocks(text)
    if not blocks:
        return {"has_code": False, "num_blocks": 0, "all_success": True, "results": []}

    results = []
    for code in blocks:
        result = execute_code_safely(code)
        results.append(result)

    return {
        "has_code": True,
        "num_blocks": len(blocks),
        "all_success": all(r["success"] for r in results),
        "results": results,
    }


# =============================================================================
# Format Compliance
# =============================================================================

def check_format_compliance(text: str) -> dict:
    """回答のフォーマット遵守率をチェック"""
    has_boxed = bool(re.search(r'\\?\\?[Bb]oxed\s*\{', text))

    # コードタグの整合性
    code_opens = len(re.findall(r'<llm-code>', text))
    code_closes = len(re.findall(r'</llm-code>', text))
    code_tags_balanced = code_opens == code_closes

    # 不正なトークンが含まれていないか
    has_bad_tokens = bool(re.search(r'<\|im_end\|>|<\|im_start\|>|<\|endoftext\|>', text))

    return {
        "has_boxed_answer": has_boxed,
        "code_tags_balanced": code_tags_balanced,
        "no_bad_tokens": not has_bad_tokens,
        "format_ok": has_boxed and code_tags_balanced and not has_bad_tokens,
    }