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

# Usar CONFIG en lugar de constantes separadas
st.set_page_config(
    page_title=CONFIG['page']['title'],
    page_icon=CONFIG['page']['icon'],
    layout=CONFIG['page']['layout']
)

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)

load_dotenv()
GOOGLE_API_KEY=os.environ.get('GOOGLE_API_KEY')
genai.configure(api_key=GOOGLE_API_KEY)

new_chat_id = f'{time.time()}'
# Configuración de avatares e identificadores
MODEL_ROLE = 'ai'
AI_AVATAR_ICON = '🤖'
USER_AVATAR_ICON = '👤'

# Función para manejar títulos de chat de manera unificada
def get_chat_title(messages):
    if not messages:
        return "Nuevo Chat"
    first_msg = messages[0]['content'] if messages else ""
    # Limitar el título a los primeros 30 caracteres
    title = first_msg[:30] + "..." if len(first_msg) > 30 else first_msg
    return title

def get_formulas_for_prompt():
    """Genera texto formateado con las fórmulas PUV para incluir en el prompt del sistema"""
    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"
        
        # Extraer estructura de la descripció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():  # Ignorar líneas vacías
                    prompt_text += f"     * {line.strip()}\n"
        
        # Manejar ejemplos
        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

# Definición COMPLETA del prompt multipersona para el 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?
"""

# Inicializar el estado de la sesión
# Inicialización unificada
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 allows a list of past chats
with st.sidebar:
    # Centrar el logo
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.image("assets/robocopy_logo.png", width=300)
    
    st.write('# Chats Anteriores')
    
    # Simplificar la lógica del selector de chat
    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
    
    # Modificar a:
    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'
    )
    
    # Botón para borrar el historial
    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]
            st.session_state.messages = []
            st.session_state.gemini_history = []
            st.session_state.chat_title = 'Nuevo Chat'
            st.experimental_rerun()  # Evitar loops

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

# Si es un chat nuevo, enviar el prompt del sistema como primer mensaje
if not st.session_state.gemini_history:
    # Enviamos el prompt del sistema como primer mensaje (invisible para el usuario)
    st.session_state.chat.send_message(SYSTEM_PROMPT)
    st.session_state.gemini_history = st.session_state.chat.history

# Mostrar mensajes del historial
for message in st.session_state.messages:
    with st.chat_message(name=message['role'], avatar=message.get('avatar')):
        st.markdown(message['content'])

# React to user input
# Funciones auxiliares para manejo de mensajes y errores
def add_message(role, content, avatar):
    """Función centralizada para añadir mensajes al historial"""
    message = {
        'role': role,
        'content': content,
        'avatar': avatar
    }
    st.session_state.messages.append(message)
    return message

def update_chat_memory():
    """Función centralizada para actualizar el chat en memoria"""
    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):
    """Función centralizada para manejar errores del modelo"""
    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

# === Añadir esta función en la sección de helpers ===
def mostrar_con_efecto_escritura(texto, velocidad=0.05):
    """Muestra texto con efecto de escritura en Streamlit"""
    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

# === Modificar en process_model_response ===
def process_model_response(prompt, max_retries=3):
    """Función centralizada para procesar respuestas del modelo"""
    retry_count = 0
    while retry_count < max_retries:
        try:
            response = st.session_state.chat.send_message(prompt, stream=True)
            full_text = "".join(chunk.text for chunk in response)
            
            # Llamada corregida
            mensaje_mostrado = mostrar_con_efecto_escritura(full_text, velocidad=0.05)
            
            add_message(MODEL_ROLE, mensaje_mostrado, 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 = (2 ** retry_count) + (time.time() % 1)
            time.sleep(wait_time)
    return False

# Modificar la función handle_example_click para usar las nuevas funciones
# Simplificar a:
def handle_example_click(prompt_text):
    st.session_state.update({
        'show_examples': False,
        'messages': [],
        'current_chat_id': str(time.time()),
        'gemini_history': [],
        'chat_title': get_chat_title([])
    })
    
    # Añadir esta verificación crítica
    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'
        }
    
    process_model_response(prompt_text)

# Modificar la sección principal de manejo de chat
if prompt := st.chat_input('¿En qué puedo ayudarte hoy?'):
    is_first_message = len(st.session_state.messages) == 0
    
    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'
        }
        
    add_message('user', prompt, USER_AVATAR_ICON)
    with st.chat_message('user', avatar=USER_AVATAR_ICON):
        st.markdown(prompt)
    
    if is_first_message:
        add_message(MODEL_ROLE, WELCOME_MESSAGE, AI_AVATAR_ICON)
        with st.chat_message(name=MODEL_ROLE, avatar=AI_AVATAR_ICON, key="welcome_msg"):
            st.markdown(WELCOME_MESSAGE)
        update_chat_memory()
        st.experimental_rerun()
    
    process_model_response(prompt)

# Mostrar ejemplos solo si show_examples es True y no hay mensajes previos
if st.session_state.show_examples and not st.session_state.messages:
    # Usar contenedor para ancho completo
    main_container = st.container()
    
    with main_container:
        st.title("💡 RoboCopy - Asistente de PUVs")
        st.markdown("### Tu experto en crear Propuestas Únicas de Valor que convierten")
        
        st.markdown("### 🎯 Prueba estos ejemplos:")
        
        # Usar columnas con proporción ajustada para mejor distribución
        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("---")