from dotenv import load_dotenv
import streamlit as st
import os
import google.generativeai as genai
import random
from streamlit import session_state as state
from formulas import headline_formulas
from angles import angles

# Cargar las variables de entorno
load_dotenv()

# Configurar la API de Google
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Fórmulas con ejemplos y explicaciones
# headline_formulas dictionary has been moved to formulas/headline_formulas.py

def generate_headlines(number_of_headlines, target_audience, product, temperature, selected_formula, selected_angle, file_content="", image_parts=None, is_image=False):
    # Crear la configuración del modelo
    generation_config = {
        "temperature": temperature,
        "top_p": 0.65,
        "top_k": 360,
        "max_output_tokens": 8196,
    }

    model = genai.GenerativeModel(
        model_name="gemini-2.0-flash",
        generation_config=generation_config,
    )
    
    # Angle dictionaries have been moved to angles/angle_data.py
    
    # Incluir las instrucciones del sistema en el prompt principal
    system_prompt = f"""You are a world-class copywriter, with expertise in crafting hooks, headlines, and subject lines that immediately capture the reader's attention, prompting them to open the email or continue reading.

FORMAT RULES:
- Each headline must start with number and period
- One headline per line
- No explanations or categories
- Add a line break between each headline
- Avoid unnecessary : symbols
- Each headline must be a complete and intriguing sentence

IMPORTANT ANGLE INSTRUCTIONS:
- The selected angle MUST be applied to EVERY headline
- The angle modifies HOW the formula is expressed, not its structure
- Think of the angle as a "tone overlay" on the formula
- The formula provides the structure, the angle provides the style
- Both must work together seamlessly

FORMAT EXAMPLE:
1. Titular 1.

2. Titular 2.

3. Titular 3.

4. Titular 4.

5. Titular 5.

IMPORTANT:
- Each headline must be unique and memorable
- Avoid clichés and generalities
- Maintain an intriguing but credible tone
- Adapt speaking language from the audience
- Focus on transformative benefits
- Follow the selected angle style while maintaining formula structure"""

    # Iniciar el prompt con las instrucciones del sistema
    headlines_instruction = f"{system_prompt}\n\n"

    # Añadir contenido del archivo si existe
    if file_content:
        headlines_instruction += f"""
REFERENCE CONTENT:
Carefully analyze the following content as a reference for generating headlines:
{file_content[:3000]}

ANALYSIS INSTRUCTIONS:
1. Extract key information about the product or service mentioned
2. Identify the tone, style, and language used
3. Detect any data about the target audience or customer avatar
4. Look for benefits, features, or pain points mentioned
5. Use relevant terms, phrases, or concepts from the content
6. Maintain consistency with the brand identity or main message
7. Adapt the headlines to resonate with the provided content
8. Identify any creative elements like characters, movies, cultural references, or storytelling devices
9. Incorporate these creative elements into your headlines when appropriate

CREATIVE CONTEXT GUIDELINES:
- If the content mentions specific characters, incorporate their traits or quotes in headlines
- If movie or pop culture references exist, use similar tone or style in headlines
- Match the emotional tone and creative direction of the original content
- Maintain the same storytelling approach if present in the content
- Use metaphors, analogies or creative devices similar to those in the content

IMPORTANT COMBINATIONS:
"""
        # Updated conditions for specific input combinations
        if product and not target_audience:
            headlines_instruction += f"""- FILE + PRODUCT: You have a reference document and product ({product}). Create headlines that highlight this specific product's benefits and features using insights from the document. Extract audience information from the document to better target the headlines. Ensure headlines match the creative context and tone of the original content.
"""
        elif target_audience and not product:
            headlines_instruction += f"""- FILE + TARGET AUDIENCE: You have a reference document and target audience ({target_audience}). Create headlines tailored to this specific audience using language and concepts from the document. Identify products or services from the document that would appeal to this audience. Maintain the creative direction and contextual elements from the original content.
"""
        elif product and target_audience:
            headlines_instruction += f"""- PRODUCT + TARGET AUDIENCE: You have both product ({product}) and target audience ({target_audience}). Create headlines that connect this specific product with this specific audience, using insights from the document to strengthen the connection. Incorporate any creative elements, characters, or thematic approaches from the original content.
"""
        
        headlines_instruction += """
IMPORTANT: Naturally integrate the elements found in the content with the selected formula and angle. If the content has a specific creative direction (storytelling, character-based, movie reference, etc.), ensure your headlines reflect this creative context while maintaining the formula structure.
"""
    
    # Añadir instrucciones de ángulo solo si no es "NINGUNO"
    if selected_angle != "NINGUNO":
        headlines_instruction += f"""
ÁNGULO PRINCIPAL: {selected_angle}
INSTRUCCIONES DE ÁNGULO ESPECÍFICAS:
{angles[selected_angle]["instruction"]}

IMPORTANTE: El ángulo {selected_angle} debe aplicarse como una "capa de estilo" sobre la estructura de la fórmula:
1. Mantén la estructura base de la fórmula intacta
2. Aplica el tono y estilo del ángulo {selected_angle}
3. Asegura que cada elemento de la fórmula refleje el ángulo
4. El ángulo afecta al "cómo" se dice, no al "qué" se dice

EJEMPLOS EXITOSOS DEL ÁNGULO {selected_angle}:
"""
        for example in angles[selected_angle]["examples"]:
            headlines_instruction += f"- {example}\n"

    headlines_instruction += (
        f"\nTu tarea es crear {number_of_headlines} titulares irresistibles para {target_audience} "
        f"que capturen la atención instantáneamente y generen curiosidad sobre {product}. "
    )

    if selected_angle != "NINGUNO":
        headlines_instruction += f"IMPORTANTE: Cada titular DEBE seguir el ángulo {selected_angle} de manera clara y consistente.\n\n"

    headlines_instruction += (
        f"Evita menciones obvias de {product} y enfócate en despertar interés genuino"
    )

    if selected_angle != "NINGUNO":
        headlines_instruction += f" usando el ángulo seleccionado"

    headlines_instruction += ".\n\n"

    headlines_instruction += (
        f"IMPORTANTE: Estudia cuidadosamente estos ejemplos de la fórmula seleccionada. "
        f"Cada ejemplo representa el estilo y estructura a seguir"
    )

    if selected_angle != "NINGUNO":
        headlines_instruction += f", adaptados al ángulo {selected_angle}"

    headlines_instruction += ":\n\n"

    # Agregar 5 ejemplos aleatorios de la fórmula
    random_examples = random.sample(selected_formula['examples'], min(5, len(selected_formula['examples'])))

    headlines_instruction += "EJEMPLOS DE LA FÓRMULA A SEGUIR:\n"
    for i, example in enumerate(random_examples, 1):
        headlines_instruction += f"{i}. {example}\n"

    headlines_instruction += "\nINSTRUCCIONES ESPECÍFICAS:\n"
    headlines_instruction += "1. Mantén la misma estructura y longitud que los ejemplos anteriores\n"
    headlines_instruction += "2. Usa el mismo tono y estilo de escritura\n"
    headlines_instruction += "3. Replica los patrones de construcción de frases\n"
    headlines_instruction += "4. Conserva el nivel de especificidad y detalle\n"
    headlines_instruction += f"5. Adapta el contenido para {target_audience} manteniendo la esencia de los ejemplos\n\n"

    headlines_instruction += f"FÓRMULA A SEGUIR:\n{selected_formula['description']}\n\n"

    # CORRECTO (con indentación):
    if selected_angle != "NINGUNO":
        headlines_instruction += f"""
    RECORDATORIO FINAL:
    1. Sigue la estructura de la fórmula seleccionada
    2. Aplica el ángulo como una "capa de estilo"
    3. Mantén la coherencia entre fórmula y ángulo
    4. Asegura que cada titular refleje ambos elementos

GENERA AHORA:
Crea {number_of_headlines} titulares que sigan fielmente el estilo y estructura de los ejemplos mostrados.
"""
    else:
        headlines_instruction += f"""
GENERA AHORA:
Crea {number_of_headlines} titulares que sigan fielmente el estilo y estructura de los ejemplos mostrados.
"""

    # Modificar la forma de enviar el mensaje según si hay imagen o no
    if is_image and image_parts:
        chat_session = model.start_chat(
            history=[
                {
                    "role": "user",
                    "parts": [
                        headlines_instruction,
                        image_parts
                    ],
                },
            ]
        )
        response = chat_session.send_message("Genera los titulares siguiendo exactamente el estilo de los ejemplos mostrados, inspirándote en la imagen proporcionada.")
    else:
        chat_session = model.start_chat(
            history=[
                {
                    "role": "user",
                    "parts": [headlines_instruction],
                },
            ]
        )
        response = chat_session.send_message("Genera los titulares siguiendo exactamente el estilo de los ejemplos mostrados.")
    
    return response.text

