import gradio as gr
from openai import OpenAI
import time

# Set a constant temperature for all API calls
TEMPERATURE = 0.0
MODEL = "gpt-3.5-turbo"

def break_down_problem(client, problem_description):
    prompt = f"Break down the following problem into small, solvable tasks. Do not attempt to solve the problem. Problem:\n\n{problem_description}"
    
    response = client.chat.completions.create(
        model= MODEL,
        messages=[
            {"role": "system", "content": "You are a helpful assistant that breaks down complex problems into smaller, manageable tasks."},
            {"role": "user", "content": prompt}
        ],
        temperature=TEMPERATURE
    )
    
    return response.choices[0].message.content

def solve_problem(api_key, problem_description):
    client = OpenAI(api_key=api_key)
    
    # Step 1: Break down the problem
    tasks = break_down_problem(client, problem_description)
    
    # Step 2: Solve each task
    solutions = []
    for task in tasks.split('\n'):
        if task.strip():
            prompt = f"Solve the following task, not the whole context problem.:\n\n{task}\n\nUse the following context if needed:\n{problem_description}"
            
            response = client.chat.completions.create(
                model= MODEL,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant that solves problems step by step."},
                    {"role": "user", "content": prompt}
                ],
                temperature=TEMPERATURE
            )
            
            solutions.append(f"Task: {task}\nSolution: {response.choices[0].message.content}\n")
    
    # Step 3: Compile final answer
    final_prompt = f"Given the following problem and its step-by-step solutions, provide a final answer and explanation:\n\nProblem:\n{problem_description}\n\nStep-by-step solutions:\n{''.join(solutions)}"
    
    final_response = client.chat.completions.create(
        model= MODEL,
        messages=[
            {"role": "system", "content": "You are a helpful assistant that provides final answers and explanations based on step-by-step solutions."},
            {"role": "user", "content": final_prompt}
        ],
        temperature=TEMPERATURE
    )
    
    return f"Problem Breakdown:\n\n{tasks}\n\nStep-by-step Solutions:\n\n{''.join(solutions)}\n\nFinal Answer and Explanation:\n\n{final_response.choices[0].message.content}"

def gradio_interface(api_key, problem_description):
    try:
        return solve_problem(api_key, problem_description)
    except Exception as e:
        return f"An error occurred: {str(e)}"

iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="OpenAI API Key", type="password"),
        gr.Textbox(label="Problem Description", lines=10)
    ],
    outputs=gr.Textbox(label="Solution", lines=20),
    title="OpenAI Problem Solver",
    description="Enter your OpenAI API key and a problem description to get a step-by-step solution."
)

iface.launch()