"""
MathVerifier — Symbolic math verification for math-domain LLM agent chains.
============================================================================

Addresses math AUROC = 0.500 (random) found in exp_code_math_farl:
  - SC_OLD behavioral features (step count, thought variance, Jaccard overlap)
    are blind to numeric correctness.
  - This module adds a symbolic verification layer that directly checks the
    final numeric answer against expressions found in the reasoning chain.

Verification pipeline (priority order)
---------------------------------------
  1. Python exec in subprocess sandbox — arithmetic, algebra, fractions,
     percentages, scientific notation. 2-second timeout, no imports except
     math and fractions.
  2. SymPy symbolic evaluation — algebra, calculus, CAS simplification.
     Falls back gracefully if sympy is not installed.
  3. Regex numeric extraction + direct comparison — last-resort for
     plain arithmetic when no expression can be constructed.

Risk output
-----------
  0.10  answer verified correct  (strong signal → chain is likely good)
  0.50  could not verify          (neutral      → fall through to SC_OLD)
  0.80  answer appears incorrect  (strong signal → chain is likely wrong)

Integration with AgentGuard (score_chain)
-----------------------------------------
  When is_math_chain() returns True, score_chain() blends math_risk into
  the behavioral sc_old_risk using a 50/50 soft blend — ONLY when the
  verifier is confident (math_risk < 0.3 or math_risk > 0.7).
  Non-math chains are completely unaffected.

Usage
-----
    from llm_guard.verification.math_verifier import MathVerifier

    mv = MathVerifier()
    risk = mv.score("What is 15% of 240?", steps, "36")
    # 0.10 → verified correct
    risk = mv.score("What is 15% of 240?", steps, "48")
    # 0.80 → verified wrong (expected 36, got 48)
Purpose:    Adds a symbolic verification layer (Python exec sandbox, SymPy, regex fallback) to detect numeric incorrectness. TPR=100%, FPR=15% on GSM8K-style math; anti-correlated on competition math.
Exports:    MathVerifier
Added:      v0.25.0
Status:     Validated
"""

from __future__ import annotations

import re
import subprocess
from typing import List, Optional, Tuple


# ── Regex patterns ────────────────────────────────────────────────────────────

# Matches numbers in various forms: integer, decimal, fraction, percentage,
# scientific notation, negative.
_NUM_RE = re.compile(
    r"(?<![a-zA-Z\d])"          # negative look-behind: not after letters/digits
    r"(-?)"                      # optional leading minus
    r"(\d[\d,]*"                 # integer part (with optional thousands commas)
    r"(?:\.\d+)?"                # optional decimal part
    r"(?:[eE][+-]?\d+)?"         # optional scientific notation
    r")"
    r"(?:\s*%)?",                # optional trailing percentage sign
    re.UNICODE,
)

# Fraction: "3/4", "22/7"
_FRAC_RE = re.compile(r"(-?\d+)\s*/\s*(\d+)")

# Equation: "x = 36", "answer = 36.0", "result = ..."
_EQ_RE = re.compile(
    r"(?:answer|result|value|total|sum|product|=)\s*[=:]\s*(-?[\d,]+\.?\d*(?:[eE][+-]?\d+)?)",
    re.IGNORECASE,
)

# Math keyword detector for is_math_chain()
_MATH_KW_RE = re.compile(
    r"\b(calculat|comput|how many|what is \d|percent|total|sum|average|"
    r"mean|ratio|equation|solve|simplif|factor|integrat|derivat|probability|"
    r"area|volume|perimeter|circumference|hypotenuse|distance|speed|rate|"
    r"profit|loss|discount|interest|tax|cost|price|revenue|dividend|"
    r"multiply|divide|subtract|add|square root|cube root)\b",
    re.IGNORECASE,
)

# Expression patterns that look like arithmetic (e.g. "15 * 240 / 100")
_ARITH_OP_RE = re.compile(r"\d[\d\s]*[+\-*/^%][\d\s]*\d")