# Configurar la interfaz de usuario con Streamlit
st.set_page_config(page_title="Enchanted Hooks", layout="wide")

# Leer el contenido del archivo manual.md
with open("manual.md", "r", encoding="utf-8") as file:
    manual_content = file.read()

# Mostrar el contenido del manual en el sidebar
st.sidebar.markdown(manual_content)

# Load CSS from file
with open("styles/main.css", "r") as f:
    css = f.read()
    
# Apply the CSS
st.markdown(f"<style>{css}</style>", unsafe_allow_html=True)

# Centrar el título y el subtítulo
st.markdown("<h1 style='text-align: center;'>Enchanted Hooks</h1>", unsafe_allow_html=True)
st.markdown("<h4 style='text-align: center;'>Imagina poder conjurar títulos que no solo informan, sino que encantan. Esta app es tu varita mágica en el mundo del copywriting, transformando cada concepto en un titular cautivador que deja a todos deseando más.</h4>", unsafe_allow_html=True)

# Crear columnas
col1, col2 = st.columns([1, 2])  

# Columnas de entrada
with col1:
    target_audience = st.text_input("¿Quién es tu público objetivo?", placeholder="Ejemplo: Estudiantes Universitarios")
    product = st.text_input("¿Qué producto tienes en mente?", placeholder="Ejemplo: Curso de Inglés")
    number_of_headlines = st.selectbox("Número de Titulares", options=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], index=4)

    # Crear un único acordeón para fórmula, creatividad y ángulo
    with st.expander("Personaliza tus titulares"):
        # Añadir área de texto para copy base
        base_copy = st.text_area(
            "Correo o copy para inspiración de titulares",
            placeholder="Agrega aquí un texto que sirva como base para generar los titulares. Si lo dejas vacío, se generarán basados en el producto y público objetivo.",
            height=150
        )
        
        temperature = st.slider("Creatividad", min_value=0.0, max_value=2.0, value=1.0, step=0.1)

        selected_formula_key = st.selectbox(
            "Selecciona una fórmula para tus titulares",
            options=list(headline_formulas.keys())
        )

        # Automatically use the keys from the angles dictionary
        # Make sure "NINGUNO" appears first, then the rest alphabetically
        angle_keys = ["NINGUNO"] + sorted([key for key in angles.keys() if key != "NINGUNO"])
        selected_angle = st.selectbox(
            "Selecciona el ángulo para tus titulares",
            options=angle_keys
        )
        
        # Añadir cargador de archivos dentro del acordeón
        uploaded_file = st.file_uploader("📄 Archivo o imagen de referencia", 
                                        type=['txt', 'pdf', 'docx', 'jpg', 'jpeg', 'png'])
        
        file_content = ""
        is_image = False
        image_parts = None
        
        if uploaded_file is not None:
            file_type = uploaded_file.name.split('.')[-1].lower()
            
            # Manejar archivos de texto
            if file_type in ['txt', 'pdf', 'docx']:
                if file_type == 'txt':
                    try:
                        file_content = uploaded_file.read().decode('utf-8')
                        st.success(f"Archivo TXT cargado correctamente: {uploaded_file.name}")
                    except Exception as e:
                        st.error(f"Error al leer el archivo TXT: {str(e)}")
                        file_content = ""
                
                elif file_type == 'pdf':
                    try:
                        import PyPDF2
                        pdf_reader = PyPDF2.PdfReader(uploaded_file)
                        file_content = ""
                        for page in pdf_reader.pages:
                            file_content += page.extract_text() + "\n"
                        st.success(f"Archivo PDF cargado correctamente: {uploaded_file.name}")
                    except Exception as e:
                        st.error(f"Error al leer el archivo PDF: {str(e)}")
                        file_content = ""
                        
                elif file_type == 'docx':
                    try:
                        import docx
                        doc = docx.Document(uploaded_file)
                        file_content = "\n".join([para.text for para in doc.paragraphs])
                        st.success(f"Archivo DOCX cargado correctamente: {uploaded_file.name}")
                    except Exception as e:
                        st.error(f"Error al leer el archivo DOCX: {str(e)}")
                        file_content = ""
            
            # Manejar archivos de imagen
            elif file_type in ['jpg', 'jpeg', 'png']:
                try:
                    from PIL import Image
                    image = Image.open(uploaded_file)
                    image_bytes = uploaded_file.getvalue()
                    image_parts = {
                        "mime_type": uploaded_file.type,
                        "data": image_bytes
                    }
                    is_image = True
                    st.image(image, caption="Imagen cargada", use_column_width=True)
                except Exception as e:
                    st.error(f"Error al procesar la imagen: {str(e)}")
                    is_image = False

    selected_formula = headline_formulas[selected_formula_key]

    # Botón de enviar
    submit = st.button("Generar Titulares")

