import sympy as sp

class MathContext:
    """
    A sandbox for the LLM to execute mathematics.
    Each operation is recorded with Hebrew context and exact LaTeX representation.
    """
    def __init__(self):
        self.steps = []
        
        # Injected SymPy helpers
        self.sp = sp
        self.Eq = sp.Eq
        self.solve = sp.solve
        self.expand = sp.expand
        self.simplify = sp.simplify
        self.sqrt = sp.sqrt
        self.diff = sp.diff

    def _format_latex(self, expr) -> str:
        """Converts SymPy to precise LaTeX, without $$ so it renders correctly in Flutter."""
        if isinstance(expr, str):
            # If the LLM passed a string instead of SymPy object, try parsing it
            try:
                expr = sp.sympify(expr)
            except:
                pass # Return as-is if not parseable

        latex_str = sp.latex(expr)
        return latex_str

    def explain(self, text: str):
        """Adds a pure Hebrew explanation step without math."""
        self.steps.append({
            "content_mixed": text,
            "block_math": ""
        })

    def declare_equation(self, text: str, eq: sp.Eq):
        """Prints a known equation with its explanation."""
        self.steps.append({
            "content_mixed": text,
            "block_math": self._format_latex(eq)
        })
        return eq

    def expand_expr(self, text: str, expr):
        """Expands an expression completely (e.g. squaring a root)."""
        expanded = sp.expand(expr)
        self.steps.append({
            "content_mixed": text,
            "block_math": self._format_latex(expanded)
        })
        return expanded

    def solve_equation(self, text: str, eq: sp.Eq, var):
        """Solves an equation for a specific variable."""
        solutions = sp.solve(eq, var)
        
        # Format solutions nicely
        if isinstance(solutions, list):
            if len(solutions) == 1:
                math_result = sp.Eq(var, solutions[0])
            else:
                # E.g. x_1 = 2, x_2 = -2
                parts = [f"{sp.latex(var)}_{{{i+1}}} = {sp.latex(sol)}" for i, sol in enumerate(solutions)]
                math_result = " \\text{ ואו } ".join(parts)
        else:
            math_result = sp.Eq(var, solutions)

        self.steps.append({
            "content_mixed": text,
            "block_math": math_result if isinstance(math_result, str) else self._format_latex(math_result)
        })
        return solutions

    def finish(self, final_answer: str, teacher_summary: str = ""):
        """Sets the final human-readable answer for the UI."""
        self.final_answer = final_answer
        self.teacher_summary = teacher_summary

def run_llm_code(python_code: str) -> dict:
    """
    Executes the LLM-generated Python code in our secure MathContext.
    Returns the step-by-step UI format required by BuddyMath.
    """
    ctx = MathContext()
    
    # Secure Globals the LLM is allowed to use
    safe_globals = {
        "ctx": ctx,
        "x": sp.Symbol('x'),
        "y": sp.Symbol('y'),
        "a": sp.Symbol('a'),
        "b": sp.Symbol('b'),
        "c": sp.Symbol('c'),
        "m": sp.Symbol('m'),
        "R": sp.Symbol('R'),
        "sp": sp,
    }

    try:
        # Execute the LLM's dynamically generated mathematics
        exec(python_code, safe_globals)
        
        return {
            "success": True,
            "steps": ctx.steps,
            "final_answer": getattr(ctx, 'final_answer', "הגענו לפתרון."),
            "teacher_summary": getattr(ctx, 'teacher_summary', "")
        }
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }

if __name__ == "__main__":
    # Simulate LLM output to solve a Locus / Parabola problem:
    # "The distance from (x, y) to (2, 0) is equal to its distance to x = -2."
    
    llm_code = """
ctx.explain("נסמן את הנקודה הכללית על המקום הגיאומטרי כ- (x,y). לפי הנתון, המרחק מהמוקד שווה למרחק מהמדריך.")
d1 = sp.sqrt((x - 2)**2 + (y - 0)**2)  # מוקד
d2 = sp.sqrt((x - (-2))**2)          # מדריך

eq1 = ctx.declare_equation("המשוואה המשווה בין המרחקים היא:", ctx.Eq(d1, d2))

ctx.explain("נעלה את שני האגפים בריבוע כדי להיפטר מהשורש:")
# SymPy understands squaring both sides! We square them and declare equality.
squared_eq = ctx.Eq(d1**2, d2**2)
ctx.steps[-1]["block_math"] = ctx._format_latex(squared_eq) # Override previous block math

# Now expand and simplify it beautifully
expanded_eq = ctx.declare_equation("נרחיב את הביטויים (פתיחת סוגריים מלאה):", ctx.Eq(ctx.expand(squared_eq.lhs), ctx.expand(squared_eq.rhs)))

# SymPy's powerful simplify equation solver (subtract RHS from LHS)
simplified_expr = ctx.simplify(expanded_eq.lhs - expanded_eq.rhs)
final_eq = ctx.declare_equation("לאחר כינוס איברים והעברת אגפים, נקבל את צורת הפרבולה הפשוטה:", ctx.Eq(simplified_expr, 0))

# Also isolate y^2 if needed
y_sq_isolated = ctx.solve(final_eq, y**2)
if y_sq_isolated:
    ctx.declare_equation("נבודד את y^2 במשוואה:", ctx.Eq(y**2, y_sq_isolated[0]))

ctx.finish("$$ y^2 = 8x $$")
"""
    
    print("🚀 Running LLM Mathematics Script:")
    result = run_llm_code(llm_code)
    
    import json
    # Print the resulting UI object
    print(json.dumps(result, indent=2, ensure_ascii=False))