import gradio as gr
import google.generativeai as genai
from google.api_core.exceptions import InvalidArgument

# Global variables
model = None
chat_session = None

# Function to set and verify API key
def set_api_key(user_key):
    global model, chat_session
    try:
        # Configure Gemini model
        genai.configure(api_key=user_key)
        model = genai.GenerativeModel("models/gemini-1.5-flash-latest")
        chat_session = model.start_chat()
        
        # Test key by sending harmless ping
        _ = model.generate_content("Hello")
        
        return (
            gr.update(visible=False),               # Hide API input
            gr.update(visible=True),                # Show chat UI
            gr.update(value="", visible=False),     # Hide error
            []                                      # Reset state
        )
    except InvalidArgument:
        return (
            gr.update(visible=True),                # Keep API input visible
            gr.update(visible=False),               # Keep chat UI hidden
            gr.update(value="❌ Invalid API Key", visible=True),  # Show error
            []
        )
    except Exception as e:
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(value=f"❌ Error: {str(e)}", visible=True),
            []
        )

# Chat function
def chat_with_gemini(user_input, history):
    global chat_session
    response = chat_session.send_message(user_input)
    history.append({"role": "user", "content": user_input})
    history.append({"role": "assistant", "content": response.text})
    return history, history, ""  # clear textbox

# Quit function
def handle_quit(history):
    history.append({"role": "user", "content": "quit"})
    history.append({"role": "assistant", "content": "Goodbye! 👋"})
    return history, gr.update(interactive=False)

# Reset chat
def reset_chat():
    global chat_session
    if model:
        chat_session = model.start_chat()
    return [], gr.update(interactive=True)

# UI
with gr.Blocks() as demo:
    gr.Markdown("![Mstar Logo](https://huggingface.co/spaces/Yatheshr/Mstar_Chatbot/resolve/main/mstar_chat.JPG)")
    gr.Markdown("## Mstar Chatbot")

    # API key section
    api_key_box = gr.Textbox(label="🔑 Enter Gemini API Key", type="password", placeholder="Paste your Gemini API Key here…")
    api_error_box = gr.Textbox(visible=False, interactive=False, show_label=False)
    api_submit_btn = gr.Button("Start Chat")

    # Chat UI (hidden until key is valid)
    with gr.Column(visible=False) as main_ui:
        chatbot = gr.Chatbot(label="Chat History", type="messages")
        with gr.Row():
            msg = gr.Textbox(label="You:", placeholder="Type your message here…")
        with gr.Row():
            send_btn = gr.Button("Send")
            quit_btn = gr.Button("Quit")

    state = gr.State([])

    # Events
    api_submit_btn.click(set_api_key, inputs=api_key_box, outputs=[api_key_box, main_ui, api_error_box, state])
    send_btn.click(chat_with_gemini, inputs=[msg, state], outputs=[chatbot, state, msg])
    msg.submit(chat_with_gemini, inputs=[msg, state], outputs=[chatbot, state, msg])  # ENTER key to send
    quit_btn.click(handle_quit, inputs=[state], outputs=[chatbot, state])

demo.launch()