import io
import globales
import herramientas
import gradio_client
from huggingface_hub import InferenceClient
import conexion_firebase

def genera_platillo_gpu(platillo):

    prompt = globales.previo + platillo
    print("Platillo enviado:", platillo)    

    kwargs = {
            "prompt": prompt,
            "api_name": "/infer"
        }    

    try:

        client = gradio_client.Client(globales.espacio, hf_token=globales.llave)
        result = client.predict(**kwargs,
        # prompt=prompt,
        # negative_prompt="",
        # seed=42,
        # randomize_seed=True,
        width=786,
        height=568,
        # guidance_scale=3.5,
        # num_inference_steps=28,
        # api_name="/infer"       
        )

        #Cuando es GPU, debe de restar segundos disponibles de HF 
        herramientas.restaSegundosGPU(globales.work_cost)

        print("Platillo generado:", platillo)
        return result[0]

    except Exception as e:     
        print("Excepción es: ", e)
        # Opción para regresar imagen genérica.
        # return "default.png"
        return '{"Error 500": e}'


def genera_platillo_inference(platillo):

    print("Proveedor:", globales.proveedor)
    modelo_actual = conexion_firebase.obtenDato('nowme', 'huggingface', 'modelo_actual')
    modelo = modelo_actual

    print("Modelo:", modelo)

    prompt = globales.previo + platillo
    print("Platillo enviado:", platillo)
       
    client = InferenceClient(
            provider= globales.proveedor,
            api_key=globales.llave
        )    

    try: 
        image = client.text_to_image(
        prompt,
        model=modelo
        )                 
        
    except Exception as e:
        print("Excepción es: ", e)
        if "Gateway Time-out" in str(e): 
            print("GATEWAY TIME-OUT 💀")
            modelo=globales.inferencia_backup
            #Escribe en txt el nuevo modelo.
            herramientas.modificaModeloActual(modelo)        
        return f"Error: {e}"
    
    #Cuando es GPU, debe de restar segundos disponibles de HF 
    herramientas.restaSegundosGPU(globales.work_cost)
    img_io = io.BytesIO()
    image.save(img_io, "PNG")
    img_io.seek(0)
    print("Platillo generado:", platillo)
    return img_io