import os
import gradio as gr
from dotenv import load_dotenv

from agents.openai_agent import query_openai
from data_loaders.load_datasets import load_all_datasets
from tools.math_tools import run_math_analysis

load_dotenv()

# Callback functions
def openai_response(prompt):
    return query_openai(prompt)

def run_tool(topic):
    return run_math_analysis(topic)

def load_datasets_button():
    datasets = load_all_datasets()
    return "✅ Datasets loaded successfully:\n" + "\n".join(datasets.keys())

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 AI Business Agent (Math + Models + Data)")

    with gr.Tab("1. Ask OpenAI"):
        q = gr.Textbox(label="Your question to GPT-4")
        a = gr.Textbox(label="Response", lines=10)
        q_btn = gr.Button("Ask")
        q_btn.click(openai_response, inputs=q, outputs=a)

    with gr.Tab("2. Run Math Tool"):
        tool = gr.Dropdown(
            ["linear_algebra", "probability", "optimization", "game_theory", "time_series", "markov"],
            label="Choose Math Topic"
        )
        tool_btn = gr.Button("Run")
        tool_out = gr.Textbox(label="Output", lines=10)
        tool_btn.click(run_tool, inputs=tool, outputs=tool_out)

    with gr.Tab("3. Load Datasets"):
        load_btn = gr.Button("Load All Datasets")
        load_out = gr.Textbox(label="Loaded Datasets")
        load_btn.click(load_datasets_button, outputs=load_out)

demo.launch()