import os
import subprocess
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
import sysconfig

# Inicializar la aplicación FastAPI
app = FastAPI()

# Definir la estructura de entrada para el prompt de la API
class PromptRequest(BaseModel):
    prompt: str

# Ruta principal para ejecutar el modelo
@app.post("/generate/")
async def generate_text(request: PromptRequest):
    model_path = "/content/executorch/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8.pte"
    tokenizer_path = "/content/executorch/tokenizer.model"
    
    # Descargar los archivos si no existen
    download_files(model_path, tokenizer_path)
    
    prompt = request.prompt
    result = run_llama_model(model_path, tokenizer_path, prompt)
    
    return {"generated_text": result}

def run_command(command):
    """Ejecutar un comando en el shell y devolver la salida."""
    result = subprocess.run(command, shell=True, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if result.returncode != 0:
        return f"Error ejecutando el comando: {result.stderr}"
    return result.stdout

def download_files(model_path, tokenizer_path):
    """Descargar el modelo y tokenizador si no están presentes."""
    if not os.path.exists(model_path):
        print(f"Descargando el modelo desde Hugging Face: {model_path}")
        run_command(f"wget https://huggingface.co./executorch-community/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8-ET/resolve/main/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8.pte -O {model_path}")
    else:
        print(f"El modelo ya está presente en: {model_path}")
    
    if not os.path.exists(tokenizer_path):
        print(f"Descargando el tokenizador desde Hugging Face: {tokenizer_path}")
        run_command(f"wget https://huggingface.co./executorch-community/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8-ET/resolve/main/tokenizer.model -O {tokenizer_path}")
    else:
        print(f"El tokenizador ya está presente en: {tokenizer_path}")

def run_llama_model(model_path, tokenizer_path, prompt):
    """Ejecutar el modelo Llama y generar texto."""
    cmd = f"cd /content/executorch/cmake-out/examples/models/llama && ./llama_main --model_path={model_path} --tokenizer_path={tokenizer_path} --prompt='{prompt}'"
    return run_command(cmd)

def install_requirements(lib_path):
    """Instalar dependencias de ejecutorch y del modelo Llama."""
    # Instalar dependencias para Executorch
    print("Instalando dependencias de Executorch...")
    run_command(f"cd /content/executorch && CMAKE_PREFIX_PATH={lib_path} EXECUTORCH_BUILD_XNNPACK=ON bash ./install_requirements.sh --pybind")
    print("Dependencias de Executorch instaladas.")

    # Instalar requerimientos adicionales para el modelo Llama
    print("Instalando dependencias para el modelo Llama...")
    run_command("cd /content/executorch/examples/models/llama && ./install_requirements.sh")
    print("Requerimientos de Llama instalados.")

# Iniciar el servidor Uvicorn directamente desde el código Python
if __name__ == "__main__":
    # Obtener la ruta del entorno Python
    lib_path = sysconfig.get_paths()["purelib"]
    print(f"Usando la ruta de la biblioteca Python: {lib_path}")

    # Instalar las dependencias necesarias
    install_requirements(lib_path)

    # Ejecutar el servidor Uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)