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

# Cargar las variables de entorno
load_dotenv()

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

# Función para inicializar el modelo
def initialize_model(temperature=1.0):
    return genai.GenerativeModel(
        model_name="gemini-2.0-flash",
        generation_config={"temperature": temperature}
    )

# Initialize session state variables at the app level
session_state_vars = {
    'headline_content': "", 
    'problem_content': "", 
    'solution_content': "",
    'avatar_details': {}, 
    'product_details': {},
    'generated_sections': {},
    'metadata': {}
}

for var, default in session_state_vars.items():
    if var not in st.session_state:
        st.session_state[var] = default

# Función para generar contenido y procesar el JSON
def generate_content(prompt, temperature=1.0):
    model = initialize_model(temperature)
    
    response = model.generate_content(prompt)
    result = response.text
    
    # Try to extract JSON metadata if present and remove it from the displayed text
    try:
        import json
        import re
        
        # Look for JSON data in the response (both with and without code blocks)
        json_pattern = r'```json\s*(.*?)\s*```|(\{[\s\S]*"character"[\s\S]*\})'
        json_match = re.search(json_pattern, result, re.DOTALL)
        
        if json_match:
            # Get the matched JSON string (either from code block or raw JSON)
            json_str = json_match.group(1) if json_match.group(1) else json_match.group(2)
            
            # Parse the JSON data
            json_data = json.loads(json_str.strip())
            
            # Store the JSON data in session state
            st.session_state.metadata = json_data
            
            # Update avatar details and other metadata elements
            metadata_fields = ['character', 'pain_points', 'emotional_triggers', 'obstacles', 'failed_solutions']
            for field in metadata_fields:
                if field in json_data:
                    if field == 'character' and 'avatar_details' in st.session_state:
                        st.session_state.avatar_details.update(json_data[field])
                    else:
                        st.session_state[field] = json_data[field]
            
            # Remove the JSON block from the displayed text
            result = re.sub(json_pattern, '', result, flags=re.DOTALL).strip()
    except Exception as e:
        # Log the error but don't disrupt the user experience
        print(f"Error processing JSON metadata: {str(e)}")
    
    return result

# Modify the display function to preserve content
def display_generated_content(column, content, key_prefix):
    """Display generated content in the specified column with download button."""
    if content:
        column.markdown("### Contenido generado:")
        column.markdown(content)
        
        # Add download button
        download_button_key = f"{key_prefix}_download_button"
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        section_name = key_prefix.replace("_", "-")
        filename = f"sales_page_{section_name}_{timestamp}.txt"
        
        column.download_button(
            label="Descargar contenido",
            data=content,
            file_name=filename,
            mime="text/plain",
            key=f"{key_prefix}_actual_download"
        )

# Add this function to validate inputs
def validate_inputs(audience, product):
    """Validate that required inputs are provided."""
    return audience.strip() != "" and product.strip() != ""

# Make sure to define your UI elements before using them
# Add this before the submit button
st.set_page_config(layout="wide")  # Set the page layout to wide
st.title("Generador de Páginas de Ventas")

# Load custom CSS from the styles folder
def load_css(css_file):
    with open(css_file, 'r') as f:
        css = f.read()
    st.markdown(f'<style>{css}</style>', unsafe_allow_html=True)

# Load the CSS file
css_path = os.path.join(os.path.dirname(__file__), "styles", "styles.css")
load_css(css_path)

# Create two columns for layout with adjusted widths (40% left, 60% right)
col1, col2 = st.columns([4, 6])

with col1:
    st.subheader("Configuración")
    
    # Input fields
    sales_page_audience = st.text_input("Público objetivo:", 
                                        help="Describe a quién va dirigida tu página de ventas")
    
    sales_page_product = st.text_input("Producto/Servicio:", 
                                       help="Describe el producto o servicio que estás vendiendo")
    
    sales_page_offer = st.text_area("Oferta específica (opcional):", 
                                    help="Detalles específicos de tu oferta, como precio, bonos, etc.")
    
    sales_page_section = st.selectbox(
        "Sección a generar:",
        [
            "Above the Fold (Encabezado)",
            "Problem (Problema)",
            "Solution & Benefits (Solución y Beneficios)",
            "Authority (Autoridad)",
            "Offer & Bonus (Oferta y Bonus)",
            "Social Proof (Prueba Social)",
            "Guarantees (Garantías)",
            "Call to Action (Llamada a la Acción)",
            "P.S. (Post-Data)",
            "Final Call to Action (Llamada Final)"
        ]
    )
    
    sales_page_temperature = st.slider(
        "Creatividad:",
        min_value=0.0,
        max_value=1.0,
        value=0.7,
        step=0.1,
        help="Valores más altos = más creatividad, valores más bajos = más consistencia"
    )
    
    # Now create the submit button
    submit_sales_page = st.button("Generar Sección", key="generate_sales_page")

# Add these section generator functions before the section_functions dictionary

def load_prompt_template(filename):
    """Load a prompt template from the prompts directory."""
    prompt_path = os.path.join(os.path.dirname(__file__), "prompts", filename)
    with open(prompt_path, "r", encoding="utf-8") as file:
        return file.read()

