import gradio as gr
import google.generativeai as genai
import os
import faiss
import gradio as gr
from sentence_transformers import SentenceTransformer
from google import genai
from google.genai import types

# --- Load Gemini API key from user ---
def set_api_key(api_key):
    os.environ["GEMINI_API_KEY"] = api_key
    genai.configure(api_key=api_key)
    return "✅ API Key set successfully!"

# --- Load text file content ---
def load_text_file(file_obj):
    global knowledge_base
    with open(file_obj.name, "r", encoding="utf-8") as f:
        knowledge_base = f.read()
    return "📂 File uploaded and knowledge base loaded!"

# --- Chat function ---
def chat(user_input):
    global knowledge_base
    if "knowledge_base" not in globals():
        return "⚠️ Please upload a text file first."
    
    # Check if user input is in knowledge base
    if user_input.lower() in knowledge_base.lower():
        return "📖 Based on knowledge base: " + user_input
    
    # Fallback to Gemini model
    model = genai.GenerativeModel("gemini-pro")
    response = model.generate_content(user_input)
    return response.text

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 Arduino Experts Chatbot")

    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)

    file_input = gr.File(label="Upload your TXT file")
    file_status = gr.Textbox(label="File Status", interactive=False)

    user_input = gr.Textbox(label="Ask a question")
    output = gr.Textbox(label="Answer")

    set_key_btn.click(fn=set_api_key, inputs=api_key, outputs=api_status)
    file_input.upload(fn=load_text_file, inputs=file_input, outputs=file_status)
    user_input.submit(fn=chat, inputs=user_input, outputs=output)

demo.launch(share=True)