from dotenv import load_dotenv
import streamlit as st
import os
import google.generativeai as genai
import time
import datetime
import random  # Add this import
from ads_formulas import ads_formulas  # Import the ads formulas
from style import styles
from prompts import create_fb_ad_instruction
from emotional_angles import emotional_angles
from copywriter_personas import copywriter_personas
from ad_objectives import ad_objectives  # Import ad objectives
import PyPDF2
import docx
from PIL import Image
import io
from headline_generator import generate_random_headlines, generate_headlines_with_model, headline_formulas  # Import headline_formulas too

# Cargar las variables de entorno
load_dotenv()

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

# Función para generar modelo
@st.cache_resource
def get_model(temperature):
    generation_config = {
        "temperature": temperature,
    }
    return genai.GenerativeModel('gemini-2.0-flash', generation_config=generation_config)

# Function to generate headlines and select one randomly
def generate_and_select_headline(target_audience, product, temperature, selected_angle, file_content="", image_parts=None):
    # Get the model
    model = get_model(temperature)
    
    # Select a random formula for the headline
    formula_name = random.choice(list(headline_formulas.keys()))
    selected_formula = headline_formulas[formula_name]
    
    # Create instruction with file content if available
    headline_instruction = f"Generate 3 compelling headlines for {product} targeting {target_audience}. Use this formula: {selected_formula}. Apply this emotional angle: {selected_angle}."
    
    # Add file content as context if available
    if file_content:
        headline_instruction += f"\n\nUse the following information as additional context:\n{file_content[:5000]}"
    
    # Generate headlines with or without image
    if image_parts:
        # If we have an image, use it for headline generation
        response = model.generate_content([headline_instruction, image_parts], generation_config={"temperature": temperature})
        headlines = response.parts[0].text if response and response.parts else ""
    else:
        # Generate 3 headlines using text only
        headlines = generate_headlines_with_model(
            model=model,
            target_audience=target_audience,
            product=product,
            selected_formula=selected_formula,
            selected_angle=selected_angle,
            number_of_headlines=3
        )
    
    # Parse the headlines (assuming they come in a numbered format)
    headline_list = []
    for line in headlines.strip().split('\n'):
        if line.strip() and line[0].isdigit() and '.' in line:
            # Extract the headline text after the number and period
            headline_text = line.split('.', 1)[1].strip()
            headline_list.append(headline_text)
    
    # If we couldn't parse any headlines, return a default one
    if not headline_list:
        return f"Descubre cómo {product} puede transformar tu vida como {target_audience}"
    
    # Select a random headline from the generated ones
    selected_headline = random.choice(headline_list)
    return selected_headline

# Move the clean_ad_text function up with other utility functions
# After the generate_headline function and before generate_fb_ad

def generate_headline(target_audience=None, product=None, temperature=1.0, angle=None):
    # Check if UI variables exist, otherwise use parameters
    try:
        target_audience = target_audience or ad_target_audience
        product = product or ad_product
        temperature = temperature or ad_temperature
        angle = angle or emotional_angle_key
    except NameError:
        # If UI variables don't exist yet, use the provided parameters
        pass
    
    # Obtener el modelo usando la función existente
    model = get_model(temperature)
    
    # Remove duplicate imports and use the ones at the top of the file
    formula_name = random.choice(list(headline_formulas.keys()))
    selected_formula = headline_formulas[formula_name]
    
    # Generar el titular usando el modelo ya inicializado
    headline = generate_headlines_with_model(
        model=model,
        target_audience=target_audience,
        product=product,
        selected_formula=selected_formula,
        selected_angle=angle
    )
    
    return headline

# Function to clean generated ad text
def clean_ad_text(text):
    """
    Cleans the generated ad text by removing unwanted characters and formatting issues.
    
    Args:
        text (str): The raw generated ad text
        
    Returns:
        str: Cleaned ad text
    """
    if not isinstance(text, str):
        return ""
    
    # Split text into lines
    lines = text.split('\n')
    
    # Clean lines
    cleaned_lines = []
    for line in lines:
        # Skip lines that are just a single letter
        if len(line.strip()) == 1 and line.strip().isalpha():
            continue
            
        # Remove triple backticks and language indicators (like ```text)
        if line.strip().startswith('```'):
            line = line.replace('```', '').strip()
            # If line now only contains a language indicator, skip it
            if line.strip() in ['text', 'markdown', 'md', 'html']:
                continue
                
        # Remove any other markdown artifacts
        line = line.replace('```', '')
        
        cleaned_lines.append(line)
    
    # Join lines back together
    cleaned_text = '\n'.join(cleaned_lines)
    
    # Remove any consecutive newlines (more than 2)
    import re
    cleaned_text = re.sub(r'\n{3,}', '\n\n', cleaned_text)
    
    return cleaned_text.strip()
    
