import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Usar un modelo más pequeño pero compatible con español
model_name = "DeepESP/gpt2-spanish"
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
model = AutoModelForCausalLM.from_pretrained(model_name)

# Conocimiento sobre ecommerce y ventas online
ECOMMERCE_INFO = {
    "productos": "Para vender productos online necesitas una tienda virtual. Los precios dependen del tipo de producto y mercado objetivo.",
    "tienda": "Una tienda online básica puede costar desde $500 hasta $5000 dependiendo de funcionalidades.",
    "plataforma": "Puedes usar plataformas como Shopify, WooCommerce, PrestaShop o Magento para tu tienda.",
    "pagos": "Los métodos de pago más comunes son tarjetas de crédito, PayPal, transferencias y pago contra reembolso.",
    "marketing": "Es importante invertir en SEO, redes sociales y posiblemente publicidad para atraer clientes.",
    "envíos": "Los costos de envío dependen del peso, dimensiones y distancia. Puedes trabajar con empresas de logística.",
    "legal": "Necesitas cumplir con normativas de comercio electrónico, protección de datos y derechos del consumidor.",
    "cómo": "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?",
    "cómo empezar": "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?",
    "cómo vender": "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?",
    "cómo vender productos": "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?",
    "cómo vender productos online": "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?",
    "cómo vender productos online fácilmente": "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?",
    "cómo vender productos online fácilmente y barato": "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?",
    "cómo vender productos online fácilmente y barato y rápidamente": "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?",
    
}

def get_ecommerce_response(query):
    """Genera una respuesta basada en palabras clave del ecommerce."""
    query = query.lower()
    
    # Buscar palabras clave en la consulta
    for keyword, info in ECOMMERCE_INFO.items():
        if keyword in query:
            return info
    
    # Respuesta genérica si no hay coincidencias
    return "Para empezar una tienda online necesitas elegir una plataforma, configurar métodos de pago, añadir productos y gestionar envíos. ¿Tienes alguna pregunta específica sobre precios, plataformas o marketing?"

def generate_model_response(prompt, max_length=100):
    """Genera texto usando el modelo (backup si las reglas no coinciden)."""
    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            max_length=max_length,
            do_sample=True,
            temperature=0.7,
            top_p=0.92,
            top_k=50,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Extraer solo la parte después del prompt
    if prompt in response:
        response = response[len(prompt):].strip()
    
    return response

def answer(message, history):
    """Genera respuestas para preguntas de ecommerce."""
    # Si es una pregunta sobre ecommerce, usar respuestas predefinidas
    ecommerce_response = get_ecommerce_response(message)
    
    # Construir prompt con contexto
    prompt = "Esta es una conversación sobre ventas online y comercio electrónico:\n\n"
    
    # Añadir historial reciente (últimos 3 intercambios como máximo)
    recent_history = history[-3:] if len(history) > 3 else history
    for h in recent_history:
        prompt += f"Cliente: {h['content']}\n"
        if h['role'] == 'assistant':
            prompt += f"Asistente: {h['content']}\n"
    
    # Añadir mensaje actual
    prompt += f"Cliente: {message}\nAsistente: "
    
    # Si tenemos respuesta de ecommerce, usarla directamente
    if ecommerce_response:
        return ecommerce_response
    
    # Si no, generar respuesta con el modelo
    model_response = generate_model_response(prompt)
    
    # Limpiar la respuesta
    if "\n" in model_response:
        model_response = model_response.split("\n")[0]
    
    # Si la respuesta está vacía o es muy corta, dar respuesta por defecto
    if len(model_response) < 10:
        return "Puedo ayudarte con preguntas sobre cómo vender tus productos online, crear una tienda virtual o estrategias de marketing digital. ¿Qué te gustaría saber específicamente?"
    
    return model_response

demo = gr.ChatInterface(
    fn=answer,
    title="IA de Cikode - Asistente de Ecommerce",
    description="Haz preguntas sobre cómo vender tus productos online, costos de implementación, plataformas y más.",
    type="messages")

demo.queue(default_concurrency_limit=1).launch()