import os
import gradio as gr
import toml
from math import ceil
from pnginfo import read_info_from_image
from huggingface_hub import HfApi

client_config = toml.load("config.toml")['client']
num_of_imgs_per_page = 60
api = HfApi()
    
def get_img_list(folder, page):
    img_list = [os.path.join(client_config['save_path'], folder, f) for f in os.listdir(os.path.join(client_config['save_path'], folder))]
    img_list.sort(key=lambda f: os.path.getctime(f), reverse=True)
    if page == 0:
        page = 1
    elif page*num_of_imgs_per_page > len(img_list) or page < 0:
        page = ceil(len(img_list)/num_of_imgs_per_page)
    res = img_list[(page - 1)*num_of_imgs_per_page: page*num_of_imgs_per_page]
    return gr.Gallery(value=res, selected_index=None), page

def update_img_dir():
    all_dirs = os.listdir(client_config['save_path'])
    all_dirs.sort(reverse=True)
    img_dir_path = gr.update(choices=all_dirs)
    return img_dir_path

def get_img_info(evt: gr.SelectData):
    info, items = read_info_from_image(evt.value['image']['path'])
    index = evt.index
    return index, info, items

def del_img(index, folder, page):
    if index < 0:
        return gr.Gallery(selected_index=None)
    img_list = [os.path.join(client_config['save_path'], folder, f) for f in os.listdir(os.path.join(client_config['save_path'], folder))]
    img_list.sort(key=lambda f: os.path.getctime(f), reverse=True)
    if page == 0:
        page = 1
    elif page*num_of_imgs_per_page > len(img_list) or page < 0:
        page = ceil(len(img_list)/num_of_imgs_per_page)
    res = img_list[(page - 1)*num_of_imgs_per_page: page*num_of_imgs_per_page]
    os.remove(res[index])
    api.delete_file(path_in_repo=res[index], repo_id="P01yH3dr0n/naimages", repo_type="dataset", token=os.environ.get("hf_token"))
    res.pop(index)
    return gr.Gallery(value=res, selected_index=None)

def img_history_ui(tab):
    os.makedirs(client_config['save_path'], exist_ok=True)
    all_dirs = os.listdir(client_config['save_path'])
    all_dirs.sort(reverse=True)
    with gr.Row(): 
        img_dir_path = gr.Dropdown(choices=all_dirs, value=None, label='图片文件夹')
    
    tab.select(fn=update_img_dir, inputs=None, outputs=img_dir_path)
    with gr.Row():
        with gr.Column(scale=2):
            with gr.Row():
                first_page = gr.Button('第一页')
                prev_page = gr.Button('前一页')
                page_index = gr.Number(value=1, label="页码", precision=0)
                next_page = gr.Button('下一页')
                end_page = gr.Button('末尾页')
            history_gallery = gr.Gallery(show_label=False, format='png', columns=6, height=1024, interactive=False, show_share_button=False)
        with gr.Column(scale=1):
            info = gr.HTML()
            items = gr.JSON(value=None, visible=False)
            selected_index = gr.Number(value=-1, visible=False, precision=0)
            with gr.Row():
                send_btn = gr.Button('参数发送到文生图', visible=False)
                history2ref = gr.Button("图片发送到风格迁移", visible=False)
                history2i2i = gr.Button("图片发送到图生图", visible=False)
                history2inp = gr.Button("图片发送到局部重绘", visible=False)
                history2dtl = gr.Button("图片发送到引导工具", visible=False)
                del_btn = gr.Button('删除', visible=False)
    
    img_dir_path.change(lambda path: get_img_list(path, 1) + ('', None), inputs=img_dir_path, outputs=[history_gallery, page_index, info, items])
    img_dir_path.blur(update_img_dir, inputs=None, outputs=img_dir_path)
    first_page.click(lambda path: get_img_list(path, 1) + ('', None), inputs=img_dir_path, outputs=[history_gallery, page_index, info, items])
    prev_page.click(lambda path, p: get_img_list(path, p - 1) + ('', None), inputs=[img_dir_path, page_index], outputs=[history_gallery, page_index, info, items])
    page_index.submit(lambda path, p: get_img_list(path, p) + ('', None), inputs=[img_dir_path, page_index], outputs=[history_gallery, page_index, info, items])
    next_page.click(lambda path, p: get_img_list(path, p + 1) + ('', None), inputs=[img_dir_path, page_index], outputs=[history_gallery, page_index, info, items])
    end_page.click(lambda path: get_img_list(path, -1) + ('', None), inputs=img_dir_path, outputs=[history_gallery, page_index, info, items])
    history_gallery.select(get_img_info, inputs=None, outputs=[selected_index, info, items])
    history_gallery.change(lambda: -1, outputs=selected_index)
    del_btn.click(del_img, inputs=[selected_index, img_dir_path, page_index], outputs=history_gallery)
    items.change(lambda s: [gr.Button(visible=False) for _ in range(6)] if s is None else [gr.Button(visible=True) for _ in range(6)], inputs=items, outputs=[send_btn, history2ref, history2i2i, history2inp, history2dtl, del_btn])
    return history_gallery, selected_index, send_btn, items, history2ref, history2i2i, history2inp, history2dtl