import gradio as gr

from config import DEFAULT_PERSONA
from prompts import PERSONAS
from utils import generate_response


def handle_message(persona, user_input, history):
    """Process a chat message and return updated history.

    Args:
        persona: Selected chat persona
        user_input: User's message
        history: Previous chat messages

    Returns:
        Updated chat history with new messages
    """
    # Validate inputs

    if not persona or persona not in PERSONAS:
        return history + [
            {"role": "assistant", "content": "Please select a valid persona"}
        ]

    if not user_input.strip():
        return history + [{"role": "assistant", "content": "Please enter a message"}]

    try:
        # Get response from model
        return generate_response(PERSONAS[persona], user_input, history)
    except Exception as e:
        # Return error message in proper format
        return history + [{"role": "assistant", "content": f"Error: {str(e)}"}]


with gr.Blocks() as demo:
    gr.Markdown("## 🤖 LLMates: Persona-based Chat Assistant")

    persona = gr.Dropdown(
        choices=list(PERSONAS.keys()), value=DEFAULT_PERSONA, label="Choose Persona"
    )
    chatbox = gr.Chatbot(label="LLMates", type="messages")
    msg = gr.Textbox(label="Type your message...")
    chat_history = gr.State([])

    def user_submit(user_input, history, persona):
        if not user_input.strip():
            return (
                [],
                history,
                "",
            )  # Return empty chatbox update but keep history and clear input

        if not persona or persona not in PERSONAS:
            return (
                [],
                history
                + [{"role": "assistant", "content": "Please select a valid persona"}],
                "",
            )

        try:
            # Get response from model
            new_history = generate_response(PERSONAS[persona], user_input, history)
            return (
                new_history,
                new_history,
                "",
            )  # Clear input after successful submission
        except Exception as e:
            return (
                [],
                history + [{"role": "assistant", "content": f"Error: {str(e)}"}],
                "",
            )

    def clear_history():
        """Clear the chat history and reset message input when persona changes"""
        return ([], [], "")  # Clear chatbox, history state, and message input

    # Clear chat history and reset message input when persona changes
    persona.change(clear_history, outputs=[chatbox, chat_history, msg])

    # Handle message submission
    msg.submit(user_submit, [msg, chat_history, persona], [chatbox, chat_history, msg])

# main function
if __name__ == "__main__":
    demo.launch()