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

def initialize_model(temperature=1.0):
    return genai.GenerativeModel(
        model_name="gemini-2.0-flash",
        generation_config={"temperature": temperature}
    )

def initialize_session_state():
    """Initialize all session state variables."""
    defaults = {
        'sections': {
            'content': {},     # Store generated content
            'context': {},     # Store context for other sections
            'metadata': {}     # Store JSON metadata
        },
        'inputs': {
            'audience': '',
            'product': '',
            'offer': ''
        }
    }
    
    for var, default in defaults.items():
        if var not in st.session_state:
            st.session_state[var] = default

initialize_session_state()

def generate_content(prompt, temperature=1.0):
    model = initialize_model(temperature)
    
    response = model.generate_content(prompt)
    result = response.text
    
    try:
        import json
        import re
        
        json_pattern = r'```json\s*(.*?)\s*```|(\{[\s\S]*"character"[\s\S]*\})'
        json_match = re.search(json_pattern, result, re.DOTALL)
        
        if json_match:
            json_str = json_match.group(1) if json_match.group(1) else json_match.group(2)
            
            json_data = json.loads(json_str.strip())
            
            st.session_state.sections['metadata'] = json_data
            
            metadata_fields = ['character', 'pain_points', 'emotional_triggers', 'obstacles', 'failed_solutions']
            for field in metadata_fields:
                if field in json_data:
                    st.session_state.sections['metadata'][field] = json_data[field]
        
            result = re.sub(json_pattern, '', result, flags=re.DOTALL).strip()
    except Exception as e:
        print(f"Error processing JSON metadata: {str(e)}")
    
    return result

def display_generated_content(column, content, key_prefix, section_name="sección"):
    """Display generated content with download buttons at top and bottom."""
    if not content:
        return
    
    clean_section_name = section_name.split(" (")[0].lower() if " (" in section_name else section_name.lower()
    
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"sales_page_{key_prefix.replace('_', '-')}_{timestamp}.txt"
    
    column.download_button(
        label=f"Descargar sección de {clean_section_name} de la Sale Page",
        data=content,
        file_name=filename,
        mime="text/plain",
        key=f"{key_prefix}_top_download"
    )
    
    column.markdown("### Contenido generado:")
    column.markdown(content)
    
    column.download_button(
        label=f"Descargar sección de {clean_section_name} de la Sale Page",
        data=content,
        file_name=filename,
        mime="text/plain",
        key=f"{key_prefix}_bottom_download"
    )
    
def validate_inputs(audience, product):
    """Validate that required inputs are provided."""
    return audience.strip() != "" and product.strip() != ""

st.set_page_config(layout="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.")
    
    # Add to the section selection list
    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)",
            "Offer Summary (Resumen de Oferta)",
            "Guarantees (Garantías)",
            "FAQ (Preguntas Frecuentes)",
            "Closing (Cierre)"
        ]
    )

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

# First define all generator functions
def load_prompt_template(filename):
    """Load a prompt template from the prompts directory."""
    if not filename:
        raise ValueError("Filename cannot be None")
        
    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 create_base_prompt(template_name, audience, product, offer=None, extra_params=None):
    """Create base prompt with context for all sections."""
    prompt_template = load_prompt_template(template_name)
    prompt = f"{prompt_template}\n\nAUDIENCE: {audience}\nPRODUCT/SERVICE: {product}"
    
    if offer:
        prompt += f"\nOFFER{'_DETAILS' if 'offer' in template_name else ''}: {offer}"
    
    if extra_params:
        for key, value in extra_params.items():
            prompt += f"\n{key}: {value}"
    
    return prompt

def update_section_context(section_name, content):
    """Update the context for a section after it's generated."""
    section_key = normalize_section_key(section_name)
    
    # Store in content dictionary
    if 'content' not in st.session_state.sections:
        st.session_state.sections['content'] = {}
    st.session_state.sections['content'][section_key] = content
    
    # Store in context dictionary with timestamp
    if 'context' not in st.session_state.sections:
        st.session_state.sections['context'] = {}
    
    st.session_state.sections['context'][section_key] = {
        'content': content,
        'timestamp': datetime.datetime.now().isoformat()
    }

