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

# Configuración del modelo (Placeholder para tu nuevo ID)
MODEL_ID = "N11100/deepseek-r1-sysadmin-pro" # Cambia esto por tu nuevo ID tras el entrenamiento

def load_model():
    print("Iniciando SysAdmin Engine...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16
        )
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True
        )
        return tokenizer, model
    except Exception as e:
        print(f"Error al cargar el modelo: {e}")
        return None, None

# CSS Personalizado para estética Next-Gen (Terminal Dark Mode)
custom_css = """
.gradio-container {
    background-color: #0d1117 !important;
    color: #c9d1d9 !important;
    font-family: 'Fira Code', 'Courier New', monospace !important;
}
.main_title {
    color: #58a6ff !important;
    text-align: center;
    text-transform: uppercase;
    letter-spacing: 2px;
}
.terminal_box {
    background-color: #161b22 !important;
    border: 1px solid #30363d !important;
    border-radius: 8px !important;
}
.unfiltered_badge {
    background-color: #da3633 !important;
    color: white !important;
    padding: 4px 8px;
    border-radius: 4px;
    font-weight: bold;
    font-size: 0.8em;
}
footer {display: none !important;}
"""

def process_sysadmin_task(instruction, temp, tokens):
    # Simulación de respuesta si el modelo no está cargado, o ejecución real
    if model is None:
        return "⚠️ Error: Modelo no detectado. Por favor, asegúrate de haber entrenado y subido el modelo con el ID correcto."
    
    messages = [{"role": "user", "content": instruction}]
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=tokens,
            do_sample=True,
            temperature=temp,
            pad_token_id=tokenizer.eos_token_id
        )
    
    return tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)

# Carga inicial (Opcional, puede fallar si aún no has entrenado el modelo)
tokenizer, model = None, None # load_model()

with gr.Blocks(css=custom_css, title="SysAdmin Next-Gen Dashboard") as demo:
    gr.Markdown("# 🚀 SysAdmin Next-Gen Dashboard", elem_classes="main_title")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🛡️ Status: <span class='unfiltered_badge'>UNFILTERED OBEDIENCE</span>")
            prompt = gr.Textbox(
                label="Comando o Instrucción de Automatización",
                placeholder="Ej: Script de Bash para rotar logs y enviarlos por SCP...",
                lines=8,
                elem_classes="terminal_box"
            )
            
            with gr.Row():
                temp = gr.Slider(0.1, 1.0, value=0.4, label="Precisión (Temperature)")
                tokens = gr.Slider(128, 4096, value=1024, step=128, label="Max Tokens")
            
            run_btn = gr.Button("EJECUTAR ACCIÓN", variant="primary")
            
            gr.Markdown("### ⚡ Presets de Administración")
            with gr.Row():
                btn_backup = gr.Button("Backup Script", size="sm")
                btn_docker = gr.Button("Docker Stack", size="sm")
                btn_ansible = gr.Button("Ansible Playbook", size="sm")

        with gr.Column(scale=2):
            gr.Markdown("### 🖥️ Salida del Sistema (Código Limpio + Markdown)")
            output = gr.Markdown(elem_classes="terminal_box")
            
            with gr.Row():
                copy_btn = gr.Button("📋 Copiar Código", size="sm")
                save_btn = gr.Button("💾 Guardar como .py/.sh", size="sm")

    # Lógica de los botones de Preset
    btn_backup.click(lambda: "Crea un script de Bash profesional para realizar un backup de la carpeta /var/www/html, comprimirlo en .tar.gz con la fecha actual y subirlo a un servidor remoto vía SCP.", outputs=prompt)
    btn_docker.click(lambda: "Genera un archivo docker-compose.yml para un entorno de producción con Nginx, PostgreSQL y Redis, incluyendo volúmenes y redes persistentes.", outputs=prompt)
    btn_ansible.click(lambda: "Escribe un playbook de Ansible para actualizar todos los paquetes del sistema en servidores Ubuntu y configurar un firewall básico con UFW permitiendo SSH y HTTP.", outputs=prompt)

    # Ejecución principal
    run_btn.click(fn=process_sysadmin_task, inputs=[prompt, temp, tokens], outputs=output)

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