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

load_dotenv()

genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Function to get response from Gemini model
def get_gemini_response(input_prompt, image_data, genre, length, language, mood):
    model = genai.GenerativeModel('gemini-1.5-flash')
    # If both image and text are provided
    if image_data and input_prompt:
        full_prompt = f"""
        You are a famous creative writer. Look at the image provided and also consider the following prompt: "{input_prompt}". 
        Create a {length} {genre} in {language}. The {genre} should be {mood}, based on both the image and the prompt, and contain realistic and emotional elements.
        """
        response = model.generate_content([full_prompt, image_data[0]])

    # If only image is provided
    elif image_data:
        full_prompt = f"""
        You are a famous creative writer. Look at the image provided and create a {length} {genre} in {language}. 
        The {genre} should be {mood}. The {genre} should be based on the image and contain realistic and emotional elements.
        """
        response = model.generate_content([full_prompt, image_data[0]])

    # If only text is provided
    elif input_prompt:
        full_prompt = f"""
        You are a creative writer. Create a {length} {genre} in {language}. The {genre} should be {mood}. The {genre} should be based on the following prompt:
        
        "{input_prompt}"
        
        Make sure it contains realistic and emotional elements.
        """
        response = model.generate_content([full_prompt])
    
    # If neither image nor text is provided
    else:
        raise ValueError("Please provide either an image or a text prompt.")
    
    # Check if response is valid and return text
    if response and response.parts:
        return response.parts[0].text
    else:
        raise ValueError("Sorry, please try a different combination of inputs.")
        
# Function to setup uploaded image
def input_image_setup(uploaded_file):
    if uploaded_file is not None:
        bytes_data = uploaded_file.getvalue()
        image_parts = [
            {
                "mime_type": uploaded_file.type,
                "data": bytes_data
            }
        ]
        return image_parts
    else:
        return None

# Initialize Streamlit app
st.set_page_config(page_title="Poetic Vision", page_icon=":pencil:", layout="wide")  # Set layout to wide

# Main UI components
st.markdown("<h1 style='text-align: center;'>VisionTales</h1>", unsafe_allow_html=True)
st.markdown("<h3 style='text-align: center;'>Convierte tus ideas en historias inolvidables.</h3>", unsafe_allow_html=True)

# Add custom CSS for the button
st.markdown("""
    <style>
    div.stButton > button {
        background-color: #FFCC00; /* Color llamativo */
        color: black;  /* Texto en negro */
        width: 90%;
        height: 60px;
        font-weight: bold;
        font-size: 22px; /* Tamaño más grande */
        text-transform: uppercase; /* Texto en mayúsculas */
        border: 1px solid #000000; /* Borde negro de 1px */
        border-radius: 8px;
        display: block;
        margin: 0 auto;  /* Centramos el botón */
    }
    div.stButton > button:hover {
        background-color: #FFD700; /* Color al pasar el mouse */
        color: black;  /* Texto sigue en negro */
    }
    </style>
    """, unsafe_allow_html=True)

# Create two columns for layout (40% and 60%)
col1, col2 = st.columns([2, 3])  # 2 + 3 = 5 parts total

with col1:
    # Subir imagen primero
    uploaded_file = st.file_uploader("Elegir una imagen...", type=["jpg", "jpeg", "png"])

    # Agrupar opciones en un desplegable con un nuevo título
    with st.expander("Clic para personalizar tu historia", expanded=False):  # El desplegable está cerrado por defecto
        input_prompt = st.text_input("Escribe de qué quieres que trate tu historia (opcional):", placeholder="Escribe tu mensaje aquí...")
        genre = st.selectbox("Tipo de texto:", ["Historia", "Shayari", "Sher", "Poema", "Cita"])
        length = st.selectbox("Longitud del texto:", ["Corto", "Largo"])
        language = st.selectbox("Idioma del texto:", ["Inglés", "Español"])
        mood = st.selectbox("Estado de ánimo:", ["Emocional", "Triste", "Feliz", "Horror", "Comedia", "Romántico"])

    # Generate button con el nuevo estilo, centrado y ajustado
    submit = st.button("Escribir mi historia")

# Display uploaded image
if uploaded_file is not None:
    image = Image.open(uploaded_file)
    col1.image(image, caption="Imagen subida", use_column_width=True)

# Handling the button click
if submit:
    if uploaded_file or input_prompt:
        try:
            image_data = input_image_setup(uploaded_file)
            response = get_gemini_response(input_prompt, image_data, genre, length, language, mood)
            col2.subheader("Contenido generado:")
            col2.write(response)
        except ValueError as e:
            col2.error(f"Error: {str(e)}")
    else:
        col2.error("Por favor, sube una imagen o proporciona un mensaje.")