import time
import os
import joblib
import streamlit as st
import google.generativeai as genai
from dotenv import load_dotenv
from puv_formulas import puv_formulas

# Configuración global
CONFIG = {
    'model': {
        'role': 'ai',
        'avatar': '🤖'
    },
    'user': {
        'avatar': '👤'
    },
    'page': {
        'title': "RoboCopy - Creador de PUVs",
        'icon': "🚀",
        'layout': "wide"
    }
}

# Configuración de página
st.set_page_config(
    page_title=CONFIG['page']['title'],
    page_icon=CONFIG['page']['icon'],
    layout=CONFIG['page']['layout']
)

# Configuración de avatares e identificadores
MODEL_ROLE = 'ai'
AI_AVATAR_ICON = '🤖'
USER_AVATAR_ICON = '👤'

# === FUNCIONES AUXILIARES ===
def load_css(file_path):
    try:
        with open(file_path) as f:
            st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
    except Exception as e:
        st.markdown("""
        <style>
        .robocopy-title {
            color: #4ECDC4 !important;
            font-weight: bold;
            font-size: 2em;
        }
        </style>
        """, unsafe_allow_html=True)

def get_chat_title(messages):
    if not messages:
        return "Nuevo Chat"
    first_msg = messages[0]['content'] if messages else ""
    title = first_msg[:30] + "..." if len(first_msg) > 30 else first_msg
    return title

def get_formulas_for_prompt():
    prompt_text = "\nFÓRMULAS DE PROPUESTAS ÚNICAS DE VALOR (PUVs):\n\n"
    
    for key, formula in puv_formulas.items():
        prompt_text += f"🔹 {key}:\n"
        prompt_text += f"   - Descripción: {formula.get('description', 'Descripción no disponible').strip()}\n"
        
        if 'Structure:' in formula.get('description', ''):
            estructura = formula['description'].split('Structure:')[1].split('Key elements:')[0].strip()
            prompt_text += "   - Estructura Base:\n"
            for line in estructura.split('\n'):
                if line.strip():
                    prompt_text += f"     * {line.strip()}\n"
        
        ejemplos = formula.get('examples', [])[:2]
        if ejemplos:
            prompt_text += "   - Ejemplos destacados:\n"
            for i, ejemplo in enumerate(ejemplos):
                prompt_text += f"     {i+1}. Público objetivo: {ejemplo.get('target_audience', 'No especificado')}\n"
                prompt_text += f"       Servicio: {ejemplo.get('product_service', 'No especificado')}\n"
                prompt_text += f"       PUV: {ejemplo.get('uvp', 'Ejemplo no disponible')}\n"
        
        prompt_text += "\n" + "-"*50 + "\n"
    
    return prompt_text

def add_message(role, content, avatar):
    message = {
        'role': role,
        'content': content,
        'avatar': avatar
    }
    st.session_state.messages.append(message)
    return message

def update_chat_memory():
    st.session_state.chats_in_memory[st.session_state.current_chat_id].update({
        'messages': st.session_state.messages,
        'gemini_history': st.session_state.gemini_history,
        'title': st.session_state.chat_title
    })

def handle_model_error(error, retry_count, max_retries):
    if retry_count >= max_retries:
        error_message = f"Lo siento, estoy experimentando problemas para procesar tu solicitud. Por favor, intenta de nuevo más tarde. Error: {str(error)}"
        with st.chat_message(name=MODEL_ROLE, avatar=AI_AVATAR_ICON):
            st.error(error_message)
        add_message(MODEL_ROLE, error_message, AI_AVATAR_ICON)
        update_chat_memory()
        return True
    return False

def mostrar_con_efecto_escritura(texto, velocidad=0.05):
    placeholder = st.empty()
    contenido_actual = ""
    
    for caracter in texto:
        contenido_actual += caracter
        placeholder.markdown(contenido_actual + "▌")
        time.sleep(velocidad)
    
    placeholder.markdown(contenido_actual)
    return contenido_actual

def process_model_response(prompt, max_retries=3):
    retry_count = 0
    while retry_count < max_retries:
        try:
            # Crear el mensaje del modelo primero
            with st.chat_message(name=MODEL_ROLE, avatar=AI_AVATAR_ICON):
                # Añade un indicador de "escribiendo..."
                typing_indicator = st.empty()
                typing_indicator.markdown("*RoboCopy está escribiendo...*")
                
                response = st.session_state.chat.send_message(prompt, stream=True)
                mensaje_completo = ""
                mensaje_actual = ""
                
                # Procesar la respuesta por chunks con efecto de escritura
                for chunk in response:
                    mensaje_completo += chunk.text
                    for caracter in chunk.text:
                        mensaje_actual += caracter
                        typing_indicator.markdown(mensaje_actual + "▌")
                        time.sleep(0.01)  # Velocidad de escritura ajustable
                
                # Mostrar mensaje final
                typing_indicator.markdown(mensaje_completo)
            
            add_message(MODEL_ROLE, mensaje_completo, AI_AVATAR_ICON)
            st.session_state.gemini_history = st.session_state.chat.history
            update_chat_memory()
            return True
            
        except Exception as e:
            retry_count += 1
            if handle_model_error(e, retry_count, max_retries):
                return False
            wait_time = (1.5 ** retry_count)
            time.sleep(wait_time)
    return False

