import gradio as gr
import pandas as pd
import requests
import os
import numpy as np
import re
from tqdm import tqdm
from time import sleep
from PIL import Image
import requests
from io import BytesIO
from datasets import Dataset, load_dataset
import json
import cv2
import pathlib

import imagehash

API_TOKEN = os.environ.get("HF_READ_TOKEN")

MAX_MODEL_NUM = 300
'''
Yntec digplay
'''

hf_civitai_image_info_dataset = load_dataset("svjack/hf_civitai_image_info_v0", token = API_TOKEN)
hf_civitai_image_info_df = hf_civitai_image_info_dataset["train"].to_pandas()

def gen_interface(model_name, max_times = 3):
    times = 0
    gr_model_interface = None
    while gr_model_interface is None and times < max_times:
        try:
            gr_model_interface = gr.load("models/{}".format(model_name),live=True,preprocess = False)
        except:
            print("error {} times {}".format(model_name, times))
            sleep(2)
        times += 1
    return gr_model_interface

#gr_model_interface.title
def toImgOpenCV(imgPIL): # Conver imgPIL to imgOpenCV
    i = np.array(imgPIL) # After mapping from PIL to numpy : [R,G,B,A]
                         # numpy Image Channel system: [B,G,R,A]
    red = i[:,:,0].copy(); i[:,:,0] = i[:,:,2].copy(); i[:,:,2] = red;
    return i;

def toImgPIL(imgOpenCV): return Image.fromarray(cv2.cvtColor(imgOpenCV, cv2.COLOR_BGR2RGB));

def jpg_val_to_img(jpg_bytes):
    img_buf = np.frombuffer(jpg_bytes, np.uint8)
    img = cv2.imdecode(img_buf, cv2.IMREAD_UNCHANGED)
    return toImgPIL(img)


model_list = hf_civitai_image_info_df["hf_repo_id"].drop_duplicates().values.tolist()
model_interface_list = []
for model_name in tqdm(model_list):
    gr_model_interface = gen_interface(model_name)
    if gr_model_interface is not None:
        model_interface_list.append(gr_model_interface)
        if len(model_interface_list) >= MAX_MODEL_NUM:
            break
print("load model num : {}".format(len(model_interface_list)))

def get_civitai_iframe(url, width = 1400, height = 768, as_html = True, visible = False):
    html= '''
            <div style="justify-content: center; display: flex;">
                    <iframe
                    src="{}"
                    frameborder="0"
                    width="{}"
                    height="{}"
                    ></iframe>
            </div>
        '''.format(url, width, height)
    if as_html:
        html = gr.HTML(html, visible = visible)
    return html

def get_info_by_interface(gr_interface, model_interface_list = model_interface_list):
    #### out: (gr_interface, civitai_url, civitai_info)
    if hasattr(gr_interface, "app"):
        civitai_url = hf_civitai_image_info_df[
            hf_civitai_image_info_df["hf_repo_id"] == gr_interface.title
        ]["civital_url"].iloc[0]
        civitai_info = hf_civitai_image_info_df[
            hf_civitai_image_info_df["hf_repo_id"] == gr_interface.title
        ][["prompt", "image"]].values.tolist()
        return gr_interface ,civitai_url, civitai_info
    else:
        civitai_url = hf_civitai_image_info_df[
            hf_civitai_image_info_df["hf_repo_id"] == gr_interface
        ]["civital_url"].iloc[0]
        civitai_info = hf_civitai_image_info_df[
            hf_civitai_image_info_df["hf_repo_id"] == gr_interface
        ][["prompt", "image"]].values.tolist()
        return list(filter(lambda x:x.title == gr_interface, model_interface_list))[0] ,civitai_url, civitai_info

def read_image_from_url(url):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
    return img

