# 📌 INSTALAR DEPENDENCIAS (si no están en tu entorno)
# pip install transformers torch gradio whisper gtts accelerate

import gradio as gr
import torch
import random
from transformers import AutoModelForCausalLM, AutoTokenizer
import whisper
from gtts import gTTS

# 📌 Verificar si se tiene GPU disponible
device = "cuda" if torch.cuda.is_available() else "cpu"

# 📌 MODELO ESPAÑOL PÚBLICO (con código personalizado)
modelo_id = "fbarragan/helloworld_model"

# Intentar importar accelerate sin fallar
try:
    import accelerate
    accelerate_installed = True
except ImportError:
    accelerate_installed = False

# Cargar el modelo de manera segura
modelo = AutoModelForCausalLM.from_pretrained(
    modelo_id,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    trust_remote_code=True
).to(device)

# 📌 LISTA NEGRA DE INSULTOS
palabras_prohibidas = [
    "Abollado", "Arrecho(a)", "Azote de barrio", "Bolsa", "Cagao E zamuro",
    "Ladilla", "Mojón", "Ñero", "Piche", "Rata(ón)", "Vuelto mierda"
]

# 📌 DICCIONARIO VENEZOLANO
diccionario_venezolano = {
    "dinero": "palo",
    "amigo": "pana",
    "genial": "chévere",
    "problema": "culebra",
    "fiesta": "bonche",
    "hola": "¡Epa mi pana!",
    "adiós": "¡Váyalo, chamo!",
}

# 📌 CARGAR MODELO (requiere autenticación si es privado)
modelo_id = "deepseek-ai/DeepSeek-R1"
tokenizer = AutoTokenizer.from_pretrained(modelo_id, trust_remote_code=True)  # 🌟 Clave para modelos con código personalizado
modelo = AutoModelForCausalLM.from_pretrained(
    modelo_id,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True  # 🌟 Necesario para DeepSeek-R1
)

# 📌 LISTA NEGRA DE INSULTOS (previamente generada)
palabras_prohibidas = [
    "Abollado", "Arrecho(a)", "Azote de barrio", "Bolsa", "Cagao E zamuro", "Caimanera",
    "Cariaquito morao", "Chimbo, todo", "Chimbo", "Chola", "Conejo", "Coño'e madre",
    "Culebra", "Flux", "Friqueado", "Furruqueado", "Guaya", "Güevo", "Hediondo", "Ladilla",
    "Landro", "Mamagüevo", "Mandinga", "Mapanare", "Marico triste", "Metra",
    "Mojón", "Muérgano", "Ñero", "Oracio", "Pasoneado(a)", "Pava", "Peaje", "Pecueca",
    "Piche", "Rata(ón)", "Raya'o", "Rolo", "Salado", "Tachón", "Taguara", "Vergajo",
    "Violín", "Vuelto mierda"
]

# 📌 CARGAR DICCIONARIO VENEZOLANO COMPLETO
def cargar_diccionario(ruta_archivo):
    diccionario = {
        "saludo": [],
        "despedida": [],
        "afirmacion": [],
        "negacion": [],
        "colloquial": {}  # Palabras venezolanas y sus equivalentes formales
    }
    with open(ruta_archivo, "r", encoding="utf-8") as f:
        for linea in f:
            linea = linea.strip()
            if linea and not linea.startswith("#"):
                partes = linea.split(".-", 1)
                if len(partes) >= 2:
                    palabra = partes[0].strip()
                    definicion = partes[1].strip().lower()
                    # Clasificar palabras según contexto
                    if "saludo" in definicion:
                        diccionario["saludo"].append(palabra)
                    elif "despedida" in definicion:
                        diccionario["despedida"].append(palabra)
                    elif "afirmación" in definicion or "positivo" in definicion:
                        diccionario["afirmacion"].append(palabra)
                    elif "negación" in definicion or "negativo" in definicion:
                        diccionario["negacion"].append(palabra)
                    else:
                        # Agregar palabras coloquiales con su significado
                        if palabra not in palabras_prohibidas:
                            diccionario["colloquial"][palabra] = definicion
    return diccionario

diccionario_venezolano = cargar_diccionario("Diccionario_Venezolano.txt")

# 📌 Función para adaptar respuestas
def adaptar_respuesta(respuesta):
    # Filtrar insultos
    for palabra in palabras_prohibidas:
        respuesta = respuesta.replace(palabra, "****")
    
    # Reemplazar términos formales por coloquiales
    for formal, venezolano in diccionario_venezolano.items():
        respuesta = respuesta.replace(formal, venezolano)
    
    return respuesta

# 📌 Función principal del asistente
def chamo_asistente(pregunta):
    inputs = tokenizer(pregunta, return_tensors="pt").to(device)  # Enviar a GPU si está disponible
    inputs = tokenizer(pregunta, return_tensors="pt")
    output = modelo.generate(**inputs, max_new_tokens=150)
    respuesta = tokenizer.decode(output[0], skip_special_tokens=True)
    return adaptar_respuesta(respuesta)

# 📌 Procesar audio
def procesar_audio(archivo):
    modelo_whisper = whisper.load_model("small")  # Cargar modelo una sola vez
    texto = modelo_whisper.transcribe(archivo)["text"]
    respuesta = chamo_asistente(texto)
    
    # Convertir texto a audio
    texto = whisper.load_model("small").transcribe(archivo)["text"]
    respuesta = chamo_asistente(texto)
    tts = gTTS(respuesta, lang="es")
    tts.save("respuesta.mp3")
    return respuesta, "respuesta.mp3"

# 📌 Interfaz Gradio
iface_texto = gr.Interface(fn=chamo_asistente, inputs="text", outputs="text")
iface_voz = gr.Interface(fn=procesar_audio, inputs="microphone", outputs=["text", "audio"])
demo = gr.TabbedInterface([iface_texto, iface_voz], ["Texto", "Voz"])
demo.launch()
demo.launch()