# 2) The actual app
import os
from getpass import getpass
from openai import OpenAI
import gradio as gr

# ——— Configure your OpenRouter key ———
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")

# Check if the API key was retrieved
if not OPENROUTER_API_KEY:
    print("Error: OPENROUTER_API_KEY not found in environment.")
    print("Please set your API key in the environment as 'OPENROUTER_API_KEY'.")
else:
    client = OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=OPENROUTER_API_KEY,
    )

    def openrouter_chat(user_message, history):
        """Send user_message and history to mistralai/devstral-small:free and append to history."""
        history = history or []

        # Build the messages list from the history and the current user message
        # Gradio history is a list of tuples [(user_msg, bot_msg), ...]
        # We need to convert it to the OpenAI API format: list of dicts [{"role": "user", "content": ...}, {"role": "assistant", "content": ...}, ...]
        messages_for_api = []
        for human_message, ai_message in history:
            messages_for_api.append({"role": "user", "content": human_message})
            if ai_message is not None: # Only add assistant message if it exists
                messages_for_api.append({"role": "assistant", "content": ai_message})

        # Add the current user message
        messages_for_api.append({"role": "user", "content": user_message})

        try:
            # Call the model with the mistralai/Devstral-Small-2505 for full conversation history
            resp = client.chat.completions.create(
                model="mistralai/devstral-small:free",
                messages=messages_for_api,
                # you can tweak max_tokens, temperature, etc. here
            )
            bot_reply = resp.choices[0].message.content
            # Append the user message and bot reply to the history for Gradio display
            history.append((user_message, bot_reply))
        except Exception as e:
            # Handle potential errors and append an error message to the history
            history.append((user_message, f"Error: {e}"))


        return history, ""

    with gr.Blocks() as demo:
        gr.Markdown("## 🦜🔗 Gradio + OpenRouter Chat with Devstral-Small (with History)")
        chatbot = gr.Chatbot(label="Chat")
        msg_in = gr.Textbox(placeholder="Type your question here…", label="You")
        msg_in.submit(openrouter_chat, inputs=[msg_in, chatbot], outputs=[chatbot, msg_in])

    demo.launch()