def image_click(images, evt: gr.SelectData, gr_interface_value,
    ):
    img_selected = images[evt.index]
    #print(img_selected)
    im_data = img_selected["name"]
    im = Image.open(im_data)
    im_hash = imagehash.average_hash(
        im, hash_size = 1024
    )
    min_diff = int(1e10)
    #print(-1)
    repo_card_im_dict = dict(
    get_info_by_interface(gr_interface_value)[2]
    )
    min_repo_name = ""
    for idx ,(repo_name, repo_card_image) in enumerate(repo_card_im_dict.items()):
        repo_img = jpg_val_to_img(repo_card_image["bytes"])
        repo_img_hash = imagehash.average_hash(
            repo_img, hash_size = 1024
        )
        diff = im_hash - repo_img_hash
        if diff < min_diff:
            min_diff = diff
            min_repo_name = repo_name
        #print(idx)
    prompt = min_repo_name
    return prompt
    #return prompt, im

def try_repo_act_func(civitai_url, show_civitai_button):
    repo_html_iframe_hide = get_civitai_iframe(civitai_url, visible = True if show_civitai_button == "Show Civitai Page" else False)
    return repo_html_iframe_hide,  gr.Button("Hide Civitai Page" if show_civitai_button == "Show Civitai Page" else "Show Civitai Page")

with gr.Blocks(
    css = '''
    .header img {
          float: middle;
          width: 33px;
          height: 33px;
        }
        .header h1 {
          top: 18px;
          left: 10px;
        }
    '''
) as demo:
    gr.HTML(
        '''
        <center>
        <div class="header">
        <h1 class = "logo"> <img src="https://huggingface.co/spaces/svjack/Civitai-Stable-Diffusion-HF/resolve/main/civitai_logo.webp" alt="logo" />  🤗 Civitai Model on Huggingface </h1>
        </center>
        '''
    )

    with gr.Row():
        with gr.Column():
            with gr.Row():
                hf_model_dropdown = gr.Dropdown(label = "🤗 Hf model",
                    choices=sorted(map(lambda x: x.title, model_interface_list)),
                    value=sorted(map(lambda x: x.title, model_interface_list))[0],)
            with gr.Column():
                with gr.Row():
                    civitai_prompt = gr.Textbox(label = "🤭 Civitai Prompt (Click from 👉 right gallery to get them, and You can edit ✍️ yourself) Don't use them to generate NSFW content, such as porns. 😬👎 Or you will get ❌👀",
                        interactive = True,
                    )
                gen_button = gr.Button(label = "Generate")
                hf_image = gr.Image(label = "🤭 Image generate by 🤗 Huggingface", height = 768)

        with gr.Column():
            civitai_info_gallery = gr.Gallery(
                pd.Series(
                get_info_by_interface(hf_model_dropdown.value)[2]
                ).sample(n = min(len(get_info_by_interface(hf_model_dropdown.value)[2]), 30)).map(lambda t2: t2[1]).map(lambda x: x["bytes"]).map(jpg_val_to_img).values.tolist(),
                height = 1024,
                label = "🖱️👇 ➡️ 👈 Civitai image samples",
                object_fit = "contain"
            )

    with gr.Row():
        with gr.Column():
            try_repo_button = gr.Button("Show Civitai Page")
            civitai_iframe_html = get_civitai_iframe(
                get_info_by_interface(hf_model_dropdown.value)[1]
            )

    hf_model_dropdown.change(
        lambda x: pd.Series(
        get_info_by_interface(x)[2]
        ).sample(n = min(len(get_info_by_interface(x)[2]), 30)).map(lambda t2: t2[1]).map(lambda x: x["bytes"]).map(jpg_val_to_img).values.tolist(),
        hf_model_dropdown,
        civitai_info_gallery
    )
    hf_model_dropdown.change(
        lambda _: (gr.Button("Show Civitai Page"), gr.HTML(visible = False)),
        None,
        [try_repo_button, civitai_iframe_html]
    )

    civitai_info_gallery.select(
            image_click,
            [civitai_info_gallery, hf_model_dropdown],
            civitai_prompt
    )
    gen_button.click(lambda hf_model_name, text_prompt:
            get_info_by_interface(hf_model_name)[0](text_prompt),
            [hf_model_dropdown, civitai_prompt],
            hf_image
    )

    try_repo_button.click(
        lambda hf_model_name, button: try_repo_act_func(
            get_info_by_interface(hf_model_name)[1]
            , button),
        [hf_model_dropdown, try_repo_button],
        [civitai_iframe_html, try_repo_button]
    )

demo.launch(show_api = False)