# Modify the generate_fb_ad function to accept a headline parameter
def generate_fb_ad(target_audience, product, temperature, selected_formula, selected_angle, 
                  selected_persona, story_prompt="", ad_objective=None, file_content="", 
                  image_parts=None, max_file_chars=50000, headline=None):
    if not target_audience or not product:
        return "Por favor, completa todos los campos requeridos."
    
    # Generate a headline if none is provided
    if not headline:
        headline = generate_and_select_headline(target_audience, product, temperature, selected_angle)
    
    # Enfatizar el tema de la historia si se proporciona
    emphasized_story_prompt = story_prompt
    if story_prompt and story_prompt.strip():
        emphasized_story_prompt = story_prompt.strip()
    
    model = get_model(temperature)
    
    # Crear la instrucción base, ahora incluyendo el titular
    ad_instruction = create_fb_ad_instruction(
        target_audience, 
        product, 
        selected_formula,
        selected_angle,
        selected_persona,
        ad_objective,
        language="español",
        story_prompt=emphasized_story_prompt,
        headline=headline  # Pass the headline to the instruction creator
    )
    
    # Si hay contenido de archivo, añadirlo a la instrucción
    if file_content:
        ad_instruction += f"\n\nAdemás, utiliza la siguiente información como referencia para crear el anuncio:\n\n{file_content[:max_file_chars]}"
    
    # Si hay un tema específico, ajustar la temperatura para mayor coherencia
    effective_temperature = temperature
    if story_prompt and story_prompt.strip():
        effective_temperature = max(0.1, temperature * 0.9)
    
    # Generar el contenido con o sin imagen
    if image_parts:
        response = model.generate_content([ad_instruction, image_parts], generation_config={"temperature": effective_temperature})
    else:
        response = model.generate_content([ad_instruction], generation_config={"temperature": effective_temperature})
    
    return response.parts[0].text if response and response.parts else "Error generating content."

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

# Ocultar el menú de tres puntos de Streamlit
hide_menu_style = """
        <style>
        #MainMenu {visibility: hidden;}
        header {visibility: hidden;}
        footer {visibility: hidden;}
        </style>
        """
st.markdown(hide_menu_style, unsafe_allow_html=True)

# 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)

# Ocultar elementos de la interfaz
st.markdown(styles["main_layout"], unsafe_allow_html=True)

# Centrar el título y el subtítulo con el nuevo texto
st.markdown("<h1 style='text-align: center;'>CopyLegend AI</h1>", unsafe_allow_html=True)
st.markdown("<h4 style='text-align: center;'>La única herramienta que combina la sabiduría de los maestros del copywriting con la precisión de la inteligencia artificial</h4>", unsafe_allow_html=True)

# Añadir CSS personalizado para el botón
st.markdown(styles["button"], unsafe_allow_html=True)

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

