captcha_solver/solvers/math_solver.py

"""Math captcha solver.

Strategy order (fast -> slow, accurate -> less):
  1. Pure-Python regex on OCR output (sub-millisecond)
  2. Tiny LLM (Qwen2.5 1.5B) for cleanup if OCR returns garbled math
  3. Ollama text model if enabled

For the OCR step we use Florence-2 (or fall back to Tesseract if
available). The math string is then evaluated safely.
"""

from __future__ import annotations

import ast
import operator
import re
from typing import Optional

from captcha_solver.solvers.base import BaseSolver, SolveAttempt
from captcha_solver.utils.image import decode_base64_image, image_to_pil


_SAFE_OPS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Mod: operator.mod,
    ast.Pow: operator.pow,
    ast.USub: operator.neg,
    ast.UAdd: operator.pos,
}


def _safe_eval(expr: str) -> Optional[float]:
    """Evaluate a math expression with only + - * / % ** and parens."""
    if not expr or not re.match(r"^[\d\s+\-*/%().]+$", expr):
        return None
    try:
        tree = ast.parse(expr, mode="eval")
        return _eval_node(tree.body)
    except Exception:
        return None


def _eval_node(node):
    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
        return node.value
    if isinstance(node, ast.BinOp):
        op = _SAFE_OPS.get(type(node.op))
        if op is None:
            raise ValueError("op not allowed")
        return op(_eval_node(node.left), _eval_node(node.right))
    if isinstance(node, ast.UnaryOp):
        op = _SAFE_OPS.get(type(node.op))
        if op is None:
            raise ValueError("op not allowed")
        return op(_eval_node(node.operand))
    raise ValueError("unsupported")


_WORD_TO_NUM = {
    "zero": 0, "one": 1, "two": 2, "three": 3, "four": 4,
    "five": 5, "six": 6, "seven": 7, "eight": 8, "nine": 9,
    "ten": 10, "eleven": 11, "twelve": 12, "thirteen": 13,
    "fourteen": 14, "fifteen": 15, "sixteen": 16, "seventeen": 17,
    "eighteen": 18, "nineteen": 19, "twenty": 20, "thirty": 30,
    "forty": 40, "fifty": 50, "hundred": 100,
}


def _words_to_expr(text: str) -> str:
    """Convert 'three plus five' -> '3+5'."""
    s = text.lower()
    s = s.replace("plus", "+").replace("add", "+").replace("+", " + ")
    s = s.replace("minus", "-").replace("subtract", "-")
    s = s.replace("times", "*").replace("multiplied by", "*").replace("x", "*")
    s = s.replace("divided by", "/").replace("over", "/")
    s = s.replace("equals", "=").replace("is", "")
    tokens: list[str] = []
    for w in re.findall(r"[a-z]+|\d+|[+\-*/()]", s):
        if w in _WORD_TO_NUM:
            tokens.append(str(_WORD_TO_NUM[w]))
        else:
            tokens.append(w)
    return " ".join(tokens)


def _normalize_ocr_math(s: str) -> str:
    s = s.strip()
    s = s.replace("×", "*").replace("x", "*").replace("X", "*")
    s = s.replace("÷", "/").replace("−", "-")
    # Strip trailing "=..." (handles "=?", "= 7", "=42", "= ?")
    s = re.sub(r"=.*$", "", s)
    # Strip OCR garbage that looks like a single trailing digit (often "?" read as 6/2/7)
    # Only strip if it's a single digit at the very end AND not part of an expression
    s_no_space = s.replace(" ", "")
    # If the result still ends with a lone digit that doesn't follow an operator, leave it
    # (e.g. "3+5=2" -> "3+5", but "12+8=?" -> "12+8" since ? is stripped first)
    s = s.replace("?", "")
    s = s.replace(" ", "")
    return s


def _extract_math_expression(s: str) -> Optional[str]:
    s = _normalize_ocr_math(s)
    # Try to find a chain of at least one operator between numbers
    m = re.search(r"(-?\d+(?:\.\d+)?\s*[+\-*/%]\s*-?\d+(?:\.\d+)?(?:\s*[+\-*/%]\s*-?\d+(?:\.\d+)?)*)", s)
    if m:
        return m.group(1)
    return None


