from __future__ import annotations

import math
import re
from statistics import mean, median
from typing import Dict, Optional

try:
    import sympy as sp
except Exception:
    sp = None

from models import SolverResult
from utils import clean_math_text, normalize_spaces


def extract_choices(text: str) -> Dict[str, str]:
    text = text or ""
    matches = list(
        re.finditer(
            r"(?i)\b([A-E])[\)\.:]\s*(.*?)(?=\s+\b[A-E][\)\.:]\s*|$)",
            text,
        )
    )
    return {m.group(1).upper(): normalize_spaces(m.group(2)) for m in matches}


def has_answer_choices(text: str) -> bool:
    return len(extract_choices(text)) >= 3


def is_quant_question(text: str) -> bool:
    lower = clean_math_text(text).lower()
    keywords = [
        "solve", "equation", "percent", "ratio", "probability", "mean", "median",
        "average", "sum", "difference", "product", "quotient", "triangle", "circle",
        "rectangle", "area", "perimeter", "volume", "algebra", "integer", "divisible",
        "number", "fraction", "decimal", "geometry", "distance", "speed", "work",
    ]
    if any(k in lower for k in keywords):
        return True
    if "=" in lower and re.search(r"[a-z]", lower):
        return True
    if re.search(r"\d", lower) and ("?" in lower or has_answer_choices(lower)):
        return True
    return False


def _prepare_expression(expr: str) -> str:
    expr = clean_math_text(expr).strip()
    expr = expr.replace("^", "**")
    expr = re.sub(r"(\d)\s*\(", r"\1*(", expr)
    expr = re.sub(r"\)\s*(\d)", r")*\1", expr)
    expr = re.sub(r"(\d)([a-zA-Z])", r"\1*\2", expr)
    return expr


def _extract_equation(text: str) -> Optional[str]:
    cleaned = clean_math_text(text)
    if "=" not in cleaned:
        return None
    patterns = [
        r"([A-Za-z0-9\.\+\-\*/\^\(\)\s]*[a-zA-Z][A-Za-z0-9\.\+\-\*/\^\(\)\s]*=[A-Za-z0-9\.\+\-\*/\^\(\)\s]+)",
        r"([0-9A-Za-z\.\+\-\*/\^\(\)\s]+=[0-9A-Za-z\.\+\-\*/\^\(\)\s]+)",
    ]
    for pattern in patterns:
        for m in re.finditer(pattern, cleaned):
            candidate = m.group(1).strip()
            tokens = re.findall(r"[a-z]", candidate.lower())
            if tokens and not candidate.lower().startswith(("how do", "can you", "please", "what is", "solve ")):
                return candidate
    eq_index = cleaned.find("=")
    left = re.findall(r"[A-Za-z0-9\.\+\-\*/\^\(\)\s]+$", cleaned[:eq_index])
    right = re.findall(r"^[A-Za-z0-9\.\+\-\*/\^\(\)\s]+", cleaned[eq_index + 1:])
    if left and right:
        candidate = left[0].strip().split()[-1] + " = " + right[0].strip().split()[0]
        if re.search(r"[a-z]", candidate.lower()):
            return candidate
    return None


def _parse_number(text: str) -> Optional[float]:
    raw = clean_math_text(text).strip().lower()
    pct = re.fullmatch(r"(-?\d+(?:\.\d+)?)%", raw.replace(" ", ""))
    if pct:
        return float(pct.group(1)) / 100.0
    frac = re.fullmatch(r"(-?\d+)\s*/\s*(-?\d+)", raw)
    if frac:
        den = float(frac.group(2))
        if den == 0:
            return None
        return float(frac.group(1)) / den
    try:
        return float(eval(_prepare_expression(raw), {"__builtins__": {}}, {"sqrt": math.sqrt, "pi": math.pi}))
    except Exception:
        return None


def _best_choice(answer_value: float, choices: Dict[str, str]) -> Optional[str]:
    best_letter = None
    best_diff = float("inf")
    for letter, raw in choices.items():
        parsed = _parse_number(raw)
        if parsed is None:
            continue
        diff = abs(parsed - answer_value)
        if diff < best_diff:
            best_diff = diff
            best_letter = letter
    if best_letter is not None and best_diff <= 1e-6:
        return best_letter
    return None


