import os
import gradio as gr
from google import genai
from google.genai.types import Tool, GoogleSearch, GenerateContentConfig
from dotenv import load_dotenv

# Load API key từ file .env
load_dotenv()
API_KEY = os.getenv("GEMINI_API_KEY")

# Cấu hình client của Gemini
client = genai.Client(api_key=API_KEY)

# Công cụ tìm kiếm Google
google_search_tool = Tool(google_search=GoogleSearch())

# Biến lưu trữ model hiện tại
current_model = "gemini-2.0-flash-exp"

# Hàm xử lý hội thoại
def chatbot_gemini(user_input):
    try:
        response = client.models.generate_content(
            model=current_model,
            contents=user_input,
        )
        reply = response.candidates[0].content.parts[0].text
        return reply, ""
    except Exception as e:
        return f"Đã xảy ra lỗi: {str(e)}", ""

# Hàm chuyển đổi model
def set_model(selected_model):
    global current_model
    current_model = selected_model
    return f"Model hiện tại: {current_model}"

# Giao diện Gradio
with gr.Blocks() as demo:
    gr.Markdown("## Gemini 2.0 Flash Chatbot")

    chatbot = gr.Chatbot(label="Chatbot")
    user_input = gr.Textbox(label="Nhập câu hỏi của bạn", placeholder="Hỏi bất kỳ điều gì...")
    submit_btn = gr.Button("Gửi")

    model_dropdown = gr.Dropdown(
        choices=["gemini-2.0-flash-exp", "gemini-2.0-flash-thinking-exp"],
        label="Chọn Model",
        value="gemini-2.0-flash-exp"
    )
    model_status = gr.Textbox(label="Model hiện tại", value=f"Model hiện tại: {current_model}", interactive=False)

    model_dropdown.change(
        fn=set_model,
        inputs=model_dropdown,
        outputs=model_status
    )

    submit_btn.click(
        fn=lambda user_input: chatbot_gemini(user_input),
        inputs=user_input,
        outputs=[chatbot, user_input]
    )

if __name__ == "__main__":
    demo.launch()