import datetime
import os
import re
import gc
import json
import time
import base64
import io
import tempfile
import zipfile
import PIL
from dataclasses import dataclass
from io import BytesIO
def sanitize_filename(filename):
    """Sanitizes a filename by replacing special characters with underscores"""
    return re.sub(r'[\\/*?:"<>|]', "_", filename)

from typing import Optional, Literal, Union
from diffusers import DiffusionPipeline, DDIMScheduler
HF_TOKEN = os.environ.get("HF_TOKEN")
import streamlit as st
st.set_page_config(layout="wide")
import torch
from diffusers import (
    StableDiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionInpaintPipeline,
    StableDiffusionImg2ImgPipeline,
)
from PIL import Image
from PIL.PngImagePlugin import PngInfo

from datetime import datetime
from threading import Thread

import requests

from huggingface_hub import HfApi
from huggingface_hub.utils._errors import RepositoryNotFoundError
from huggingface_hub.utils._validators import HFValidationError
from loguru import logger
from PIL.PngImagePlugin import PngInfo
from st_clickable_images import clickable_images

import streamlit.components.v1 as components

prefix = 'image_generation'

def dict_to_style(d):
    return ';'.join(f'{k}:{v}' for k, v in d.items())

def clickable_images(images, titles, div_style={}, img_style={}):
    """Generates a component with clickable images"""
    img_tag = "".join(
        f'<a href="{img}" target="_blank"><img src="{img}" title="{title}" style="{dict_to_style(img_style)}"></a>'
        for img, title in zip(images, titles)
    )
    return components.html(f'<div style="{dict_to_style(div_style)}">{img_tag}</div>', scrolling=True)

def display_and_download_images(output_images, metadata):
    with st.spinner("Preparing images..."):
        # save images to a temporary directory
        with tempfile.TemporaryDirectory() as tmpdir:
            gallery_images = []
            for i, image in enumerate(output_images):
                image.save(os.path.join(tmpdir, f"{i + 1}.png"), pnginfo=metadata)
                with open(os.path.join(tmpdir, f"{i + 1}.png"), "rb") as img:
                    encoded = base64.b64encode(img.read()).decode()
                    gallery_images.append(f"data:image/png;base64,{encoded}")

            _ = clickable_images(
                gallery_images,
                titles=[f"Image #{str(i + 1)}" for i in range(len(gallery_images))],
                div_style={"display": "flex", "justify-content": "center", "flex-wrap": "wrap"},
                img_style={"margin": "5px", "height": "200px"},
            )

            
PIPELINE_NAMES = Literal["txt2img", "inpaint", "img2img"]

DEFAULT_PROMPT = "a sprinkled donut sitting on top of a purple cherry apple with ice cubes, colorful hyperrealism, digital explosion of vibrant colors and abstract digital elements"
DEFAULT_WIDTH, DEFAULT_HEIGHT = 512, 512
OUTPUT_IMAGE_KEY = "output_img"
LOADED_IMAGE_KEY = "loaded_image"



def get_image(key: str) -> Optional[Image.Image]:
    if key in st.session_state:
        return st.session_state[key]
    return None


def set_image(key: str, img: Image.Image):
    st.session_state[key] = img


@st.cache_resource(max_entries=1)
def get_pipeline(
    name: PIPELINE_NAMES,
) -> Union[
    StableDiffusionPipeline,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipeline,
]:
    if name in ["txt2img", "img2img"]:
        model_id = "FFusion/FFusion-BaSE"
        
        pipeline = DiffusionPipeline.from_pretrained(model_id)
        # switch the scheduler in the pipeline to use the DDIMScheduler
        pipeline.scheduler = DDIMScheduler.from_config(
            pipeline.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing"
        )
        pipeline = pipeline.to("cuda")
        return pipeline


def generate(
    prompt,
    pipeline_name: PIPELINE_NAMES,
    num_images=1,
    negative_prompt=None,
    steps=22,
    width=896,
    height=1024,
    guidance_scale=6,
    enable_attention_slicing=True,
    enable_xformers=True
):
    """Generates an image based on the given prompt and pipeline name"""
    negative_prompt = negative_prompt if negative_prompt else None
    p = st.progress(0)
    callback = lambda step, *_: p.progress(step / steps)

    pipe = get_pipeline(pipeline_name)
    torch.cuda.empty_cache()

    if enable_attention_slicing:
        pipe.enable_attention_slicing()
    else:
        pipe.disable_attention_slicing()

    if enable_xformers:
        pipe.enable_xformers_memory_efficient_attention()
        
    kwargs = dict(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=steps,
        callback=callback,
        guidance_scale=guidance_scale,
        guidance_rescale=0.7
    )
    print("kwargs", kwargs)

    if pipeline_name == "txt2img":
        kwargs.update(width=width, height=height)

    elif pipeline_name in ["inpaint", "img2img"]:
        kwargs.update(image_input=image_input)

    else:
        raise Exception(
            f"Cannot generate image for pipeline {pipeline_name} and {prompt}"
        )

    output_images = []  # list to hold output image objects
    for _ in range(num_images):  # loop over number of images
        result = pipe(**kwargs)  # generate one image at a time
        images = result.images
        for i, image in enumerate(images):  # loop over each image
            filename = (
                "data/"
                + sanitize_filename(re.sub(r"\s+", "_", prompt)[:50])
                + f"_{i}_{datetime.now().timestamp()}"
            )
            image.save(f"{filename}.png")
            output_images.append(image)  # add the image object to the list

    for image in output_images:
        with open(f"{filename}.txt", "w") as f:
            f.write(prompt)

    return output_images  # return the list of image objects



