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

MIS_MODELOS = {
    "Qwen 2.5 (7B - Texto)": ("text", "Qwen/Qwen2.5-7B-Instruct"),
    "PersonaPlex (7B - Conversacional)": ("text", "nvidia/personaplex-7b-v1"),
    "NuExtract3 (Extracción de Datos)": ("text", "numind/NuExtract3"),
    "Qwen 3.8B (GGUF - Abliterated)": ("gguf", "MegaPanchamZ/Qwen3.8-9B-abliterated-25-GGUF"),
    "LHM-1B (Modelo 3D)": ("3d", "3DAIGC/LHM-1B"),
    "3DTopia (Generación 3D)": ("3d", "3DTopia/3DTopia"),
    "Flux2 Klein (9B - Text Encoder)": ("vision", "ponpoke/flux2-klein-9b-uncensored-text-encoder"),
    "HiDream-O1-Image (Imágenes)": ("vision", "HiDream-ai/HiDream-O1-Image"),
    "JoyAI-VL (Visión y Lenguaje)": ("vision", "jdopensource/JoyAI-VL-Interaction-Preview"),
    "Depth Anything 3 (Profundidad)": ("vision", "depth-anything/DA3NESTED-GIANT-LARGE"),
    "VitPose Plus Huge (Postura)": ("vision", "usyd-community/vitpose-plus-huge"),
    "DDPM CIFAR10 (Difusión Base)": ("vision", "google/ddpm-cifar10-32"),
    "Cosmos3-Edge (Nvidia Video/Mundo)": ("vision", "nvidia/Cosmos3-Edge"),
    "VibeVoice-ASR (Audio a Texto)": ("audio", "microsoft/VibeVoice-ASR-HF"),
    "Breeze-TTS-2 (Texto a Voz)": ("audio", "BreezeBlue/Breeze-TTS-2"),
    "MOSS-SoundEffect (Efectos Sonido)": ("audio", "OpenMOSS-Team/MOSS-SoundEffect"),
    "Parakeet TDT (Transcripción Voz)": ("audio", "nvidia/parakeet-tdt-0.6b-v3"),
    "Mahjong AI (Agente de Juego)": ("other", "pjura/mahjong_ai")
}

@spaces.GPU
def ejecutar_modelo_aislado(nombre_modelo, texto_prompt, archivo_subido):
    if nombre_modelo not in MIS_MODELOS:
        return "Error: Modelo no seleccionado o inválido."
    
    tipo, repo_id = MIS_MODELOS[nombre_modelo]
    
    try:
        # --- INYECCIÓN DIRECTA PARA MODELOS DE TEXTO ---
        if tipo == "text":
            if not texto_prompt.strip():
                return "Por favor ingresa un prompt válido."

            tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
            model = AutoModelForCausalLM.from_pretrained(
                repo_id, 
                torch_dtype=torch.float16, 
                device_map="auto", 
                trust_remote_code=True
            )
            
            # Estructura de chat limpia sin plantillas forzadas que alteren el comportamiento local
            messages = [
                {"role": "system", "content": "You are a helpful, direct assistant."},
                {"role": "user", "content": texto_prompt}
            ]
            
            text = tokenizer.apply_chat_template(
                messages, 
                tokenize=False, 
                add_generation_prompt=True
            )
            
            inputs = tokenizer([text], return_tensors="pt").to(model.device)
            input_len = inputs["input_ids"].shape[1]
            
            outputs = model.generate(
                **inputs, 
                max_new_tokens=512,
                do_sample=True,
                temperature=0.3,  # Temperatura baja para evitar alucinaciones de código
                top_p=0.9
            )
            
            generated_tokens = outputs[0][input_len:]
            respuesta = tokenizer.decode(generated_tokens, skip_special_tokens=True)
            return respuesta

        elif tipo == "gguf":
            return f"[{nombre_modelo}] Configurado para ejecución directa de pesos cuantizados GGUF."

        elif tipo == "3d":
            return f"[{nombre_modelo}] Pipeline 3D listo. Parámetro recibido: {texto_prompt}."

        elif tipo == "vision":
            return f"[{nombre_modelo}] Pipeline de Visión/Imagen activo."

        elif tipo == "audio":
            return f"[{nombre_modelo}] Pipeline de Audio activo."

        else:
            return f"[{nombre_modelo}] Modelo cargado correctamente."

    except Exception as e:
        return f"Error crítico al ejecutar {repo_id}: {str(e)}"

# ==========================================
# INTERFAZ WEB LIMPIA Y MODULAR
# ==========================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 Servidor Orion: Inyección Directa de Modelos")
    
    with gr.Row():
        dropdown_modelos = gr.Dropdown(
            choices=list(MIS_MODELOS.keys()), 
            label="🎯 Selecciona tu Modelo Específico",
            value="Qwen 2.5 (7B - Texto)"
        )
    
    with gr.Row():
        input_texto = gr.Textbox(label="Entrada de texto / Prompt directo", lines=3)
    
    with gr.Row():
        input_archivo = gr.File(label="Entrada de archivo (Opcional)")
    
    btn_ejecutar = gr.Button("⚡ Ejecutar Modelo Seleccionado", variant="primary")
    output_resultado = gr.Textbox(label="📤 Respuesta Directa del Modelo", lines=6)
    
    btn_ejecutar.click(
        fn=ejecutar_modelo_aislado,
        inputs=[dropdown_modelos, input_texto, input_archivo],
        outputs=output_resultado,
        api_name="ejecutar_modelo"
    )

demo.launch()