def handle_example_click(prompt_text):
    """Función para manejar clicks en ejemplos"""
    st.session_state.update({
        'show_examples': False,
        'messages': [],
        'current_chat_id': str(time.time()),
        'gemini_history': [],
        'chat_title': 'Nuevo Chat',
        'user_input': prompt_text  # Aquí pasamos el texto del ejemplo directamente
    })
    
    # Inicializar nuevo chat si es necesario
    if st.session_state.current_chat_id not in st.session_state.chats_in_memory:
        st.session_state.chats_in_memory[st.session_state.current_chat_id] = {
            'messages': [],
            'gemini_history': [],
            'title': 'Nuevo Chat'
        }
    
    st.rerun()  # Forzar recarga de la página
    
    # Eliminar todo el código restante de esta función
    # El procesamiento se hará en la sección de entrada principal
    if st.session_state.current_chat_id not in st.session_state.chats_in_memory:
        st.session_state.chats_in_memory[st.session_state.current_chat_id] = {
            'messages': [],
            'gemini_history': [],
            'title': 'Nuevo Chat'
        }
    
    try:
        title_response = st.session_state.model.generate_content(
            f"Título para consulta de ejemplo: '{prompt_text}' (máximo 4 palabras)"
        )
        st.session_state.chat_title = title_response.text.strip()[:25]
    except Exception as e:
        st.session_state.chat_title = f"Ejemplo-{time.strftime('%H:%M')}"
    
    st.session_state.chats_in_memory[st.session_state.current_chat_id]['title'] = st.session_state.chat_title
    process_model_response(prompt_text)

# === CONFIGURACIÓN INICIAL ===
load_dotenv()
GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
genai.configure(api_key=GOOGLE_API_KEY)

# Definición del prompt del sistema
SYSTEM_PROMPT = f"""
Eres un equipo colaborativo de expertos de clase mundial trabajando juntos para crear Propuestas Únicas de Valor (PUVs) excepcionales que conviertan a la audiencia en clientes.

EL EQUIPO DE EXPERTOS:
1. ESTRATEGA MAESTRO DE MARKETING:
   - Experto en marcos de propuestas de valor y estrategias de conversión
   - Asegura que las PUVs sigan la estructura de fórmula seleccionada con precisión
   - Se enfoca en la colocación estratégica de elementos clave de conversión

2. COPYWRITER ELITE DE RESPUESTA DIRECTA:
   - Crea ganchos, historias y elementos persuasivos convincentes
   - Elabora propuestas de valor irresistibles que impulsan conversiones
   - Asegura que el lenguaje resuene con la audiencia objetivo

3. ESPECIALISTA EN PSICOLOGÍA DE AUDIENCIA:
   - Experto en comprender las motivaciones y objeciones de la audiencia
   - Crea contenido que construye conexión genuina y confianza
   - Identifica y aborda miedos y deseos ocultos

{get_formulas_for_prompt()}

INSTRUCCIONES PARA CREAR PUVs:
1. Si el usuario no ha proporcionado información sobre su producto/servicio y audiencia objetivo, solicítala de manera amable y directa.
2. Si el usuario ha proporcionado información pero no ha elegido fórmula específica, pregúntale qué fórmula le gustaría utilizar.
3. Una vez con toda la información, crear propuestas de valor utilizando ÚNICAMENTE la fórmula elegida.
"""

WELCOME_MESSAGE = """
¡Hola! 👋 Soy RoboCopy, tu asistente especializado en crear Propuestas Únicas de Valor impactantes. 
¿En qué puedo ayudarte hoy?
"""

# === INICIALIZACIÓN DEL ESTADO ===
new_chat_id = str(time.time())

if 'chats_in_memory' not in st.session_state:
    st.session_state.update({
        'chats_in_memory': {},
        'current_chat_id': new_chat_id,
        'chat_title': 'Nuevo Chat',
        'messages': [],
        'show_examples': True,
        'gemini_history': []
    })

# === SIDEBAR ===
with st.sidebar:
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.image("assets/robocopy_logo.png", width=300)
    
    st.write('# Chats Anteriores')
    
    chat_options = [new_chat_id] + list(st.session_state.chats_in_memory.keys())
    current_index = chat_options.index(st.session_state.current_chat_id) if st.session_state.current_chat_id in chat_options else 0
    
    st.session_state.current_chat_id = st.selectbox(
        label='Selecciona un chat anterior',
        options=chat_options,
        index=current_index,
        format_func=lambda x: st.session_state.chats_in_memory.get(x, {}).get('title', 'Nuevo Chat'),
        key='chat_selector_unique'
    )
    
    if st.button('🗑️ Borrar Historial de Chat Actual'):
        if st.session_state.current_chat_id in st.session_state.chats_in_memory:
            del st.session_state.chats_in_memory[st.session_state.current_chat_id]
        
        # Reiniciar completamente el estado del chat actual
        st.session_state.update({
            'messages': [],
            'gemini_history': [],
            'chat_title': 'Nuevo Chat',
            'show_examples': True
        })
        
        # Reiniciar el chat con el modelo
        st.session_state.chat = st.session_state.model.start_chat(history=[])
        st.session_state.chat.send_message(SYSTEM_PROMPT)
        st.session_state.gemini_history = st.session_state.chat.history
        st.rerun()

