import streamlit as st
from gtts import gTTS
import google.generativeai as genai
from io import BytesIO

# Set your API key
api_key = "AIzaSyC70u1sN87IkoxOoIj4XCAPw97ae2LZwNM"  # Replace with your actual API key

genai.configure(api_key=api_key)

# Configure the generative AI model
generation_config = genai.GenerationConfig(
    temperature=0.9,
    max_output_tokens=3000
)

# Safety settings configuration
safety_settings = [
    {
        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
        "threshold": "BLOCK_NONE",
    },
    {
        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
        "threshold": "BLOCK_NONE",
    },
    {
        "category": "HARM_CATEGORY_HATE_SPEECH",
        "threshold": "BLOCK_NONE",
    },
    {
        "category": "HARM_CATEGORY_HARASSMENT",
        "threshold": "BLOCK_NONE",
    },
]

# Initialize session state for chat history
if 'chat_history' not in st.session_state:
    st.session_state['chat_history'] = []

st.title("Gemini Chatbot")

# Display chat history
def display_chat_history():
    for entry in st.session_state['chat_history']:
        st.markdown(f"{entry['role'].title()}: {entry['parts'][0]['text']}")

# Function to clear conversation history
def clear_conversation():
    st.session_state['chat_history'] = []

# Send message function with sequential AI model interaction
def send_message():
    user_input = st.session_state.user_input
    if user_input:
        # Initial system prompt for the chatbot interaction
        initial_system_prompt = "You are a knowledgeable and helpful chatbot. Respond to the user queries informatively and politely."
        
        # AI Writer System Prompt for generating text based on the outline
        ai_writer_system_prompt = "As the AI Writer, your main objective is to generate the actual text of the book based on the outline provided by the AI Planner. You will use natural language generation techniques to produce coherent and readable prose that follows the structure and narrative defined by the AI Planner. Your output should adhere to the user's style and tone preferences, and you should incorporate any specific information or prompts provided by the user to create a captivating and immersive story."
        
        prompts = [entry['parts'][0]['text'] for entry in st.session_state['chat_history']]
        prompts.append(user_input)
        
        # Combine initial system prompt with the chat history
        chat_history_str = initial_system_prompt + "\n" + "\n".join(prompts)

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

        # First model generation call
        initial_response = model.generate_content([{"role": "user", "parts": [{"text": chat_history_str}]}])
        initial_response_text = initial_response.text if hasattr(initial_response, "text") else "No response text found."

        if initial_response_text:
            # Use the output of the first model call as input for the second, applying the AI Writer System Prompt
            final_chat_history_str = ai_writer_system_prompt + "\n" + initial_response_text
            
            # Second model generation call
            #final_response = model.generate_content([{"role": "model", "parts": [{"text": final_chat_history_str}]}])
            # Corrected second model generation call
            final_response = model.generate_content([{"role": "user", "parts": [{"text": final_chat_history_str}]}])
            final_response_text = final_response.text if hasattr(final_response, "text") else "No response text found."

            if final_response_text:
                st.session_state['chat_history'].append({"role": "model", "parts":[{"text": final_response_text}]})

                # Convert the final response text to speech
                tts = gTTS(text=final_response_text, lang='en')
                tts_file = BytesIO()
                tts.write_to_fp(tts_file)
                tts_file.seek(0)
                st.audio(tts_file, format='audio/mp3')

        st.session_state.user_input = ''

    display_chat_history()

# User input text area
user_input = st.text_area(
    "Enter your message here:",
    value="",
    key="user_input"
)

# Send message button
send_button = st.button(
    "Send",
    on_click=send_message
)

# Clear conversation button
clear_button = st.button("Clear Conversation", on_click=clear_conversation)