import os
import subprocess
import time
import gradio as gr
import requests
from huggingface_hub import hf_hub_download

# 1. Ajuste al repositorio oficial de Google (GGUF Balanceado)
REPO_ID = "google/gemma-4-12B-it-qat-q4_0-gguf"
FILENAME = "gemma-4-12b-it-qat-q4_0.gguf"
BIN_PATH = "./serverLlama"

# Definimos la ruta de destino limpia
MODEL_PATH = "/app/gemma-4-12b.gguf" if os.path.exists("/app") else "./gemma-4-12b.gguf"

if not os.path.exists(MODEL_PATH):
    print(f"Descargando {FILENAME} desde el repositorio oficial de Google...")
    descarga = hf_hub_download(
        repo_id=REPO_ID, 
        filename=FILENAME,
        local_dir="/app" if os.path.exists("/app") else "."
    )
    # Renombramos el archivo para que coincida con la ruta de tu serverLlama
    os.rename(descarga, MODEL_PATH)
    print("¡Descarga oficial completada con éxito!")

# 2. Darle permisos de ejecución a tu binario subido
print("Configurando permisos para serverLlama...")
os.chmod(BIN_PATH, 0o755)

# 3. Arrancar tu binario serverLlama en segundo plano con 32k de contexto
print("Encendiendo serverLlama con 32k de contexto en segundo plano...")
subprocess.Popen([
    BIN_PATH,
    "--model", MODEL_PATH,
    "--host", "127.0.0.1",
    "--port", "8080",
    "--threads", "2",
    "--ctx-size", "32768",
# 🔥 EXTREMO - SOLO PARA PRUEBAS (Necesita 65GB+ RAM)
# "--ctx-size", "262144",        # 262K - Límite oficial del modelo
# "--ctx-size", "229376",        # 224K - Casi máximo
# "--ctx-size", "196608",        # 192K - Muy pesado
# "--ctx-size", "163840",        # 160K - Extremo

# ⚠️ MUY ALTO (Necesita 40GB+ RAM)
# "--ctx-size", "131072",        # 128K - Análisis de libros completos
# "--ctx-size", "114688",        # 112K - Documentos muy largos
# "--ctx-size", "98304",         # 96K  - Investigación profunda

# 📊 ALTO (Necesita 27GB+ RAM)
# "--ctx-size", "81920",         # 80K  - Análisis extenso
# "--ctx-size", "73728",         # 72K  - Documentos largos
# "--ctx-size", "65536",         # 64K  - Límite recomendado para 32GB RAM

    
 #  "--ctx-size", "24576",    
    "--n-predict", "4096",
 # Optimizaciones de caché
    "--cache-type-k", "q4_0",
    "--cache-type-v", "q4_0",
    
    # Batch size para 2 vCPUs
    # "--batch-size", "4096",
    # "--ubatch-size", "4096",
    "--batch-size", "1536",
    "--ubatch-size", "768",
    
])

# Esperamos a que el backend en C++ reserve la memoria y levante el puerto
time.sleep(10)
print("¡Servidor backend en C++ listo!")

# 4. Función de comunicación entre Gradio y tu binario local
def predict_api(prompt, system_prompt="Eres un asistente experto y conciso."):
    if not prompt.strip():
        return "Por favor, escribe un prompt válido."

    headers = {"Content-Type": "application/json"}
    data = {
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.7
    }
    try:
        response = requests.post("http://127.0.0.1:8080/v1/chat/completions", json=data, headers=headers)
        if response.status_code == 200:
            return response.json()["choices"][0]["message"]["content"].strip()
        return f"Error del backend (Código {response.status_code}): {response.text}"
    except Exception as e:
        return f"Error al conectar con serverLlama: {e}"

# 5. Interfaz Gráfica Estándar de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Backend de IA - Sara (serverLlama + Gradio 32k)")
    
    with gr.Row():
        with gr.Column():
            txt_prompt = gr.Textbox(label="Prompt / Pregunta", lines=4, placeholder="Escribe aquí...")
            txt_system = gr.Textbox(label="System Prompt", value="Eres un asistente experto.", lines=2)
            btn_enviar = gr.Button("Enviar", variant="primary")
        with gr.Column():
            txt_salida = gr.Textbox(label="Respuesta de Sara", lines=8, interactive=False)
            
    btn_enviar.click(fn=predict_api, inputs=[txt_prompt, txt_system], outputs=txt_salida)
    
    # Endpoint expuesto para tu script local
    api_endpoint = gr.Interface(
        fn=predict_api,
        inputs=[gr.Textbox(elem_id="prompt"), gr.Textbox(elem_id="system_prompt")],
        outputs=gr.Textbox(),
        api_name="predict_api"
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)