from generate_person import init_display_person_dataframe, schedule_to_df, update_schedules, visualize_person_info
from generate_scene import person_info_to_description, reset_person_page, visualize_dynamic_generate, person_info_to_elements
from io import BytesIO
import pickle
from utils import read_character
import gradio as gr
from pathlib import Path

def add_character_click(character_file: str, person_name_to_file: dict[str, list[str|Path|dict]], person_name_to_info: dict[str, dict], person_name_to_description: dict[str, str], *reset_inputs):
    # we can test this by add person and check if it is one the person choice dropdown
    character_dict = pickle.load(open(character_file, 'rb'))
    _, person_info, _ = read_character(character_dict)
    person_name = person_info['name']
    person_name_to_file[person_name] = character_dict
    person_name_to_info[person_name] = person_info
    person_name_to_description[person_name] = person_info_to_description(person_info)
    return [
        person_name_to_file, person_name_to_info, person_name_to_description, gr.update(choices=sorted(person_name_to_description.values())), gr.update(visible=False)
    ] + reset_person_page(*reset_inputs)


def display_person_file(person_description_dropdown: str, person_name_to_file: dict[str, list[str|Path|dict]], person_name_to_info: dict[str, dict], person_name_to_description: dict[str, str]):
    # get person name
    person_name = [person_name for person_name, description in person_name_to_description.items() if description == person_description_dropdown][0]
    # get person file
    person_file = person_name_to_file[person_name]
    _, _, schedules = read_character(person_file)
    schedule_updates = update_schedules(schedules)
    
    person_info = person_name_to_info[person_name]
    return list(person_info_to_elements(person_info)) + schedule_updates
    
def create_person_page():
    with gr.Row(visible=False) as row:
        image = gr.Image(width=200, scale=0.25, show_label=False, interactive=False)
        with gr.Column():
            name = gr.Markdown()
            age = gr.Markdown()
            gender = gr.Markdown()
            personality = gr.Markdown()
            routine = gr.Markdown()
    return row, (image, name, age, gender, personality, routine)

def person_info_page(person_name_to_file: gr.State, person_name_to_info: gr.State, person_name_to_description: gr.State):
    with gr.Blocks():
        upload_button = gr.UploadButton()
        person_description_dropdown = gr.Dropdown(label="Person Description", choices=sorted(person_name_to_description.value.values()))
        person_row, person_elements = create_person_page()
        schedule_df_list, schedule_tab_list = init_display_person_dataframe()
    person_description_dropdown.change(
        fn=display_person_file,
        inputs=[person_description_dropdown, person_name_to_file, person_name_to_info, person_name_to_description],
        outputs=list(person_elements) + schedule_df_list + schedule_tab_list,
    )
    person_description_dropdown.change(
        fn=lambda : gr.update(visible=True),
        outputs=[person_row],
    )
    return person_description_dropdown, upload_button


def visualize():
    with gr.Blocks() as demo:
        with gr.Tab("Create Person"):
            add_character_button, character_file, person_name_to_file, person_name_to_info, person_name_to_description = visualize_person_info()
        with gr.Tab("Person Info"):
            person_description_dropdown, upload_button = person_info_page(person_name_to_file, person_name_to_info, person_name_to_description)
        with gr.Tab("Generate Scene"):
            reset_inputs, reset_outputs = visualize_dynamic_generate(person_name_to_file, person_name_to_info, person_name_to_description)
        add_character_button.click(
            fn=add_character_click,
            inputs=[character_file, person_name_to_file, person_name_to_info, person_name_to_description] + reset_inputs,
            outputs=[person_name_to_file, person_name_to_info, person_name_to_description, person_description_dropdown, add_character_button] + reset_outputs
        )
        upload_button.upload(
            fn=add_character_click,
            inputs=[upload_button, person_name_to_file, person_name_to_info, person_name_to_description] + reset_inputs,
            outputs=[person_name_to_file, person_name_to_info, person_name_to_description, person_description_dropdown, add_character_button] + reset_outputs
        )
    demo.launch()

if __name__ == "__main__":
    # with gr.Blocks() as demo:
    #     with gr.Tab("Create Person"):
    #         visualize_person_info()
    #     with gr.Tab("Generate Scene"):
    #         visualize_dynamic_generate()
    visualize()

# your_dictionary = {
#     'key1': 'value1',
#     'key2': 'value2',
#     'key3': 'value3'
# }

# # 将字典序列化为 pickle 格式的 bytes 对象
# def serialize_dictionary_to_bytes_obj(dictionary):
#     # 使用 pickle 将字典序列化为 bytes 对象
#     pickle_bytes = pickle.dumps(dictionary)
    
#     # 返回 bytes 对象和文件名
#     return pickle_bytes


# def process_and_return_file_path(dictionary):
#     # 创建一个临时文件，该文件在关闭时会被自动删除
#     with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
#         # 向临时文件写入数据
#         tmp_file.write(serialize_dictionary_to_bytes_obj(dictionary))
#         # 获取临时文件的路径
#         temp_file_path = tmp_file.name
        
#     # 返回文件路径
#     # 注意：这里我们不关闭文件，因为我们需要文件路径有效
#     # 但这意味着文件在操作系统级别不会自动删除，需要手动管理
#     return temp_file_path

# with gr.Blocks() as demo:
#     file_path = process_and_return_file_path(your_dictionary)
#     gr.File(file_path, label="Download File")
# demo.launch()