class MathSolver(BaseSolver):
    name = "math"
    captcha_type = "math"

    def attempts(self):
        return [
            self._solve_regex,
            self._solve_tesseract_fallback,
            self._solve_ollama_if_enabled,
        ]

    def _solve_regex(self) -> SolveAttempt:
        """Pure regex path: try to extract and eval the math string from the raw input.

        Note: when the input is an image, we need OCR first. We delegate to
        a small helper that uses Florence-2 and a cleanup LLM if needed.
        """
        raw = self._raw_text
        if not raw:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.regex", error="no text input")

        for src in (raw, _words_to_expr(raw)):
            expr = _extract_math_expression(src)
            if not expr:
                continue
            val = _safe_eval(expr)
            if val is None:
                continue
            conf = 0.92 if src is raw else 0.8
            return SolveAttempt(answer=_format_int(val), confidence=conf, solver_name="math.regex")
        return SolveAttempt(answer="", confidence=0.0, solver_name="math.regex", error="no math found")

    def _solve_tesseract_fallback(self) -> SolveAttempt:
        """Use Tesseract OCR when Florence-2 isn't available.

        Useful when HF models can't be downloaded (expired token, offline).
        """
        if not self._pil_image:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.tesseract", error="no image")
        try:
            from captcha_solver.utils.tesseract_ocr import ocr_captcha_math
            text = ocr_captcha_math(self._pil_image)
        except Exception as exc:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.tesseract", error=str(exc))
        if not text:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.tesseract", error="empty OCR")
        expr = _extract_math_expression(text)
        if not expr:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.tesseract", error=f"no math in: {text!r}")
        val = _safe_eval(expr)
        if val is None:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.tesseract", error=f"eval failed: {expr}")
        return SolveAttempt(answer=_format_int(val), confidence=0.6, solver_name="math.tesseract")

    def _solve_ollama_if_enabled(self) -> SolveAttempt:
        if not self.ctx.ollama.enabled:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.ollama", error="ollama disabled")
        if not self._raw_text:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.ollama", error="no text")
        out = self.ctx.ollama.generate_text(
            f"Extract and solve the math expression. Return ONLY the final integer answer.\nExpression: {self._raw_text}",
            system="You solve simple arithmetic. Reply with one integer, nothing else.",
            max_tokens=8,
        )
        m = re.search(r"-?\d+", out)
        if not m:
            return SolveAttempt(answer="", confidence=0.0, solver_name="math.ollama", error="no number in output")
        return SolveAttempt(answer=m.group(0), confidence=0.88, solver_name="math.ollama")

    def prepare(self, image_b64: Optional[str], audio_b64: Optional[str], hint: Optional[str]) -> None:
        """OCR the image to text (or accept a pre-OCR'd hint).

        Stores the result in self._raw_text and keeps the PIL image in
        self._pil_image for the tesseract fallback.
        """
        self._pil_image = None
        if hint:
            self._raw_text = hint
            return
        if not image_b64:
            self._raw_text = ""
            return
        try:
            data = decode_base64_image(image_b64)
            img = image_to_pil(data)
            self._pil_image = img
        except Exception as exc:
            self._raw_text = ""
            self._last_error = f"decode: {exc}"
            return
        try:
            raw = self.ctx.florence.ocr(img, task="<OCR>")
            self._raw_text = _clean_ocr_text(raw)
        except Exception as exc:
            self._raw_text = ""
            self._last_error = f"florence: {exc}"

    def __init__(self, ctx) -> None:
        super().__init__(ctx)
        self._raw_text: str = ""
        self._pil_image = None
        self._last_error: str = ""


def _clean_ocr_text(s: str) -> str:
    s = s.strip()
    s = s.replace("\n", " ").replace("\r", " ")
    s = re.sub(r"\s+", " ", s)
    return s


def _format_int(v: float) -> str:
    if abs(v - round(v)) < 1e-9:
        return str(int(round(v)))
    return ("%g" % v)