# Mostrar los titulares generados
if submit:
    # Check if we have a valid combination of inputs
    has_file = 'file_content' in locals() and file_content.strip() != ""
    has_product = product.strip() != ""
    has_audience = target_audience.strip() != ""
    has_base_copy = 'base_copy' in locals() and base_copy.strip() != ""
    
    # Valid combinations:
    # 1. File + Product (no audience needed)
    # 2. File + Audience (no product needed)
    # 3. Product + Audience (traditional way)
    # 4. Base copy alone (new option)
    valid_inputs = (
        (has_file and has_product) or
        (has_file and has_audience) or
        (has_product and has_audience) or
        has_base_copy
    )
    
    if valid_inputs and selected_formula:
        try:
            # If base_copy is provided, use it as file_content
            content_for_generation = file_content if 'file_content' in locals() and file_content.strip() != "" else ""
            if has_base_copy:
                content_for_generation = base_copy
                
            generated_headlines = generate_headlines(
                number_of_headlines,
                target_audience,
                product,
                temperature,
                selected_formula,
                selected_angle,
                content_for_generation,
                image_parts if 'image_parts' in locals() else None,
                is_image if 'is_image' in locals() else False
            )
            col2.markdown(f"""
                <div class="results-container">
                    <h4>Observa la magia en acción:</h4>
                    <p>{generated_headlines}</p>
                </div>
            """, unsafe_allow_html=True)
        except ValueError as e:
            col2.error(f"Error: {str(e)}")
    else:
        if not selected_formula:
            col2.error("Por favor, selecciona una fórmula.")
        else:
            col2.error("Por favor, proporciona al menos una de estas combinaciones: archivo + producto, archivo + público objetivo, producto + público objetivo, o un texto de inspiración.")