# Columnas de entrada
with col1:
    ad_target_audience = st.text_input("¿Quién es tu público objetivo?", placeholder="Ejemplo: Madres trabajadoras de 30-45 años")
    ad_product = st.text_input("¿Qué producto tienes en mente?", placeholder="Ejemplo: Curso de gestión del tiempo")
    input_prompt = st.text_area("Escribe de qué quieres que trate la historia:", placeholder="Escribe aquí tu idea...")
    
    # Añadir cargador de archivos
    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')
                except Exception as e:
                    st.error(f"Error al leer el archivo TXT: {str(e)}")
                    file_content = ""
                
            elif file_type == 'pdf':
                try:
                    pdf_reader = PyPDF2.PdfReader(uploaded_file)
                    file_content = ""
                    for page in pdf_reader.pages:
                        file_content += page.extract_text() + "\n"
                except Exception as e:
                    st.error(f"Error al leer el archivo PDF: {str(e)}")
                    file_content = ""
                    
            elif file_type == 'docx':
                try:
                    doc = docx.Document(uploaded_file)
                    file_content = "\n".join([para.text for para in doc.paragraphs])
                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:
                image = Image.open(uploaded_file)
                image_bytes = uploaded_file.getvalue()
                image_parts = [
                    {
                        "mime_type": uploaded_file.type,
                        "data": image_bytes
                    }
                ]
                is_image = True
            except Exception as e:
                st.error(f"Error al procesar la imagen: {str(e)}")
                is_image = False
    
    # Mover el botón aquí, después del cargador de archivos
    submit_ad = st.button("GENERAR ANUNCIO")
    
    with st.expander("Opciones avanzadas"):
        # Selector de fórmula de anuncio
        ad_formula_key = st.selectbox(
            "Fórmula de anuncio",
            options=list(ads_formulas.keys()),
            label_visibility="visible"
        )
        ad_formula = ads_formulas[ad_formula_key]
        
        # Selector de ángulo emocional
        emotional_angle_key = st.selectbox(
            "Ángulo emocional",
            options=list(emotional_angles.keys()),
            label_visibility="visible"
        )
        emotional_angle = emotional_angles[emotional_angle_key]
        
        # Selector de personalidad de copywriter
        ad_persona_key = st.selectbox(
            "Estilo de copywriter",
            options=list(copywriter_personas.keys()),
            index=0,
            format_func=lambda x: x.replace("_", " "),
            label_visibility="visible"
        )
        ad_persona = copywriter_personas[ad_persona_key]
        
        # Selector de objetivo de anuncio
        ad_objective_key = st.selectbox(
            "Objetivo de la campaña",
            options=list(ad_objectives.keys()),
            label_visibility="visible"
        )
        selected_objective = ad_objectives[ad_objective_key] if ad_objective_key != "ninguno" else None

        ad_temperature = st.slider(
            "Nivel de creatividad",
            min_value=0.0, 
            max_value=1.5,  # Reduced from 2.0 to 1.5 for better policy compliance
            value=1.0, 
            step=0.1,
            label_visibility="visible"
        )
        
        # Add disclaimer about content review
        st.markdown("""
        <div style='font-size: 0.8em; color: #888;'>
        Nota: Todo contenido generado debe ser revisado para asegurar cumplimiento con las políticas de Facebook antes de publicar.
        </div>
        """, unsafe_allow_html=True)

# Mostrar el anuncio generado
if submit_ad:
    # Check for different valid input combinations
    has_file_content = file_content.strip() != "" or is_image
    has_basic_fields = ad_target_audience and ad_product
    
    # Generate ad if either file is uploaded OR basic fields are filled
    if has_file_content or has_basic_fields:
        # Mostrar el spinner en la segunda columna
        with col2:
            with st.spinner('Generando anuncio...'):
                # Use default values if fields are empty
                target_audience = ad_target_audience if ad_target_audience else "usuarios interesados"
                product = ad_product if ad_product else "este producto/servicio"
                story = input_prompt if input_prompt else ""
                
                # Generate a headline first, now passing file_content and image_parts
                headline = generate_and_select_headline(
                    target_audience,
                    product,
                    ad_temperature,
                    emotional_angle,
                    file_content,  # Pass file content to headline generation
                    image_parts if is_image else None  # Pass image data if available
                )
                
                # Generate the ad with the headline
                generated_ad = generate_fb_ad(
                    target_audience, 
                    product, 
                    ad_temperature, 
                    ad_formula,
                    emotional_angle,
                    ad_persona,
                    story,
                    selected_objective,
                    file_content,
                    image_parts if is_image else None,
                    50000,
                    headline
                )
                
                if not isinstance(generated_ad, str):
                    st.error("Error al generar el anuncio")
                else:
                    # Clean the generated ad text
                    cleaned_ad = clean_ad_text(generated_ad)
                    
                    # Display only the ad without separate headline and without storing in session state
                    st.markdown(f"""
                        <div style="{styles['results_container']}">
                            <h3>Anuncio Generado:</h3>
                            <p>{cleaned_ad}</p>
                        </div>
                    """, unsafe_allow_html=True)
                    
                    # Aplicar estilo para el botón de descarga
                    st.markdown(styles["download_button"], unsafe_allow_html=True)
                    
                    # Get current timestamp for the filename
                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                    
                    # Prepare content for download (without separate headline)
                    download_content = f"{cleaned_ad}"
                    
                    # Botón de descarga with timestamp in filename
                    st.download_button(
                        label="DESCARGAR ANUNCIO",
                        data=download_content,
                        file_name=f"anuncio_facebook_{timestamp}.txt",
                        mime="text/plain"
                    )
    else:
        col2.warning("Por favor, sube un archivo o completa al menos uno de los campos principales.")
# Remove the section that displays cached ads

# Agregar firma del autor al final de la página
st.markdown('---')
st.markdown('Made with ❤️ by Jesús Cabrera')