import gradio as gr
from huggingface_hub import InferenceClient
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM
import os
import torch

title = "Generador y clasificador de coloquialismos cubanos"
description = "Genera texto con el dialecto cubano y detecta expresiones coloquiales cubanas."
article = """
## Modelo:
Se usa el somosnlp-hackathon-2026/patriae-cuban-dialect-model.

## Fine-tuning:
Se utilizó el dataset somosnlp-hackathon-2026/patriae-cuba-literature-dataset para hacer fine-tuning del modelo gia-uh/cecilia-2b-instruct-v1.

## Participantes del equipo Patriae en SomosNLP 2026:
* **Carlos Luis Barnés Infante (https://huggingface.co/blacknoize404)**
* **Yisel Clavel Quintero (https://huggingface.co/clavel)**
* **Dionis Lopez (https://huggingface.co/inoid)**
"""

HF_TOKEN = os.environ.get("HF_TOKEN", None)

# Cargar el modelo para clasificación
classifier = pipeline(
    "text-classification", 
    model='somosnlp-hackathon-2026/patriae-cuban-dialect-model', 
    token=HF_TOKEN 
)

repo_id = 'somosnlp-hackathon-2026/patriae-cuban-dialect-model'
chat_available = False
generator = None
tokenizer_chat = None
error_message = ""

try:
    # Cargar tokenizador
    tokenizer_chat = AutoTokenizer.from_pretrained(repo_id, token=HF_TOKEN)
    print("✅ Tokenizador cargado")
    
    # Configurar token de padding
    if tokenizer_chat.pad_token is None:
        tokenizer_chat.pad_token = tokenizer_chat.eos_token
        print("✅ Token de padding configurado")
    
    # Cargar modelo para generación
    model_chat = AutoModelForCausalLM.from_pretrained(
        repo_id, 
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto" if torch.cuda.is_available() else None,
        token=HF_TOKEN,
        low_cpu_mem_usage=True
    )
    print("✅ Modelo cargado")
    
    # Crear pipeline de generación de texto con parámetros optimizados
    generator = pipeline(
        "text-generation",
        model=model_chat,
        tokenizer=tokenizer_chat,
        device=0 if torch.cuda.is_available() else -1,
        max_new_tokens=100,  # Reducido para velocidad
        temperature=0.7,
        top_p=0.95,
        do_sample=True,
        repetition_penalty=1.1  # Evita repeticiones
    )
    
    chat_available = True
    print("✅ Modelo de chat cargado correctamente")
    
except Exception as e:
    error_message = str(e)
    print(f"⚠️ Error cargando modelo de chat: {e}")
    chat_available = False
    generator = None
    tokenizer_chat = None


def chat_response(message, history):
    """
    Función para generar respuestas en dialecto cubano usando el modelo de chat
    """
    if not chat_available:
        yield f"⚠️ El modelo de chat no está disponible.\n\n**Error:** {error_message}\n\nUsa el clasificador por ahora."
        return
    
    if not message or message.strip() == "":
        yield "🤔 ¿Qué quieres que te diga? Dame algo con qué trabajar."
        return
    
    try:
        # Construir el prompt con contexto más corto para velocidad
        prompt = f"""Eres un cubano conversando con un amigo. Responde con naturalidad usando el dialecto cubano:\nUsuario: {message}\nCubano:"""
        
        # Generar respuesta con parámetros optimizados
        result = generator(
            prompt,
            max_new_tokens=80,  # Límite de tokens para respuesta más rápida
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.15,
            pad_token_id=tokenizer_chat.eos_token_id if tokenizer_chat else None,
            eos_token_id=tokenizer_chat.eos_token_id if tokenizer_chat else None
        )
        
        # Extraer la respuesta generada
        generated_text = result[0]['generated_text']
        # Quitar el prompt de la respuesta
        response = generated_text.replace(prompt, "").strip()
        
        # Si no hay respuesta, dar un mensaje por defecto
        if not response:
            response = "¡Ay, mi hermano! No sé qué decirte, pero pa' que veas que estoy aquí escuchando."
        
        # Limitar respuesta a 200 caracteres para velocidad
        if len(response) > 200:
            response = response[:200] + "..."
        
        yield response
        
    except Exception as e:
        yield f"❌ Error en el chat: {str(e)}"


def classify_text(text):
    """
    Función para clasificar texto y detectar coloquialismos cubanos
    """
    if not text or text.strip() == "":
        return "⚠️ Por favor, ingresa un texto para clasificar."
    
    try:
        # Obtener resultados del modelo
        results = classifier(text)
        
        # Verificar que results no esté vacío
        if not results or len(results) == 0:
            return "❌ No se obtuvieron resultados del modelo."
        
        # Obtener el resultado con mayor confianza
        top_result = results[0]
        label = top_result['label']
        score = top_result['score']
        
        # Determinar si es coloquial (LABEL_1 = coloquial, LABEL_0 = no coloquial)
        if label == "LABEL_1":
            emoji = "🇨🇺"
            tipo = "**COLOQUIAL CUBANO**"
            descripcion = "✅ Este texto contiene expresiones típicas del dialecto cubano."
        elif label == "LABEL_0":
            emoji = "📚"
            tipo = "**NO COLOQUIAL / ESTÁNDAR**"
            descripcion = "ℹ️ Este texto usa lenguaje estándar, sin expresiones coloquiales cubanas."
        else:
            emoji = "🔍"
            tipo = f"**{label}**"
            descripcion = f"Etiqueta desconocida: {label}"
        
        # Construir la salida formateada
        output = f"""
### {emoji} Resultado: {tipo}

**Confianza:** {score:.2%}

{descripcion}

"""
        return output
        
    except Exception as e:
        return f"❌ Error en clasificación: {str(e)}"

    
# Interfaz de Chat
chat_interface = gr.ChatInterface(
    fn=chat_response,
    title="💬 Chat en Dialecto Cubano",
    description="Conversa con un modelo entrenado en dialecto cubano.",
    examples=[
        "¿Qué bolá?",
        "Cuéntame qué hiciste hoy.",
        "Dime un refrán cubano."
    ]
)

# Interfaz de Clasificación
classifier_interface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(
        label="Texto a clasificar",
        placeholder="Escribe una frase en español cubano...",
        lines=5
    ),
    outputs=gr.Markdown(label="Resultado"),
    title="Clasificador de coloquialismos cubanos",
    description="Detecta si un texto contiene expresiones coloquiales cubanas.",
    examples=[
        ["Nos marchábamos sin ti, creyendo que se te habían pegado las sábanas."],
        ["De Varadero venía para acá porque en su casa las broncas con el padre eran de ampanga."],
        ["Los pueblos que no se conocen han de darse prisa para conocerse, como quienes van a pelear juntos."],
        ["El cielo está despejado y hace buen tiempo."],
        ["¿Qué bolá? ¿Cómo andas? ¡Tremenda jodedera!"],
    ]
)

# Combinar ambas interfaces en una sola app
with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}")
    gr.Markdown(description) 

    with gr.Tab("🔍 Clasificador"):
        classifier_interface.render()
    
    with gr.Tab("💬 Chat"):
        if chat_available:
            chat_interface.render()
        else:
            gr.Markdown(f"""
            ## ⚠️ Modelo de chat no disponible. 
            
            **Error:** `{error_message}`
            
            """)
    
    gr.Markdown(article)

if __name__ == "__main__":
    demo.launch()