import google.generativeai as genai
import streamlit as st
from dotenv import load_dotenv
import os

load_dotenv()
api_key = os.getenv("GOOGLE_API_KEY") 

MODEL_NAME = "gemini-2.0-flash"  
SAFETY_SETTINGS = [
    {
        "category": "HARM_CATEGORY_HARASSMENT",
        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
    },
    {
        "category": "HARM_CATEGORY_HATE_SPEECH",
        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
    },
    {
        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
    },
    {
        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
        "threshold": "BLOCK_MEDIUM_AND_ABOVE"
    },
]

def initialize_gemini():
    try:
        model = genai.GenerativeModel(MODEL_NAME, safety_settings=SAFETY_SETTINGS)
        return model
    except Exception as e:
        st.error(f"Error initializing Gemini model: {e}")
        return None

def get_gemini_response(model, prompt,
                        output_format=None, tone=None, length=None, num_responses=1, creative_mode=False, fact_checking=False, chat=None):  

    constructed_prompt = prompt
    if output_format:
        constructed_prompt += f" Generate this in the format of a {output_format}."
    if tone:
        constructed_prompt += f" Use a {tone} tone."
    if length:
        constructed_prompt += f" The generated text should be approximately {length} words long."
    if creative_mode:
        constructed_prompt += " Be creative."
    if fact_checking:
        constructed_prompt += " Please ensure factual accuracy."

    try:
        if chat:
            response = chat.send_message(constructed_prompt)
        else:
            response = model.generate_content(constructed_prompt, generation_config=genai.types.GenerationConfig(candidate_count=num_responses)) 
        return response.text if response else None
    except Exception as e:
        st.error(f"Error generating response: {e}")
        return None

custom_css = """
<style>
.user-message {
    background-color: #b49577; 
    border-radius: 10px;
    padding: 25px 30px;
    margin-bottom: 3px;
    text-align: right;
    width: fit-content;
    margin-left: auto; 
    word-break: break-word; 
}

.gemini-message {
    background-color: #FFFFFF; 
    border-radius: 10px;
    padding: 25px 30px;
    margin-bottom: 3px;
    text-align: left; 
    width: fit-content;
    margin-right: auto; 
    word-break: break-word; 
    font-size: 16px;
    font-family: sans-serif; 
    line-height: 1.6; 
}

.message-container {
    display: flex; 
    flex-direction: column; 
    margin-bottom: 10px;
}

.stChatMessage {
    padding: 0px; 
    border: none; 
    background-color: transparent;
}

/* CSS for the loading dots */
.loader {
  font-size: 2em; /* Adjust size as needed */
  color: #3498db; /* Loading dots color */
}

.loader:after {
  content: ' .';
  animation: dots 1s steps(5, end) infinite;
}

@keyframes dots {
  0%, 20% {
    content: ' .';
  }
  40% {
    content: ' ..';
  }
  60%, 80% {
    content: ' ...';
  }
  100% {
    content: ' .';
  }
}

/* Styles for input boxes */
div.stTextArea > div > div > textarea {
    background-color: #f0f2f6 !important; /* Light gray background */
    border-radius: 5px;
    padding: 8px;
    color: black; /* Ensuring text is readable */
}

div.stSelectbox > label {
    color: black;
}

div.stSelectbox > div > button {
    background-color: #f0f2f6 !important;
    color: black;
    border-radius: 5px;
}

div.stNumberInput > label {
    color: black;
}

div.stNumberInput > div > div input {
    background-color: #f0f2f6 !important;
    color: black;
    border-radius: 5px;
}

div.stSlider > label {
    color: black;
}

div.stSlider > div > div > div[data-baseweb="slider"] {
    background-color: #f0f2f6 !important;
    border-radius: 5px;
}

div.stCheckbox > label {
    color: black;
}

</style>
"""

def main():
    st.set_page_config(
        page_title="AI Text Generator",
        page_icon=":speech_balloon:",
        layout="wide",
        initial_sidebar_state="expanded",
    )

    st.markdown(custom_css, unsafe_allow_html=True)

    st.title("AI Text Generator :speech_balloon:")

    if "chat_history" not in st.session_state:
        st.session_state["chat_history"] = []

  
    with st.sidebar:
        st.title("Chat History")

        clear_button_key = "clear_sidebar" 
        if st.button("Clear History", key=clear_button_key):
            st.session_state["chat_history"] = []
            if "chat" in st.session_state:
                del st.session_state["chat"]
            st.rerun()

        st.subheader("Recent:") 

        for i, message in enumerate(st.session_state["chat_history"]):
            if message["role"] == "user": 
                st.markdown(f"**Input {i//2 + 1}:** {message['content']}") 

    if "generated" not in st.session_state:
        st.session_state["generated"] = False

    if not st.session_state["generated"]:
        topic = st.text_area("Enter Topic:")

        with st.container():
            col1, col2, col3 = st.columns(3)
            with col1:
                output_format = st.selectbox("Output Format:", ["Story", "Poem", "Article", "Code", None])
            with col2:
                tone = st.selectbox("Tone:", ["Formal", "Informal", "Humorous", "Technical", None])
            with col3:
                num_responses = st.number_input("Responses:", min_value=1, max_value=5, value=1, step=1)

        with st.container(): 
            length = st.slider("Length (words):", 0, 300, 100)          
            creative_mode = st.checkbox("Creative Mode")
            fact_checking = st.checkbox("Fact-Check")


        model = initialize_gemini()
        if model is None:
            st.error("Chatbot initialization failed. Check your API key.")
            return


        if "chat" not in st.session_state:
            try:
                st.session_state["chat"] = model.start_chat(history=st.session_state["chat_history"])
            except Exception as e:
                st.error(f"Error starting chat session: {e}")
                return

        chat = st.session_state["chat"]


        if st.button("Generate Text"):
            if not topic:
                st.warning("Please enter a topic.")
            else:

                with st.spinner(f"Generating"):
                    response = get_gemini_response(model, topic, output_format, tone, length, num_responses, creative_mode, fact_checking, chat)

                if response:
                    st.session_state["generated_text"] = response
                    st.session_state["topic"] = topic 
                    st.session_state["chat_history"].append({"role": "user", "content": topic})
                    st.session_state["chat_history"].append({"role": "gemini", "content": response}) 

                    st.session_state["generated"] = True
                    st.rerun()  

                else:
                    st.error("Failed to get a response from the chatbot.")

    else:
        st.subheader("Generated Text:")
        st.markdown(f'<div class="gemini-message">{st.session_state["generated_text"]}</div>', unsafe_allow_html=True)

        col1, col2 = st.columns(2)
        with col1:
            like = st.button("Like")
        with col2:
            dislike = st.button("Dislike")

        if like:
            st.success("Thanks for your feedback!")
        if dislike:
            st.error("We appreciate your feedback.")

        if st.button("Generate Again"):
            st.session_state["generated"] = False
            del st.session_state["generated_text"] 
            st.rerun()


if __name__ == "__main__":
    main()