# === CONFIGURACIÓN DEL MODELO ===
model = genai.GenerativeModel(model_name='gemini-2.0-flash')
st.session_state.model = model
st.session_state.chat = st.session_state.model.start_chat(history=st.session_state.gemini_history)

if not st.session_state.gemini_history:
    st.session_state.chat.send_message(SYSTEM_PROMPT)
    st.session_state.gemini_history = st.session_state.chat.history

# === MOSTRAR MENSAJES DEL HISTORIAL ===
# Crear el input primero (aparecerá al final)
input_placeholder = st.empty()
new_prompt = input_placeholder.chat_input('¿En qué puedo ayudarte hoy?')

# Contenedor principal para mensajes
chat_container = st.container()

with chat_container:
    # === SECCIÓN DE EJEMPLOS ===
    if st.session_state.show_examples and not st.session_state.messages:
        st.title("💡 RoboCopy - Asistente de PUVs")
        st.markdown("### Tu experto en crear Propuestas Únicas de Valor que convierten")
        st.markdown("### 🎯 Prueba estos ejemplos:")
        
        col1, col2 = st.columns([1, 1])
        
        with col1:
            if st.button("¿Cuál es la mejor fórmula para mi PUV? 🤔", use_container_width=True):
                handle_example_click("¿Podrías explicarme cuál es la mejor fórmula para crear una PUV efectiva para mi negocio?")
            
            if st.button("Necesito una PUV para mi tienda online 🛍️", use_container_width=True):
                handle_example_click("Quiero crear una PUV para mi tienda online de ropa sostenible dirigida a mujeres de 25-35 años")
        
        with col2:
            if st.button("Ayúdame a mejorar mi PUV actual ✨", use_container_width=True):
                handle_example_click("¿Podrías ayudarme a mejorar mi PUV actual para hacerla más persuasiva?")
            
            if st.button("Crear PUV para servicios profesionales 👔", use_container_width=True):
                handle_example_click("Necesito una PUV para mi servicio de consultoría en marketing digital")
    
    st.markdown("---")
    
    # Crear un contenedor para los mensajes que empuja el input hacia abajo
    messages_container = st.container()
    
    # Mostrar mensajes existentes
    for message in st.session_state.messages:
        with st.chat_message(
            name=message['role'],
            avatar=AI_AVATAR_ICON if message['role'] == MODEL_ROLE else USER_AVATAR_ICON
        ):
            st.markdown(message['content'])
    
    # Procesar entrada del usuario si existe
    if 'user_input' in st.session_state:
        prompt = st.session_state.user_input
        del st.session_state.user_input
        
        with st.chat_message("user", avatar=USER_AVATAR_ICON):
            st.markdown(prompt)
        add_message("user", prompt, USER_AVATAR_ICON)
        
        try:
            title_response = st.session_state.model.generate_content(
                f"Título para consulta: '{prompt}' (máximo 4 palabras)"
            )
            st.session_state.chat_title = title_response.text.strip()[:25]
        except Exception as e:
            st.session_state.chat_title = f"Chat-{time.strftime('%H:%M')}"
        
        st.session_state.chats_in_memory[st.session_state.current_chat_id] = {
            'messages': st.session_state.messages,
            'gemini_history': st.session_state.gemini_history,
            'title': st.session_state.chat_title
        }
        
        process_model_response(prompt)
        update_chat_memory()
    
    # Procesar nueva entrada si existe
    if new_prompt:
        with st.chat_message("user", avatar=USER_AVATAR_ICON):
            st.markdown(new_prompt)
        add_message("user", new_prompt, USER_AVATAR_ICON)
        
        try:
            title_response = st.session_state.model.generate_content(
                f"Título para consulta: '{new_prompt}' (máximo 4 palabras)"
            )
            st.session_state.chat_title = title_response.text.strip()[:25]
        except Exception as e:
            st.session_state.chat_title = f"Chat-{time.strftime('%H:%M')}"
        
        st.session_state.chats_in_memory[st.session_state.current_chat_id] = {
            'messages': st.session_state.messages,
            'gemini_history': st.session_state.gemini_history,
            'title': st.session_state.chat_title
        }
        
        process_model_response(new_prompt)
        update_chat_memory()
        # En la sección del SIDEBAR
        if st.button('📝 Nuevo Chat'):
            st.session_state.update({
                'current_chat_id': str(time.time()),
                'messages': [],
                'gemini_history': [],
                'chat_title': 'Nuevo Chat',
                'show_examples': True
            })
            st.session_state.chat = st.session_state.model.start_chat(history=[])
            st.session_state.chat.send_message(SYSTEM_PROMPT)
            st.session_state.gemini_history = st.session_state.chat.history
            st.rerun()