import gradio as gr
from google import genai
from google.genai import types

# --- Load knowledge base from repo at startup ---
with open("knowledge.txt", "r", encoding="utf-8") as f:
    knowledge_base = f.read()

# --- Store API key globally ---
global_api_key = None

# --- Load Gemini API key from user ---
def set_api_key(api_key):
    global global_api_key
    global_api_key = api_key.strip()
    return "✅ API Key set successfully!"

# --- Chat function ---
def chat(user_input, history):
    global knowledge_base, global_api_key
    history = history or []

    if not global_api_key:
        return history, "⚠️ Please set your Gemini API key first."

    if not knowledge_base:
        return history, "⚠️ Knowledge base is empty."

    # Combine knowledge base + chat history + current input
    context_text = f"Use the following knowledge base to answer the user's question:\n\n{knowledge_base}\n\n"
    context_text += "Previous chat history:\n"
    for msg in history:
        context_text += f"{msg['role'].capitalize()}: {msg['content']}\n"
    context_text += f"User: {user_input}\nAnswer based on the knowledge above."

    try:
        client = genai.Client(api_key=global_api_key)
        chat_session = client.chats.create(model="gemini-2.5-flash")
        response = chat_session.send_message(context_text)
        reply = response.text
    except Exception as e:
        reply = f"❌ Error calling Gemini model: {e}"

    # Append messages in new Gradio format
    history.append({"role": "user", "content": user_input})
    history.append({"role": "assistant", "content": reply})

    return history, ""  # "" clears the user input box


# --- Custom CSS for Arduino Expert Theme ---
custom_css = """
.gradio-container {
    background-color: #000000 !important;  /* Black background */
    font-family: 'Arial', sans-serif;
    color: white;
}

/* Header */
h1, h2, h3, .gr-markdown {
    color: #0082C9 !important; /* Arduino Blue */
    text-align: center;
    font-weight: bold;
}

/* Buttons */
button {
    background-color: #0082C9 !important;
    color: white !important;
    border-radius: 8px !important;
    font-weight: bold !important;
    border: none !important;
}
button:hover {
    background-color: #006fa8 !important;
}

/* Textboxes */
textarea, input {
    border: 2px solid #0082C9 !important;
    border-radius: 6px !important;
    background-color: #1a1a1a !important;
    color: white !important;
}

/* Chatbot */
.gr-chatbot {
    border: 2px solid #0082C9 !important;
    border-radius: 10px !important;
    background-color: #ffffff !important; /* White chat background */
    color: black !important;
    max-height: 70vh !important; 
    overflow-y: auto !important;
}
"""

# --- Gradio UI ---
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("## 🤖 Arduino Expert Chatbot")

    # API key input
    api_key = gr.Textbox(label="Enter your Gemini API Key", type="password")
    set_key_btn = gr.Button("Set API Key")
    api_status = gr.Textbox(label="Status", interactive=False)

    # Store chat history in session with initial greeting
    chat_history = gr.State([
        {"role": "system", "content": "👋 Welcome to Arduino Expert! Ask me anything about Arduino projects."}
    ])

    # Chatbot interface
    chatbot = gr.Chatbot(type="messages", label="Conversation")
    user_input = gr.Textbox(label="Ask a question")
    send_btn = gr.Button("Send")
    clear_btn = gr.Button("Clear Chat")

    # Set API key
    set_key_btn.click(fn=set_api_key, inputs=api_key, outputs=api_status)

    # Send message
    send_btn.click(fn=chat, inputs=[user_input, chat_history], outputs=[chatbot, user_input])
    user_input.submit(fn=chat, inputs=[user_input, chat_history], outputs=[chatbot, user_input])

    # Clear chat
    clear_btn.click(lambda: [{"role": "system", "content": "👋 Welcome to Arduino Expert! Ask me anything about Arduino projects."}], None, chatbot)

demo.launch(share=True)