import gradio as gr
from daggr import FnNode, InferenceNode, Graph
import random

# Generate math problems with increasing difficulty
def generate_math_problem(level):
    if level <= 3:
        # Addition problems
        a = random.randint(1, 10 * level)
        b = random.randint(1, 10 * level)
        return f"{a} + {b} = ?", a + b
    elif level <= 6:
        # Subtraction problems
        a = random.randint(10, 20 * level)
        b = random.randint(1, a)
        return f"{a} - {b} = ?", a - b
    elif level <= 9:
        # Multiplication problems
        a = random.randint(1, 5 * level)
        b = random.randint(1, 5)
        return f"{a} × {b} = ?", a * b
    else:
        # Division problems
        b = random.randint(1, 5)
        a = b * random.randint(1, 10 * level)
        return f"{a} ÷ {b} = ?", a // b

# Node to generate new math problems
problem_generator = FnNode(
    fn=lambda level: generate_math_problem(level),
    inputs={"level": gr.Number(label="Current Level", value=1)},
    outputs={
        "problem": gr.Textbox(label="Math Problem"),
        "correct_answer": gr.Number(label="Correct Answer", visible=False)
    }
)

# Node for user to input their answer
user_input = FnNode(
    fn=lambda problem: problem,
    inputs={"problem": gr.Textbox(visible=False)},
    outputs={"problem_display": gr.Textbox(label="Your Problem")}
)

# Node for AI to solve the problem
ai_solver = InferenceNode(
    model="meta-llama/Llama-3.1-8B-Instruct",
    inputs={
        "prompt": gr.Textbox(visible=False),
        "problem": gr.Textbox(visible=False)
    },
    postprocess=lambda original, inputs: {
        "ai_response": original[0]["generated_text"].split("=")[-1].strip()
    },
    outputs={"ai_response": gr.Textbox(label="AI's Answer", visible=False)},
    system_prompt="Solve this math problem. Provide ONLY the numerical answer with no explanation or formatting. Problem: {problem}"
)

# Node to check answers and update game state
def check_answers(user_answer, ai_response, correct_answer, level, user_score, ai_score):
    try:
        user_correct = int(user_answer) == correct_answer
    except:
        user_correct = False
    
    try:
        ai_correct = int(ai_response) == correct_answer
    except:
        ai_correct = False
    
    # Update scores
    user_score += 1 if user_correct else 0
    ai_score += 1 if ai_correct else 0
    
    # Increase difficulty if both got it right
    level += 1 if user_correct and ai_correct else 0
    
    # Prepare feedback message
    feedback = f"**Level {level} Results:**\n\n"
    feedback += f"✅ Your answer ({user_answer}) was {'correct!' if user_correct else 'wrong.'}\n"
    feedback += f"🤖 AI's answer ({ai_response}) was {'correct!' if ai_correct else 'wrong.'}\n\n"
    feedback += f"Score: You {user_score} - {ai_score} AI\n"
    
    if level > 10:
        feedback += "\n🎉 Congratulations! You've completed all levels!"
        level = 1  # Reset game
    
    return {
        "feedback": feedback,
        "new_level": level,
        "new_user_score": user_score,
        "new_ai_score": ai_score
    }

answer_checker = FnNode(
    fn=check_answers,
    inputs={
        "user_answer": gr.Textbox(label="Your Answer"),
        "ai_response": gr.Textbox(visible=False),
        "correct_answer": gr.Number(visible=False),
        "level": gr.Number(visible=False),
        "user_score": gr.Number(visible=False, value=0),
        "ai_score": gr.Number(visible=False, value=0)
    },
    outputs={
        "feedback": gr.Markdown(),
        "new_level": gr.Number(visible=False),
        "new_user_score": gr.Number(visible=False),
        "new_ai_score": gr.Number(visible=False)
    }
)

# Create the game workflow
graph = Graph(
    name="Math Challenge: Human vs AI",
    nodes=[problem_generator, user_input, ai_solver, answer_checker],
    connections=[
        (problem_generator.outputs["problem"], user_input.inputs["problem"]),
        (problem_generator.outputs["problem"], ai_solver.inputs["problem"]),
        (problem_generator.outputs["correct_answer"], answer_checker.inputs["correct_answer"]),
        (problem_generator.outputs["level"], answer_checker.inputs["level"]),
        (user_input.outputs["problem_display"], None),  # Display to user
        (ai_solver.outputs["ai_response"], answer_checker.inputs["ai_response"]),
        (answer_checker.outputs["feedback"], None),  # Display feedback
        (answer_checker.outputs["new_level"], problem_generator.inputs["level"])
    ],
    inputs=[problem_generator.inputs["level"]],
    outputs=[user_input.outputs["problem_display"], 
             answer_checker.outputs["feedback"]]
)

# Launch the game
graph.launch()