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:
            response = st.session_state.chat.send_message(prompt, stream=True)
            full_text = "".join(chunk.text for chunk in response)
            mensaje_mostrado = mostrar_con_efecto_escritura(full_text, velocidad=0.01)
            
            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 = (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  # Añadir el texto al estado
    })
    
    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]
            st.session_state.messages = []
            st.session_state.gemini_history = []
            st.session_state.chat_title = 'Nuevo Chat'
            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 ===
for message in st.session_state.messages:
    with st.chat_message(name=message['role'], avatar=message.get('avatar', USER_AVATAR_ICON if message['role'] == 'user' else AI_AVATAR_ICON)):
        st.markdown(message['content'])

# === SECCIÓN DE EJEMPLOS ===
if st.session_state.show_examples and not st.session_state.messages:
    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:")
        
        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("---")

# === ENTRADA DEL USUARIO ===
if 'user_input' in st.session_state:
    prompt = st.session_state.user_input
    del st.session_state.user_input  # Limpiar después de usar
else:
    prompt = st.chat_input('¿En qué puedo ayudarte hoy?')

if prompt:
    # Simular el envío del mensaje
    with st.chat_message("user", avatar=USER_AVATAR_ICON):
        st.markdown(prompt)  # Cambiado de prompt_text a prompt
    add_message("user", prompt, USER_AVATAR_ICON)  # Cambiado de prompt_text a prompt
    
    # Procesar la respuesta
    process_model_response(prompt)  # Cambiado de prompt_text a prompt
    
    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'
        }
    
    try:
        title_response = st.session_state.model.generate_content(
            f"Genera un título corto (máximo 5 palabras) que describa esta consulta, sin comillas: '{prompt}'"
        )
        generated_title = title_response.text.strip().replace('"', '')[:30]
        st.session_state.chat_title = generated_title or f"Chat-{st.session_state.current_chat_id}"
    except Exception as e:
        st.session_state.chat_title = f"Chat-{st.session_state.current_chat_id}"
    
    st.session_state.chats_in_memory[st.session_state.current_chat_id]['title'] = st.session_state.chat_title
    
    if is_first_message:
        add_message(MODEL_ROLE, WELCOME_MESSAGE, AI_AVATAR_ICON)
        with st.chat_message(name=MODEL_ROLE, avatar=AI_AVATAR_ICON):
            st.markdown(WELCOME_MESSAGE)
    update_chat_memory()
    
    process_model_response(prompt)