def get_relevant_context(current_section):
    """Get relevant context based on current section"""
    context = {}
    
    section_key = normalize_section_key(current_section)
    
    if section_key == 'above_the_fold':
        return context
    
    specific_relationships = {
        'problem': ['above_the_fold'],
        'solution_&_benefits': ['problem'],
        'authority': ['problem', 'solution_&_benefits'],
        'offer_&_bonus': ['solution_&_benefits', 'problem'],
        'social_proof': ['offer_&_bonus', 'solution_&_benefits'],
        'offer_summary': ['offer_&_bonus', 'solution_&_benefits'],
        'guarantees': ['offer_&_bonus', 'social_proof'],
        'faq': ['offer_&_bonus', 'guarantees', 'problem'],
        'closing': ['offer_&_bonus', 'guarantees', 'problem', 'solution_&_benefits']
    }
    
    section_order = [
        "above_the_fold",
        "problem",
        "solution_&_benefits",
        "authority",
        "offer_&_bonus",
        "social_proof",
        "offer_summary",
        "guarantees",
        "faq",
        "closing"
    ]
    
    # First, add specific relationships if they exist
    if section_key in specific_relationships:
        for rel_section in specific_relationships[section_key]:
            if rel_section in st.session_state.sections['context']:
                context[rel_section] = st.session_state.sections['context'][rel_section]['content']
    
    # Then, for any section after "above_the_fold", add previous sections as context
    if section_key in section_order:
        current_index = section_order.index(section_key)
        if current_index > 0:
            for prev_section in section_order[:current_index][-3:]:
                if prev_section in st.session_state.sections['context'] and prev_section not in context:
                    context[prev_section] = st.session_state.sections['context'][prev_section]['content']
    
    return context

# Función auxiliar para normalizar nombres de secciones
def normalize_section_key(section_name):
    """Convierte nombres de sección a formato de clave normalizado."""
    # Extraer solo la parte en inglés antes del paréntesis
    if " (" in section_name:
        section_name = section_name.split(" (")[0]
    return section_name.replace(" ", "_").lower()

def generate_section(section_name, audience, product, temperature=0.7, offer=None):
    """Base generator function for all sections."""
    template_map = {
        "above_the_fold": "above_the_fold.txt",
        "problem": "problem.txt",
        "solution_&_benefits": "solution_benefits.txt",
        "authority": "authority.txt",
        "offer_&_bonus": "offer_bonus.txt",
        "social_proof": "social_proof.txt",
        "offer_summary": "offer_summary.txt",
        "guarantees": "guarantees.txt",
        "faq": "faq.txt",
        "closing": "closing.txt"
    }
    
    section_key = normalize_section_key(section_name)
    template_name = template_map.get(section_key)
    
    # Verificación más detallada para depuración
    if not template_name:
        # Mostrar todas las claves disponibles para ayudar en la depuración
        available_keys = ", ".join(template_map.keys())
        raise ValueError(f"No template found for section: '{section_name}' (normalized key: '{section_key}'). Available keys: {available_keys}")
    
    # Verificar que el archivo de plantilla existe
    prompt_path = os.path.join(os.path.dirname(__file__), "prompts", template_name)
    if not os.path.exists(prompt_path):
        raise FileNotFoundError(f"Template file not found: {prompt_path}")
    
    extra_params = {}
    
    # Add context for sections that need it
    context = get_relevant_context(section_name)
    if context:
        extra_params['CONTEXT'] = context
    
    # Special handling for specific sections
    if section_key == 'faq':
        extra_params['LANGUAGE'] = 'Spanish'
    elif section_key == 'closing':
        extra_params['ADDITIONAL_INSTRUCTIONS'] = """
- Generate content directly in Spanish
- Do not include labels like 'BINARY CHOICE', 'FINAL REMINDER', etc.
- Do not include comments or observations in any language
- Do not include technical notes or instructions
- Write the closing text naturally and fluidly"""
    
    prompt = create_base_prompt(template_name, audience, product, offer, extra_params)
    content = generate_content(prompt, temperature)
    
    # Update context after generation
    update_section_context(section_name, content)
    
    return content

section_functions = {
    section_name: lambda audience, product, temperature=0.7, offer=None, section=section_name: generate_section(section, audience, product, temperature, offer)
    for section_name in [
        "Above the Fold (Encabezado)",
        "Problem (Problema)",
        "Solution & Benefits (Solución y Beneficios)",
        "Authority (Autoridad)",
        "Offer & Bonus (Oferta y Bonus)",
        "Social Proof (Prueba Social)",
        "Offer Summary (Resumen de Oferta)",
        "Guarantees (Garantías)",
        "FAQ (Preguntas Frecuentes)",
        "Closing (Cierre)"
    ]
}

if submit_sales_page:
    if validate_inputs(sales_page_audience, sales_page_product):
        try:
            st.session_state.inputs.update({
                'audience': sales_page_audience,
                'product': sales_page_product,
                'offer': sales_page_offer
            })
            
            with col2:
                with st.spinner("Generando sección de página de ventas...", show_time=True):
                    generator_func = section_functions[sales_page_section]
                    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
                    )
                
            display_generated_content(col2, generated_content, "sales_page_section", 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:
    section_key = normalize_section_key(sales_page_section)
    if section_key in st.session_state.sections['content']:
        display_generated_content(col2, st.session_state.sections['content'][section_key], "sales_page_section", sales_page_section)