import os
import zipfile
import numpy as np
import random
from pathlib import Path
from PIL import Image
import streamlit as st
from huggingface_hub import AsyncInferenceClient
import asyncio

MAX_SEED = np.iinfo(np.int32).max
client = AsyncInferenceClient()
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)

PREDEFINED_SEED = random.randint(0, MAX_SEED)

async def generate_image(prompt, width, height, seed, model_name):
    try:
        if seed == -1:
            seed = PREDEFINED_SEED
        seed = int(seed)
        image = await client.text_to_image(
            prompt=prompt, height=height, width=width, model=model_name
        )
        return image, seed
    except Exception as e:
        return f"Error al generar imagen: {e}", None

def save_prompt(prompt_text, seed):
    try:
        prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
        with open(prompt_file_path, "w") as prompt_file:
            prompt_file.write(prompt_text)
        return prompt_file_path
    except Exception as e:
        st.error(f"Error al guardar el prompt: {e}")
        return None

async def gen(prompt, width, height, model_name):
    combined_prompt = prompt  
    seed = PREDEFINED_SEED
    progress_bar = st.progress(0)
    image, seed = await generate_image(combined_prompt, width, height, seed, model_name)
    progress_bar.progress(100)

    if isinstance(image, str) and image.startswith("Error"):
        progress_bar.empty()
        return [image, None]

    image_path = save_image(image, seed)
    prompt_file_path = save_prompt(combined_prompt, seed)

    return [str(image_path), str(prompt_file_path)]

def save_image(image, seed):
    try:
        image_path = DATA_PATH / f"image_{seed}.jpg"
        image.save(image_path, format="JPEG")
        return image_path
    except Exception as e:
        st.error(f"Error al guardar la imagen: {e}")
        return None

def get_storage():
    files = [file for file in DATA_PATH.glob("*.jpg") if file.is_file()]
    files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
    usage = sum([file.stat().st_size for file in files])
    return [str(file.resolve()) for file in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"

def get_prompts():
    prompt_files = [file for file in DATA_PATH.glob("*.txt") if file.is_file()]
    return {file.stem.replace("prompt_", ""): file for file in prompt_files}

def delete_all_images():
    try:
        files = [file for file in DATA_PATH.glob("*.jpg")]
        prompts = [file for file in DATA_PATH.glob("*.txt")]
        for file in files + prompts:
            os.remove(file)
        st.success("Todas las imágenes y prompts han sido borrados.")
    except Exception as e:
        st.error(f"Error al borrar archivos: {e}")

def download_images_as_zip():
    zip_path = DATA_PATH / "images.zip"
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file in DATA_PATH.glob("*.jpg"):
            zipf.write(file, arcname=file.name)
    with open(zip_path, "rb") as zip_file:
        st.download_button(label="Descargar imágenes en .zip", data=zip_file, file_name="images.zip", mime="application/zip")

def main():
    st.set_page_config(layout="wide")
    st.title("Generador de Imágenes")
    st.warning("Este espacio contiene contenido que no es adecuado para todas las audiencias. Se recomienda discreción.")
    agree = st.checkbox("Soy mayor de 18 años y entiendo que el contenido puede no ser apropiado.")

    if agree:
        prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
        format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
        model_option = st.sidebar.selectbox("Modelo", ["enhanceaiteam/Flux-uncensored", "enhanceaiteam/Flux-Uncensored-V2"])
        
        width, height = (720, 1280) if format_option == "9:16" else (1280, 720)

        if st.sidebar.button("Generar Imagen"):
            with st.spinner("Generando imagen..."):
                result = asyncio.run(gen(prompt, width, height, model_option))
                image_paths = result[0]
                prompt_file = result[1]

            if image_paths:
                if Path(image_paths).exists():
                    st.image(image_paths, caption="Imagen Generada")
                else:
                    st.error("El archivo de imagen no existe.")

                if prompt_file and Path(prompt_file).exists():
                    prompt_text = Path(prompt_file).read_text()
                    st.write(f"Prompt utilizado: {prompt_text}")
                else:
                    st.write("El archivo del prompt no está disponible.")

        files, usage = get_storage()
        st.text(usage)
        cols = st.columns(6)
        prompts = get_prompts()

        for idx, file in enumerate(files):
            with cols[idx % 6]:
                image = Image.open(file)
                prompt_file = prompts.get(Path(file).stem.replace("image_", ""), None)
                prompt_text = Path(prompt_file).read_text() if prompt_file else "No disponible"
                
                st.image(image, caption=f"Imagen {idx+1}")
                st.write(f"Prompt: {prompt_text}")
                
                if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"):
                    try:
                        os.remove(file)  
                        if prompt_file:
                            os.remove(prompt_file)
                        st.success(f"Imagen {idx+1} y su prompt fueron borrados.")
                    except Exception as e:
                        st.error(f"Error al borrar la imagen o prompt: {e}")

        if st.button("Borrar todas las imágenes"):
            delete_all_images()

        download_images_as_zip()

    else:
        st.error("No puedes acceder a este espacio hasta que confirmes que eres mayor de edad.")

if __name__ == "__main__": 
    main()