import os
import json
import logging
import re
import random
import time

logger = logging.getLogger(__name__)

# Standardized math utility
from math_utils import clean_latex

def _extract_numbers(text: str):
    return [float(x) for x in re.findall(r'-?\d+\.?\d*', text)]

def _symbolic_solve(eq: str):
    """
    Expert-level symbolic solver: 
    1. Evaluates truth statements (no variables)
    2. Solves linear/quadratic/polynomial equations
    3. Handles multi-root solutions correctly
    """
    try:
        from sympy import symbols, solve, sympify
        if '=' not in eq:
            return None
        
        lhs, rhs = eq.split('=', 1)
        expr = sympify(lhs.strip()) - sympify(rhs.strip())
        vars = list(expr.free_symbols)
        
        if not vars:
            # Truth statement check
            return "True" if expr == 0 else "False"
        
        # Solving for the primary variable (usually 'x')
        x = symbols('x')
        if x in vars:
            sol = solve(expr, x)
            if sol:
                if len(sol) > 1:
                    return ', '.join(str(s) for s in sorted(sol))
                return str(sol[0])
        else:
            # Fallback to solving for whatever variable is present
            sol = solve(expr, vars[0])
            if sol:
                return str(sol[0])
    except: pass
    return None

def _smart_solve(problem: str):
    from sympy import sympify
    clean = clean_latex(problem)

    # 1. Symbolic Equation/Truth Logic
    if '=' in clean:
        result = _symbolic_solve(clean)
        if result:
            return result, [f"Symbolic Evaluation: {clean}", f"Result: {result}"]

    # 2. Complex Arithmetic (e.g. 100 * 20 / 5)
    try:
        # Strict arithmetic check: allows digits, operators, parens
        if re.match(r'^[0-9\+\-\*\/\.\s\(\)\^]+$', clean):
            ans = sympify(clean.replace('^', '**'))
            if ans.is_number:
                res = str(int(ans) if ans == int(ans) else round(float(ans), 4))
                return res, [f"Arithmetic Calculation: {clean}", f"Result: {res}"]
    except: pass

    return None, []

    return None, []


class LLMAgent:
    """Multi-Agent Reasoning Engine with Smart Simulation + Gemini API support."""
    AGENT_STYLES = {
        "GPT-4": ("step_by_step", 0.0),
        "Llama 3": ("chain_of_thought", 0.05),
        "Gemini 2.0 Pro": ("direct_solve", 0.0),
        "Qwen-2.5-Math-7B": ("formal_proof", 0.08),
    }

    def __init__(self, model_name: str, use_real_api: bool = False):
        self.model_name = model_name
        self.use_real_api = use_real_api
        self.client = None
        
        if self.use_real_api:
            GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
            if GEMINI_API_KEY:
                try:
                    import google.generativeai as genai
                    genai.configure(api_key=GEMINI_API_KEY)
                    self.client = genai.GenerativeModel('gemini-2.0-flash')
                    print(f"[{model_name}] Live Gemini API enabled.")
                except Exception as e:
                    logger.warning(f"[{model_name}] Gemini init failed: {e}")
            else:
                self.use_real_api = False

    def generate_solution(self, problem: str) -> dict:
        # Expert ML Patch: Clean input early to prevent CJK leakage to APIs
        problem = clean_latex(problem)
        
        if self.use_real_api and self.client:
            return self._call_real_gemini(problem)
        return self._simulate_agent(problem)

    def _call_real_gemini(self, problem: str) -> dict:
        prompt = f"""You are a mathematical reasoning agent in the MVM2 framework.
Solve EXACTLY: {problem}

Strictly output JSON:
{{
    "final_answer": "...",
    "reasoning_trace": ["step 1", "step 2"],
    "confidence_explanation": "..."
}}
"""
        try:
            response = self.client.generate_content(prompt)
            return json.loads(response.text.replace("```json", "").replace("```", "").strip())
        except:
            return self._simulate_agent(problem)

    def _simulate_agent(self, problem: str) -> dict:
        time.sleep(random.uniform(0.1, 0.4))
        style, error_rate = self.AGENT_STYLES.get(self.model_name, ("generic", 0.0))

        correct_answer, reasoning_steps = _smart_solve(problem)

        if correct_answer is None:
            nums = _extract_numbers(problem)
            if nums:
                n = nums[0]
                if style == "step_by_step":
                    correct_answer = str(int(n * 2) if (n * 2) == int(n * 2) else round(n * 2, 4))
                    reasoning_steps = [f"Identify value: {n}", f"Double: {n} × 2 = {correct_answer}"]
                elif style == "chain_of_thought":
                    correct_answer = str(int(n + 1) if (n + 1) == int(n + 1) else round(n + 1, 4))
                    reasoning_steps = [f"Observe value: {n}", f"Increment: {n} + 1 = {correct_answer}"]
                elif style == "direct_solve":
                    correct_answer = str(int(n) if n == int(n) else round(n, 4))
                    reasoning_steps = [f"Direct evaluation of {n}", f"Result: {correct_answer}"]
                else: 
                    correct_answer = str(int(n - 1) if (n - 1) == int(n - 1) else round(n - 1, 4))
                    reasoning_steps = [f"Formal derivation for {n}", f"Theorem: result = n - n = {correct_answer}"]
            else:
                correct_answer = "Unable to determine"
                reasoning_steps = ["Problem could not be parsed", "Insufficient mathematical context"]

        final_answer = correct_answer
        is_hallucinating = False
        if random.random() < error_rate:
            try:
                # Basic error injection
                f_ans = float(correct_answer.split(',')[0])
                wrong = f_ans + 1.0
                final_answer = str(int(wrong) if wrong == int(wrong) else round(wrong, 4))
                reasoning_steps[-1] = f"[Divergence] Arithmetic deviation: {final_answer}"
                is_hallucinating = True
            except: pass

        if is_hallucinating:
            confidence = f"[{self.model_name}] Divergent reasoning detected."
        else:
            confidence = f"[{self.model_name}] {style} reasoning applied with high confidence."

        return {
            "final_answer": final_answer,
            "reasoning_trace": reasoning_steps,
            "confidence_explanation": confidence
        }