import os
import gradio as gr
from dotenv import load_dotenv
from openai import OpenAI
from prompts.main_prompt import MAIN_PROMPT, get_prompt_for_method
from prompts.initial_prompt import INITIAL_PROMPT

# ✅ Load API key from .env file
if os.path.exists(".env"):
    load_dotenv(".env")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

if not OPENAI_API_KEY:
    raise ValueError("🚨 OpenAI API key is missing! Set it in your .env file.")

# ✅ Correct OpenAI Client Initialization
client = OpenAI()

def respond(user_message, history):
    if not user_message:
        return "", history

    # ✅ Ensure proper message handling
    try:
        messages = [{"role": "system", "content": INITIAL_PROMPT}]

        for user_text, assistant_text in history:
            if user_text:
                messages.append({"role": "user", "content": user_text})
            if assistant_text:
                messages.append({"role": "assistant", "content": assistant_text})

        # ✅ Handling Method Selection
        method_selection = user_message.lower().strip()
        method_prompt = get_prompt_for_method(method_selection)
        
        if method_prompt != "I didn’t understand your choice. Please type 'Bar Model,' 'Double Number Line,' or 'Equation' to proceed.":
            messages.append({"role": "assistant", "content": method_prompt})
            history.append((user_message, method_prompt))
            return "", history

        messages.append({"role": "user", "content": user_message})

        # ✅ Get AI response (Updated call to OpenAI API)
        completion = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            max_tokens=512,
            temperature=0.7
        )

        assistant_reply = completion.choices[0].message.content
        history.append((user_message, assistant_reply))

        return "", history

    except Exception as e:
        return f"⚠️ An error occurred: {str(e)}", history

# ✅ Gradio UI Setup
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 AI-Guided Math PD Chatbot")

    chatbot = gr.Chatbot(value=[("", INITIAL_PROMPT)], height=500)
    state_history = gr.State([("", INITIAL_PROMPT)])

    user_input = gr.Textbox(placeholder="Type your message here...", label="Your Input")

    user_input.submit(
        respond,
        inputs=[user_input, state_history],
        outputs=[user_input, chatbot]
    ).then(
        fn=lambda _, h: h,
        inputs=[user_input, chatbot],
        outputs=[state_history]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)