def generate_above_the_fold(audience, product, temperature=0.7, offer=None):
    """Generate the Above the Fold (headline) section."""
    prompt_template = load_prompt_template("above_the_fold.txt")
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_problem_section(audience, product, temperature=0.7, offer=None):
    """Generate the Problem section."""
    prompt_template = load_prompt_template("problem.txt")
    
    # Create context dictionary with any existing session state data
    context = {}
    if 'headline_content' in st.session_state:
        context['headline_content'] = st.session_state.headline_content
    if 'avatar_details' in st.session_state:
        context['avatar_details'] = st.session_state.avatar_details
    if 'product_details' in st.session_state:
        context['product_details'] = st.session_state.product_details
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}\nCONTEXT: {context}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_solution_section(audience, product, temperature=0.7, offer=None):
    """Generate the Solution & Benefits section."""
    prompt_template = load_prompt_template("solution_benefits.txt")
    
    # Create context dictionary with any existing session state data
    context = {}
    if 'headline_content' in st.session_state:
        context['headline_content'] = st.session_state.headline_content
    if 'problem_content' in st.session_state:
        context['problem_content'] = st.session_state.problem_content
    if 'avatar_details' in st.session_state:
        context['avatar_details'] = st.session_state.avatar_details
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}\nCONTEXT: {context}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_authority_section(audience, product, temperature=0.7, offer=None):
    """Generate the Authority section."""
    prompt_template = load_prompt_template("authority.txt")
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_offer_section(audience, product, temperature=0.7, offer=None):
    """Generate the Offer & Bonus section."""
    prompt_template = load_prompt_template("offer_bonus.txt")
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    if offer:
        prompt += f"\nOFFER DETAILS: {offer}"
    else:
        prompt += "\nOFFER DETAILS: Create a compelling offer based on the product/service and audience."
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_social_proof_section(audience, product, temperature=0.7, offer=None):
    """Generate the Social Proof section."""
    prompt_template = load_prompt_template("social_proof.txt")
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_guarantees_section(audience, product, temperature=0.7, offer=None):
    """Generate the Guarantees section."""
    prompt_template = load_prompt_template("guarantees.txt")
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_cta_section(audience, product, temperature=0.7, offer=None):
    """Generate the Call to Action section."""
    prompt_template = load_prompt_template("cta.txt")
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_ps_section(audience, product, temperature=0.7, offer=None):
    """Generate the P.S. (Post-Data) section."""
    prompt_template = load_prompt_template("ps.txt")
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

def generate_final_cta_section(audience, product, temperature=0.7, offer=None):
    """Generate the Final Call to Action section."""
    prompt_template = load_prompt_template("final_cta.txt")
    
    # Create the prompt with the audience and product information
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    if offer:
        prompt += f"\nOFFER: {offer}"
    
    # Generate the content
    return generate_content(prompt, temperature)

# Define section functions mapping once
section_functions = {
    "Above the Fold (Encabezado)": generate_above_the_fold,
    "Problem (Problema)": generate_problem_section,
    "Solution & Benefits (Solución y Beneficios)": generate_solution_section,
    "Authority (Autoridad)": generate_authority_section,
    "Offer & Bonus (Oferta y Bonus)": generate_offer_section,
    "Social Proof (Prueba Social)": generate_social_proof_section,
    "Guarantees (Garantías)": generate_guarantees_section,
    "Call to Action (Llamada a la Acción)": generate_cta_section,
    "P.S. (Post-Data)": generate_ps_section,
    "Final Call to Action (Llamada Final)": generate_final_cta_section
}

# Then use it in the conditional
if submit_sales_page:
    # Validar entradas
    if validate_inputs(sales_page_audience, sales_page_product):
        try:
            # Store current inputs in session state
            st.session_state.last_audience = sales_page_audience
            st.session_state.last_product = sales_page_product
            st.session_state.last_offer = sales_page_offer
            
            # Crear un contenedor para el spinner en la columna 2
            with col2:
                with st.spinner("Generando sección de página de ventas...", show_time=True):
                    # Obtener la función correspondiente
                    generator_func = section_functions[sales_page_section]
                    
                    # Generar el contenido
                    generated_content = generator_func(
                        audience=sales_page_audience,
                        product=sales_page_product,
                        temperature=sales_page_temperature,
                        offer=sales_page_offer if sales_page_offer.strip() else None
                    )
                    
                    # Store the generated content in session state with a key based on the section
                    section_key = sales_page_section.replace(" ", "_").lower()
                    st.session_state.generated_sections[section_key] = generated_content
                
            # Mostrar el contenido generado fuera del bloque del spinner
            display_generated_content(col2, generated_content, "sales_page_section")
                
        except Exception as e:
            st.error(f"Error al generar la sección: {str(e)}")
    else:
        st.warning("Por favor, completa los campos de público objetivo y producto/servicio.")
else:
    # Display previously generated content if available for the selected section
    section_key = sales_page_section.replace(" ", "_").lower()
    if section_key in st.session_state.generated_sections and st.session_state.generated_sections[section_key]:
        display_generated_content(col2, st.session_state.generated_sections[section_key], "sales_page_section")