class MathVerifier:
    """
    Symbolic math verification for math-domain LLM agent chains.

    Instantiate once and call score() for each chain. Fully stateless.
    No external dependencies required — SymPy is used opportunistically
    when available, but the verifier works without it.

    Parameters
    ----------
    tolerance : float
        Relative tolerance for numeric comparison.
        Default 0.02 (2%) to handle rounding in intermediate steps.
    timeout_sec : float
        Subprocess timeout for Python exec sandbox. Default 2.0 s.
    """

    def __init__(self, tolerance: float = 0.02, timeout_sec: float = 2.0) -> None:
        self._tol         = tolerance
        self._timeout_sec = timeout_sec
        # Check SymPy once at init; cache result
        self._has_sympy: Optional[bool] = None

    # ── Public API ─────────────────────────────────────────────────────────────

    def score(
        self,
        question: str,
        steps: List[dict],
        final_answer: str,
    ) -> float:
        """
        Compute a risk score for a completed math chain.

        Returns
        -------
        float
            0.10  verified correct
            0.50  could not verify (neutral)
            0.80  verified wrong
        """
        # Extract the expected numeric answer from the chain
        expected = self._extract_expected_from_chain(question, steps, final_answer)
        claimed  = self._extract_number(final_answer)

        if expected is None or claimed is None:
            # Neither extraction succeeded — neutral signal
            return 0.50

        if self._approx_equal(claimed, expected):
            return 0.10
        else:
            return 0.80

    def is_math_chain(
        self,
        question: str,
        steps: List[dict],
        final_answer: str,
    ) -> bool:
        """
        Detect whether this chain is a math / computation task.

        Heuristic: keyword match on question + final answer, OR
        the final answer looks purely numeric.

        Returns True when math routing should activate.
        """
        combined = f"{question} {final_answer}"
        if _MATH_KW_RE.search(combined):
            return True
        # Final answer is a bare number → likely math
        if self._extract_number(final_answer) is not None:
            # Only treat as math if question also has a numeric flavour
            if re.search(r"\d", question):
                return True
        return False

    # ── Number extraction ───────────────────────────────────────────────────

    def _extract_number(self, text: str) -> Optional[float]:
        """
        Extract a single representative numeric value from text.

        Handles: integers, decimals, percentages (→ float), fractions,
        scientific notation, negative numbers, comma-separated thousands.

        Returns None when no number can be found.
        """
        if not text:
            return None
        text = text.strip()

        # 1. Try whole-text as a clean number (most common: "36", "3.14", "-7")
        cleaned = text.replace(",", "").rstrip("%").strip()
        try:
            val = float(cleaned)
            # percentage → divide by 100
            if "%" in text and not re.search(r"[+\-*/]", text):
                val /= 100.0
            return val
        except ValueError:
            pass

        # 2. Try fraction form "3/4"
        m = _FRAC_RE.search(text)
        if m:
            num, den = int(m.group(1)), int(m.group(2))
            if den != 0:
                return num / den

        # 3. Equation pattern "= 36" at the end of text
        m = _EQ_RE.search(text)
        if m:
            try:
                return float(m.group(1).replace(",", ""))
            except ValueError:
                pass

        # 4. First number found anywhere in text (last resort)
        nums = _NUM_RE.findall(text)
        for sign, digits in nums:
            try:
                val = float((sign + digits).replace(",", ""))
                if "%" in text[text.find(digits) : text.find(digits) + len(digits) + 2]:
                    val /= 100.0
                return val
            except ValueError:
                continue

        return None

    # ── Expression extraction from chain ──────────────────────────────────

    def _extract_expressions(self, steps: List[dict]) -> List[str]:
        """
        Extract mathematical expressions from step observations and thoughts.

        Returns a list of string expressions suitable for eval / SymPy.
        """
        exprs: List[str] = []
        for step in steps:
            for field in ("thought", "observation"):
                text = step.get(field, "") or ""
                # Look for "X = Y" patterns
                for m in re.finditer(r"([^=\n]+?)\s*=\s*([-\d./%+\-* ()]+)", text):
                    rhs = m.group(2).strip().rstrip(".")
                    if re.search(r"\d", rhs):
                        exprs.append(rhs)
                # Look for standalone arithmetic expressions
                for m in _ARITH_OP_RE.finditer(text):
                    exprs.append(m.group(0).strip())
        return exprs

    def _extract_expected_from_chain(
        self,
        question: str,
        steps: List[dict],
        final_answer: str,
    ) -> Optional[float]:
        """
        Determine the expected numeric answer by:
          1. Evaluating arithmetic expressions found in observations
          2. Using SymPy on symbolic expressions
          3. Directly computing from the question text

        Returns the numeric value we expect, or None on failure.
        """
        # Collect all candidate expressions from step observations / thoughts
        exprs = self._extract_expressions(steps)

        # Also try to build an expression directly from the question
        q_expr = self._build_expr_from_question(question)
        if q_expr:
            exprs.insert(0, q_expr)

        # Try each expression in priority order: Python exec, then SymPy
        for expr in exprs:
            val = self._safe_exec_python(expr)
            if val is not None:
                return val
            val = self._safe_eval_sympy(expr)
            if val is not None:
                return val

        # Last resort: look for "= <number>" in observations
        for step in steps:
            obs = step.get("observation", "") or ""
            m = _EQ_RE.search(obs)
            if m:
                try:
                    return float(m.group(1).replace(",", ""))
                except ValueError:
                    pass

        return None

    def _build_expr_from_question(self, question: str) -> Optional[str]:
        """
        Attempt to construct a computable Python expression from a plain-
        English math question.

        Handles common templates:
          "What is X% of Y?"   → X / 100 * Y
          "X + Y"              → passed through directly
          "X * Y"              → passed through

        Returns an expression string or None.
        """
        q = question.strip().rstrip("?")

        # "X percent of Y" / "X% of Y"
        m = re.search(
            r"(\d+\.?\d*)\s*(?:percent|%)\s+of\s+(\d+\.?\d*)",
            q, re.IGNORECASE
        )
        if m:
            pct, base = m.group(1), m.group(2)
            return f"{pct} / 100 * {base}"

        # "what is X + Y" / "calculate X * Y" etc.
        m = re.search(r"(\d[\d\s.]*[+\-*/^][\d\s.]*\d)", q)
        if m:
            return m.group(1).strip()

        # "square root of X"
        m = re.search(r"square\s+root\s+of\s+(\d+\.?\d*)", q, re.IGNORECASE)
        if m:
            return f"math.sqrt({m.group(1)})"

        # "X squared" / "X^2"
        m = re.search(r"(\d+\.?\d*)\s*(?:squared|\*\*\s*2|\^\s*2)", q, re.IGNORECASE)
        if m:
            return f"{m.group(1)} ** 2"

        return None

    # ── Safe evaluation ────────────────────────────────────────────────────

    def _safe_eval_sympy(self, expr: str) -> Optional[float]:
        """
        Safely evaluate expression with SymPy.

        Returns a float if successful, None on failure or if SymPy is absent.
        """
        if self._has_sympy is None:
            try:
                import sympy  # noqa: F401
                self._has_sympy = True
            except ImportError:
                self._has_sympy = False

        if not self._has_sympy:
            return None

        try:
            import sympy
            # Restrict to numeric expressions only (no free symbols)
            result = sympy.sympify(expr).evalf()
            if result.is_number:
                return float(result)
        except Exception:
            pass
        return None

    def _safe_exec_python(self, code: str, timeout_sec: Optional[float] = None) -> Optional[float]:
        """
        Execute a Python arithmetic expression in a subprocess.

        The subprocess wrapper imports only math and fractions, but there is NO
        OS-level isolation (no seccomp, no chroot, no namespace). If agent chain
        observations contain crafted text (e.g. ``__import__('os').system(...)``),
        that code will execute in the subprocess before the float() cast fails.

        Security guidance:
        - Use only on trusted agent chains from your own controlled agents.
        - For untrusted/third-party agent output, use ComprehensiveMathVerifier
          instead (SymPy lambdify/sympify only — no exec/eval path).
        - For true subprocess isolation, wrap MathVerifier in a container or
          use seccomp via a tool such as nsjail.

        Returns a float result or None on failure/timeout.
        """
        timeout = timeout_sec or self._timeout_sec

        # Normalise common math notation: "^" → "**", "×" → "*", "÷" → "/"
        code = code.replace("^", "**").replace("×", "*").replace("÷", "/")
        # Strip trailing non-expression chars
        code = code.strip().rstrip(".;,")
        if not code:
            return None

        wrapper = (
            "import math, fractions\n"
            "try:\n"
            f"    _r = {code}\n"
            "    print(float(_r))\n"
            "except Exception as _e:\n"
            "    print('ERROR')\n"
        )
        try:
            out = subprocess.run(
                ["python3", "-c", wrapper],
                capture_output=True,
                text=True,
                timeout=timeout,
            )
            val = out.stdout.strip()
            if val and val != "ERROR":
                return float(val)
        except (subprocess.TimeoutExpired, ValueError, OSError):
            pass
        return None

    # ── Helpers ────────────────────────────────────────────────────────────

    def _approx_equal(self, a: float, b: float) -> bool:
        """
        True if a and b are approximately equal within relative tolerance.
        Uses absolute tolerance of 0.001 when both values are near zero.
        """
        if abs(b) < 1e-9:
            return abs(a - b) < 1e-3
        return abs(a - b) / abs(b) <= self._tol
