import os
import gradio as gr
import google.generativeai as genai
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Retrieve API key from environment variable
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

# Retrieve system content from environment variable
SYSTEM_CONTENT = os.getenv("SYSTEM_CONTENT")

# Configure Google Gemini API
genai.configure(api_key=GEMINI_API_KEY)

# Create the model
generation_config = {
    "temperature": 0.7,
    "top_p": 0.95,
    "top_k": 64,
    "max_output_tokens": 512,  # Adjust as needed
    "response_mime_type": "text/plain",
}

# Define safety settings for the model
safety_settings = [
    {
        "category": "HARM_CATEGORY_HARASSMENT",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_HATE_SPEECH",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
        "threshold": "BLOCK_NONE"
    },
    {
        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
        "threshold": "BLOCK_NONE"
    }
]

# Function to read names and characters from list.txt
def load_characters():
    characters = {}  # Use a dictionary to store names and full descriptions
    with open("list.txt", "r") as f:
        for line in f:
            if "(" in line:
                name = line.split("(")[0].strip()
                description = line.strip()  # Keep the full line with brackets
                characters[name] = description 
            else:
                characters[line.strip()] = line.strip()  # If no brackets, name is the description
    return characters

# Function to generate a response based on user input and chat history
def generate_response(user_input, chat_history, selected_character):
    """Generates a response based on user input, chat history, and selected character."""

    # Get the full character description from the dictionary
    full_character_description = characters.get(selected_character, selected_character)

    # Update system content with the full character description
    updated_system_content = SYSTEM_CONTENT.replace("godfather", full_character_description)

    # Create the generative model
    model = genai.GenerativeModel(
        model_name="gemini-1.5-pro", 
        generation_config=generation_config,
        safety_settings=safety_settings,
        system_instruction=updated_system_content,
    )

    # Add user input to history
    chat_history.append(user_input)

    # Limit history length
    if len(chat_history) > 10:
        chat_history = chat_history[-10:]

    # Start a new chat session
    chat_session = model.start_chat()

    # Send the entire chat history as the first message
    response = chat_session.send_message("\n".join(chat_history))

    # Return response and update chat history
    return response.text, chat_history

# Load characters from list.txt
characters = load_characters()

# Extract character names for the dropdown
character_names = list(characters.keys())

with gr.Blocks() as iface:
    gr.Interface(
        fn=generate_response,
        inputs=[
            gr.Textbox(lines=2, label="Talk to AI", placeholder="Enter your message here..."),
            gr.State([]),  # State input for chat history
            gr.Dropdown(choices=character_names, label="Select Character")
        ],
        outputs=[
            gr.Textbox(label="Response"),
            gr.State([])  # State output to update chat history 
        ],
        title="AI Character Chatbot",
        description="Chat with different AI characters!<br>"
                "Contact me if you want another character/voice<br>"
                "WhatsApp me: +92-332-4399819<br>"
                "Email me: aheedsajid@gmail.com<br>"
                "<b>Donate something to increase GPU power</b><br>"
                "[Click here to Donate](https://nowpayments.io/donation/aheed)")

iface.launch()