import base64
import datetime
import gradio as gr
import numpy as np
import os
import pytz
import psutil
import re
import random
import torch
import time
import shutil  # Added for zip functionality

from PIL import Image
from io import BytesIO
from diffusers import DiffusionPipeline, LCMScheduler, AutoencoderTiny

try:
    import intel_extension_for_pytorch as ipex
except:
    pass

SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)
HF_TOKEN = os.environ.get("HF_TOKEN", None)
# check if MPS is available OSX only M1/M2/M3 chips
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available()
device = torch.device(
    "cuda" if torch.cuda.is_available() else "xpu" if xpu_available else "cpu"
)
torch_device = device
torch_dtype = torch.float16


# add file save and download and clear:
# Function to create a zip file from a list of files
def create_zip(files):
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    zip_filename = f"images_{timestamp}.zip"
    with zipfile.ZipFile(zip_filename, 'w') as zipf:
        for file in files:
            zipf.write(file, os.path.basename(file))
    return zip_filename

# Function to encode a file to base64
def encode_file_to_base64(file_path):
    with open(file_path, "rb") as file:
        encoded = base64.b64encode(file.read()).decode()
    return encoded

# Function to save all images as a zip file and provide a base64 download link
def save_all_images(images):
    if len(images) == 0:
        return None, None

    zip_filename = create_zip(images)  # Create a zip file from the list of image files
    zip_base64 = encode_file_to_base64(zip_filename)  # Encode the zip file to base64
    download_link = f'<a href="data:application/zip;base64,{zip_base64}" download="{zip_filename}">Download All</a>'

    return zip_filename, download_link


# Function to clear all image files
def clear_all_images():
    base_dir = os.getcwd()  # Get the current base directory
    img_files = [file for file in os.listdir(base_dir) if file.lower().endswith((".png", ".jpg", ".jpeg"))]  # List all files ending with ".jpg" or ".jpeg"
    
    # Remove all image files
    for file in img_files:
        os.remove(file)

# Function to handle "Save All" button click
def save_all_button_click():
    images = [file for file in os.listdir() if file.lower().endswith((".png", ".jpg", ".jpeg"))]
    zip_filename, download_link = save_all_images(images)
    
    if download_link:
        gr.write(download_link)

    zip_file = create_zip_of_files(all_files)
    gr.HTML(get_zip_download_link(zip_file))



# Function to handle "Clear All" button click
def clear_all_button_click():
    clear_all_images()




print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
print(f"TORCH_COMPILE: {TORCH_COMPILE}")
print(f"device: {device}")

if mps_available:
    device = torch.device("mps")
    torch_device = "cpu"
    torch_dtype = torch.float32

if SAFETY_CHECKER == "True":
    pipe = DiffusionPipeline.from_pretrained("Lykon/dreamshaper-7")
else:
    pipe = DiffusionPipeline.from_pretrained("Lykon/dreamshaper-7", safety_checker=None)

pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to(device=torch_device, dtype=torch_dtype).to(device)
pipe.unet.to(memory_format=torch.channels_last)
pipe.set_progress_bar_config(disable=True)

# check if computer has less than 64GB of RAM using sys or os
if psutil.virtual_memory().total < 64 * 1024**3:
    pipe.enable_attention_slicing()

if TORCH_COMPILE:
    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
    pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
    pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)

# Load LCM LoRA
pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
pipe.fuse_lora()

def safe_filename(text):
    """Generate a safe filename from a string."""
    safe_text = re.sub(r'\W+', '_', text)
    timestamp = datetime.datetime.now().strftime("%Y%m%d")
    return f"{safe_text}_{timestamp}.png"
    
def encode_image(image):
    """Encode image to base64."""
    buffered = BytesIO()
    #image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

def fake_gan():
    base_dir = os.getcwd()  # Get the current base directory
    img_files = [file for file in os.listdir(base_dir) if file.lower().endswith((".png", ".jpg", ".jpeg"))]  # List all files ending with ".jpg" or ".jpeg"
    images = [(random.choice(img_files), os.path.splitext(file)[0]) for file in img_files]
    return images
    