def _solve_percent(text: str) -> Optional[SolverResult]:
    lower = clean_math_text(text).lower()
    choices = extract_choices(text)

    m = re.search(r"(\d+(?:\.\d+)?)\s*(?:%|percent)\s+of\s+(?:a\s+)?number\s+is\s+(\d+(?:\.\d+)?)", lower)
    if m:
        p = float(m.group(1))
        value = float(m.group(2))
        ans = value / (p / 100.0)
        return SolverResult(
            domain="quant",
            solved=True,
            topic="percent",
            answer_value=f"{ans:g}",
            answer_letter=_best_choice(ans, choices) if choices else None,
            internal_answer=f"{ans:g}",
            steps=[
                f"Let the number be n.",
                f"Write {p}% of n as {p/100:g}n.",
                f"Set {p/100:g}n = {value} and solve for n.",
            ],
        )

    m = re.search(r"what is\s+(\d+(?:\.\d+)?)\s*(?:%|percent)\s+of\s+(\d+(?:\.\d+)?)", lower)
    if m:
        p = float(m.group(1))
        n = float(m.group(2))
        ans = p / 100.0 * n
        return SolverResult(
            domain="quant",
            solved=True,
            topic="percent",
            answer_value=f"{ans:g}",
            answer_letter=_best_choice(ans, choices) if choices else None,
            internal_answer=f"{ans:g}",
            steps=[f"Convert {p}% to {p/100:g}.", f"Multiply by {n}."]
        )
    return None


def _solve_mean_median(text: str) -> Optional[SolverResult]:
    lower = clean_math_text(text).lower()
    nums = [float(n) for n in re.findall(r"-?\d+(?:\.\d+)?", lower)]
    if not nums:
        return None
    if "mean" in lower or "average" in lower:
        ans = mean(nums)
        return SolverResult(domain="quant", solved=True, topic="statistics", answer_value=f"{ans:g}", internal_answer=f"{ans:g}", steps=["Add the values.", f"Divide by {len(nums)}."])
    if "median" in lower:
        ans = median(nums)
        return SolverResult(domain="quant", solved=True, topic="statistics", answer_value=f"{ans:g}", internal_answer=f"{ans:g}", steps=["Order the values.", "Take the middle value."])
    return None


def _solve_linear_equation(text: str) -> Optional[SolverResult]:
    if sp is None:
        return None
    expr = _extract_equation(text)
    if not expr:
        return None
    try:
        lhs, rhs = expr.split("=", 1)
        symbols = sorted(set(re.findall(r"\b[a-z]\b", expr)))
        if not symbols:
            return None
        var_name = symbols[0]
        var = sp.symbols(var_name)
        sol = sp.solve(sp.Eq(sp.sympify(_prepare_expression(lhs)), sp.sympify(_prepare_expression(rhs))), var)
        if not sol:
            return None
        value = sol[0]
        try:
            as_float = float(value)
        except Exception:
            as_float = None
        choices = extract_choices(text)
        return SolverResult(
            domain="quant",
            solved=True,
            topic="algebra",
            answer_value=str(value),
            answer_letter=_best_choice(as_float, choices) if (as_float is not None and choices) else None,
            internal_answer=f"{var_name} = {value}",
            steps=[
                "Treat the statement as an equation.",
                "Undo operations on both sides to isolate the variable.",
                f"That gives {var_name} = {value}.",
            ],
        )
    except Exception:
        return None


def solve_quant(text: str) -> SolverResult:
    text = text or ""
    for fn in (_solve_percent, _solve_mean_median, _solve_linear_equation):
        result = fn(text)
        if result is not None:
            return result
    return SolverResult(
        domain="quant",
        solved=False,
        topic="general_quant",
        reply="This looks quantitative, but it does not match a strong rule-based pattern yet.",
        steps=[
            "Identify the quantity the question wants.",
            "Translate the wording into an equation, ratio, or diagram.",
            "Carry out the calculation carefully.",
        ],
    )