import wget
import os
import asyncio
from fastapi import FastAPI
from pydantic import BaseModel
import executorch
import uvicorn

# Definir los parámetros de la solicitud utilizando Pydantic
class Request(BaseModel):
    # Valores por defecto para las rutas de los archivos
    model_path: str = "Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8.pte"
    tokenizer_path: str = "tokenizer.model"
    
    # Valores por defecto para la temperatura y longitud de secuencia
    prompt: str
    temperature: float = 0.7
    seq_len: int = 256

# Crear una instancia de FastAPI
app = FastAPI()

# Cargar el modelo y el tokenizer globalmente para evitar recargarlo en cada solicitud
model = None
tokenizer = None

# Función para descargar el modelo y el tokenizer en la misma carpeta del código (ahora es asíncrona)
async def download_files(model_url: str, tokenizer_url: str, model_path: str, tokenizer_path: str):
    # Verificar si el archivo ya existe, si no, descargarlo
    if not os.path.exists(model_path):
        print(f"Descargando el modelo desde: {model_url}")
        # Usamos wget de manera asíncrona
        await asyncio.to_thread(wget.download, model_url, model_path)

    if not os.path.exists(tokenizer_path):
        print(f"Descargando el tokenizer desde: {tokenizer_url}")
        # Usamos wget de manera asíncrona
        await asyncio.to_thread(wget.download, tokenizer_url, tokenizer_path)

# Cargar el modelo y el tokenizer (se hace una vez al inicio) (ahora es asíncrona)
async def load_model(request: Request):
    global model, tokenizer
    
    # Si ya están cargados, no cargamos de nuevo
    if model is not None and tokenizer is not None:
        print("Modelo y tokenizer ya están cargados en la memoria.")
        return
    
    # URLs de los archivos a descargar
    model_url = "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"
    tokenizer_url = "https://huggingface.co/executorch-community/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8-ET/resolve/main/tokenizer.model"
    
    # Obtener la ruta del directorio actual
    current_dir = os.path.dirname(os.path.realpath(__file__))
    
    # Definir las rutas locales donde se guardarán los archivos
    local_model_path = os.path.join(current_dir, request.model_path)
    local_tokenizer_path = os.path.join(current_dir, request.tokenizer_path)

    # Descargar los archivos si no existen (se hace de forma asíncrona)
    await download_files(model_url, tokenizer_url, local_model_path, local_tokenizer_path)

    # Cargar el modelo y tokenizer descargados
    print("Cargando el modelo y tokenizer en memoria...")
    model = executorch.load_model(local_model_path)
    tokenizer = executorch.load_tokenizer(local_tokenizer_path)
    
    # Configurar el modelo con los parámetros recibidos
    model.set_temperature(request.temperature)
    model.set_max_length(request.seq_len)
    print("Modelo y tokenizer cargados en memoria.")

# Ruta para generar texto (ahora es asíncrona)
@app.post("/generate/")
async def generate_text(request: Request):
    # Cargar el modelo si no ha sido cargado
    if model is None or tokenizer is None:
        await load_model(request)  # Esperamos la carga del modelo de forma asíncrona
    
    # Generar el texto con el prompt recibido
    output = model.generate(request.prompt)
    
    # Devolver el texto generado
    return {"generated_text": output}

# Código para ejecutar uvicorn desde el archivo Python
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)