import gradio as gr
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold
import os

# Retrieve Gemini API key from environment (already set in Hugging Face Spaces)
gemini_api_key = os.getenv("GEMINI_API_KEY", "")

# Ensure API key is configured
if gemini_api_key:
    genai.configure(api_key=gemini_api_key)
else:
    raise ValueError("❌ No GEMINI_API_KEY found. Set it in the environment variables on Hugging Face Spaces.")

# Chatbot function
def chat_with_cardi(user_input, chat_history):
    generation_config = {
        "temperature": 0.9,  # Creative responses
        "top_p": 0.95,
        "top_k": 64,
        "max_output_tokens": 8192,
        "response_mime_type": "text/plain",
    }

    model = genai.GenerativeModel(
        model_name="gemini-pro",
        generation_config=generation_config,
    )

    try:
        # Get Gemini response
        response = model.generate_content(
            f"You are Cardi B. You are bold, flamboyant, and outspoken. Respond like Cardi B would. {user_input}",
            safety_settings={
                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
            }
        )
        return response.text
    except Exception as e:
        return f"🚨 Error: {str(e)}"

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🎤 Cardi B GPT 💅\nAsk Cardi anything and get her spicy takes!")

    chatbot = gr.Chatbot(label="Cardi B AI 💅", type="messages")  # FIXED: Added `type="messages"`

    textbox = gr.Textbox(label="Ask Cardi B something 💄")  # FIXED: Removed `input_textbox`
    submit_btn = gr.Button("Send it, babe! 💅")

    def process_chat(user_input, history):
        response = chat_with_cardi(user_input, history)
        history.append({"role": "user", "content": user_input})  # OpenAI-style format
        history.append({"role": "assistant", "content": response})
        return history, ""

    submit_btn.click(process_chat, inputs=[textbox, chatbot], outputs=[chatbot, textbox])

# Launch the Gradio app
demo.launch()