def predict(prompt, guidance, steps, seed=1231231):
    generator = torch.manual_seed(seed)
    last_time = time.time()
    results = pipe(
        prompt=prompt,
        generator=generator,
        num_inference_steps=steps,
        guidance_scale=guidance,
        width=512,
        height=512,
        # original_inference_steps=params.lcm_steps,
        output_type="pil",
    )
    print(f"Pipe took {time.time() - last_time} seconds")
    nsfw_content_detected = (
        results.nsfw_content_detected[0]
        if "nsfw_content_detected" in results
        else False
    )
    if nsfw_content_detected:
        nsfw=gr.Button("🕹️NSFW🎨", scale=1)

    try: 
        central = pytz.timezone('US/Central')
        safe_date_time = datetime.datetime.now().strftime("%Y%m%d")
        replaced_prompt = prompt.replace(" ", "_").replace("\n", "_")
        safe_prompt = "".join(x for x in replaced_prompt if x.isalnum() or x == "_")[:90]
        filename = f"{safe_date_time}_{safe_prompt}.png"
        
        # Save the image
        if len(results.images) > 0:
            image_path = os.path.join("", filename)  # Specify your directory
            results.images[0].save(image_path)
            print(f"#Image saved as {image_path}")
            encoded_image = encode_image(image)
            html_link = f'<a href="data:image/png;base64,{encoded_image}" download="{filename}">Download Image</a>'
            #gr.Markdown(html_link)
    except:
        return results.images[0]

    return results.images[0] if len(results.images) > 0 else None


css = """
#container{
    margin: 0 auto;
    max-width: 40rem;
}
#intro{
    max-width: 100%;
    text-align: center;
    margin: 0 auto;
}
"""
with gr.Blocks(css=css) as demo:
    
    # Add "Save All" button with emoji
    save_all_button = gr.Button("💾 Save All", scale=1)
    # Add "Clear All" button with emoji
    clear_all_button = gr.Button("🗑️ Clear All", scale=1)

    # Add buttons to the Streamlit app
    gr.Button(save_all_button)
    gr.Button(clear_all_button)

    # Attach click event handlers to the buttons
    save_all_button.click(save_all_button_click)
    clear_all_button.click(clear_all_button_click)

    with gr.Column(elem_id="container"):
        gr.Markdown(
            """##🕹️ Real Time 🎨 ImageGen Gallery 🌐""",
            elem_id="intro",
        )
        with gr.Row():
            with gr.Row():
                prompt = gr.Textbox(
                    placeholder="Insert your prompt here:", scale=5, container=False
                )
                generate_bt = gr.Button("Generate", scale=1)

        # Image Result from last prompt
        image = gr.Image(type="filepath")

        # Gallery of Generated Images with Image Names in Random Set to Download
        with gr.Row(variant="compact"):
            text = gr.Textbox(
                label="Image Sets",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
            )
            btn = gr.Button("Generate Gallery of Saved Images")
        gallery = gr.Gallery(
            label="Generated Images", show_label=False, elem_id="gallery"
        )

        # Advanced Generate Options
        with gr.Accordion("Advanced options", open=False):
            guidance = gr.Slider(
                label="Guidance", minimum=0.0, maximum=5, value=0.3, step=0.001
            )
            steps = gr.Slider(label="Steps", value=4, minimum=2, maximum=10, step=1)
            seed = gr.Slider(
                randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1
            )

        # Diffusers
        with gr.Accordion("Run with diffusers"):
            gr.Markdown(
                """## Running LCM-LoRAs it with `diffusers`
            ```bash
            pip install diffusers==0.23.0
            ```
            
            ```py
            from diffusers import DiffusionPipeline, LCMScheduler

            pipe = DiffusionPipeline.from_pretrained("Lykon/dreamshaper-7").to("cuda") 
            pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
            pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5") #yes, it's a normal LoRA

            results = pipe(
                prompt="ImageEditor",
                num_inference_steps=4,
                guidance_scale=0.0,
            )
            results.images[0]
            ```
            """
            )

        # Function IO Eventing and Controls
        inputs = [prompt, guidance, steps, seed]
        generate_bt.click(fn=predict, inputs=inputs, outputs=image, show_progress=False)
        btn.click(fake_gan, None, gallery)
        prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
        guidance.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
        steps.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
        seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)

demo.queue()
demo.launch()