def prompt_and_generate_button(prefix, pipeline_name: PIPELINE_NAMES, **kwargs):
    prompt = st.text_area(
        "Prompt",
        value=DEFAULT_PROMPT,
        key=f"{prefix}-prompt",
    )
    negative_prompt = st.text_area(
        "Negative prompt",
        value="(disfigured), bad quality, ((bad art)), ((deformed)), ((extra limbs)), (((duplicate))), ((morbid)), (((ugly)), blurry, ((bad anatomy)), (((bad proportions))), cloned face, body out of frame, out of frame, bad anatomy, gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), (fused fingers), (too many fingers), (((long neck))), Deformed, blurry",
        key=f"{prefix}-negative-prompt",
    )
    col1, col2 = st.columns(2)
    with col1:
        steps = st.slider("Number of inference steps", min_value=11, max_value=69, value=14, key=f"{prefix}-inference-steps")
    with col2:
        guidance_scale = st.slider(
            "Guidance scale", min_value=0.0, max_value=20.0, value=7.5, step=0.5, key=f"{prefix}-guidance-scale"
        )
   # enable_attention_slicing = st.checkbox('Enable attention slicing (enables higher resolutions but is slower)', key=f"{prefix}-attention-slicing", value=True)
   # enable_xformers = st.checkbox('Enable xformers library (better memory usage)', key=f"{prefix}-xformers", value=True)
        num_images = st.slider("Number of images to generate", min_value=1, max_value=4, value=1, key=f"{prefix}-num-images")

    images = []

    
    if st.button("Generate images", key=f"{prefix}-btn"):
        with st.spinner("Generating image..."):
            images = generate(
                prompt,
                pipeline_name,
                num_images=num_images,  # add this
                negative_prompt=negative_prompt,
                steps=steps,
                guidance_scale=guidance_scale,
                enable_attention_slicing=True,  # value always set to True
                enable_xformers=True,  # value always set to True
                **kwargs,
            )

    for i, image in enumerate(images):  # loop over each image
        set_image(f"{OUTPUT_IMAGE_KEY}_{i}", image.copy())  # save each image with a unique key


    image_indices = [int(key.split('_')[-1]) for key in st.session_state.keys() if OUTPUT_IMAGE_KEY in key]
    cols = st.columns(len(image_indices) if image_indices else 1)  # create a column for each image or a single one if no images
    for i in range(max(image_indices) + 1 if image_indices else 1):  # loop over each image index
        output_image_key = f"{OUTPUT_IMAGE_KEY}_{i}"
        output_image = get_image(output_image_key)
        if output_image:
            cols[i].image(output_image)



def width_and_height_sliders(prefix):
    col1, col2 = st.columns(2)
    with col1:
        width = st.slider(
            "Width",
            min_value=768,
            max_value=1024,
            step=128,
            value=768,
            key=f"{prefix}-width",
        )
    with col2:
        height = st.slider(
            "Height",
            min_value=768,
            max_value=1024,
            step=128,
            value=768,
            key=f"{prefix}-height",
        )
    return width, height
    
data_dir = "/data"  # Update with the correct path

# Get all file names in the data directory
file_names = os.listdir(data_dir)


def txt2img_tab():
    prefix = "txt2img"
    width, height = width_and_height_sliders(prefix)
    prompt_and_generate_button(prefix, "txt2img", width=width, height=height)


def inpainting_tab():
    col1, col2 = st.columns(2)

    with col1:
        image_input, mask_input = inpainting()

    with col2:
        if image_input and mask_input:
            prompt_and_generate_button(
                "inpaint", "inpaint", image_input=image_input, mask_input=mask_input
            )


def img2img_tab():
    col1, col2 = st.columns(2)

    with col1:
        image = image_uploader("img2img")
        if image:
            st.image(image)

    with col2:
        if image:
            prompt_and_generate_button("img2img", "img2img", image_input=image)

def main():
    st.title("FFusion AI -beta- Playground")

    tabs = ["FFusion BaSE 768+ (txt2img)"]
    selected_tab = st.selectbox("Choose a di.FFusion.ai model", tabs)

    if selected_tab == "FFusion BaSE 768+ (txt2img)":
        txt2img_tab()

    st.header("Citation")

    """
    ```
        @misc {ffusion_ai_2023,
    	author       = { {FFusion AI} },
    	title        = { FFusion-BaSE (Revision ba72848) },
    	year         = 2023,
    	url          = { https://huggingface.co/FFusion/FFusion-BaSE },
    	doi          = { 10.57967/hf/0851 },
    	publisher    = { Hugging Face }
        } http://doi.org/10.57967/hf/0851
    ```
    """

    """
        Please note that the demo is intended for academic and research purposes ONLY. Any use of the demo for generating inappropriate content is strictly prohibited. The responsibility for any misuse or inappropriate use of the demo lies solely with the users who generated such content, and this demo shall not be held liable for any such use. By interacting within this environment, you hereby acknowledge and agree to the terms of the CreativeML Open RAIL-M License.
    """
         
if __name__ == "__main__":
    main()