import os
import subprocess

# Definir la ruta de la carpeta donde se almacenan los archivos Python
carpeta = os.getcwd()  # Ruta actual de trabajo

# Definir las variables de entorno
model_hf_dir = os.getenv("MODEL_HF_DIR")
model_gguf_dir = os.getenv("MODEL_GGUF_DIR", os.path.join(carpeta, "MODEL_GGUF_DIR"))
quantization_types = os.getenv("QUANTIZATION_TYPES", "Q8_0, Q6_K, Q5_K, Q4_K, Q3_K, Q2_K")

# Comando base para ejecutar los archivos Python
comando_base = "python3"

# Lista de comandos predefinidos
comandos_predefinidos = [
    f"{comando_base} convert.py {model_hf_dir} --outtype f16",
    f"{comando_base} convert-hf-to-gguf.py {model_hf_dir}",
    f"{comando_base} convert-lora-to-ggml.py {model_hf_dir}",
]

# Variable para controlar si se encontró un archivo ejecutado exitosamente
ejecucion_exitosa = False

# Ejecutar los comandos predefinidos uno por uno
for comando in comandos_predefinidos:
    try:
        # Ejecutar el comando
        print(f"Ejecutando: {comando}")
        resultado = subprocess.run(comando, shell=True, check=True)
        # Si el comando se ejecuta sin errores, detener el bucle
        ejecucion_exitosa = True
        print(f"Ejecución exitosa del comando: {comando}")
        break
    except subprocess.CalledProcessError:
        # Si el comando falla, continuar con el siguiente comando predefinido
        print(f"Error al ejecutar {comando}, continuando con el siguiente comando predefinido...")

# Si se ejecutó exitosamente, realizar las siguientes acciones
if ejecucion_exitosa:
    # Ruta del archivo ggml-model-f16.gguf en la carpeta MODEL_HF_DIR
    archivo_gguf = os.path.join(model_hf_dir, "ggml-model-f16.gguf")

    # Crear la carpeta MODEL_GGUF_DIR si no existe
    os.makedirs(model_gguf_dir, exist_ok=True)

    # Mover el archivo ggml-model-f16.gguf a MODEL_GGUF_DIR
    nueva_ubicacion_gguf = os.path.join(model_gguf_dir, "ggml-model-f16.gguf")
    os.rename(archivo_gguf, nueva_ubicacion_gguf)
    print(f"Archivo ggml-model-f16.gguf movido a {nueva_ubicacion_gguf}")

    # Obtener los tipos de cuantización de QUANTIZATION_TYPES y separarlos por comas
    tipos_cuantizacion = quantization_types.split(",")

    # Ejecutar ./quantize para cada tipo de cuantización
    for tipo in tipos_cuantizacion:
        try:
            comando_quantize = f"./quantize {nueva_ubicacion_gguf} {tipo}"
            print(f"Ejecutando: {comando_quantize}")
            # Ejecutar el comando quantize
            subprocess.run(comando_quantize, shell=True, check=True)
        except subprocess.CalledProcessError:
            # Si el comando falla, continuar con el siguiente tipo de cuantización
            print(f"Error al ejecutar {comando_quantize}, continuando con el siguiente tipo de cuantización...")

# Mostrar los archivos en la carpeta model_hf_dir
print(f"\nContenido de la carpeta {model_hf_dir}:")
contenido = os.listdir(model_hf_dir)
for archivo in contenido:
    print(archivo)

# Mostrar el contenido del archivo de registro "convert.log" si existe
log_path = os.path.join(carpeta, "convert.log")
if os.path.exists(log_path):
    print("\nContenido de convert.log:")
    with open(log_path, 'r') as log_file:
        print(log_file.read())
else:
    print("\nEl archivo convert.log no existe.")