from os import rename

from uuid import uuid4

import gradio as gr
from zipfile import ZipFile

output_kind = [
    'Both',
    'Bordered',
    'Borderless'
]

def rename_file(file_path: str) -> str:
    new_file_path = file_path.replace(".png", "") + str(uuid4()) + ".png"
    try: 
        rename(file_path, new_file_path)
        return new_file_path
    except Exception as e:
        print(e)
        return file_path


def extract_snapshots(watchface_link: str, output_type='Both') -> list[str]:
    assert(watchface_link != None)

    with ZipFile(watchface_link, 'r') as zObject:
        if output_type == output_kind[0]:
            snapshot = zObject.extract('snapshot.png', path='outputs')
            bordless_snapshot = zObject.extract('no_borders_snapshot.png', path='outputs')
            return [snapshot, bordless_snapshot]
        elif output_type == output_kind[1]:
            return [zObject.extract('snapshot.png', path='outputs')]
        else:
            return [zObject.extract('no_borders_snapshot.png', path='outputs')]
    

def extract_snapshots_from(watchface_links: list[str], output_kind='Both') -> list[str]:
    assert(watchface_links != None and len(watchface_links) >= 1)

    filtered_watchface_links = [link for link in watchface_links if link.endswith('watchface')]
    snapshots = []
    for watchface_link in filtered_watchface_links:
        images = extract_snapshots(watchface_link, output_kind)
        images = list(map(rename_file, images))
        snapshots += images
    
    return snapshots


with gr.Blocks() as interface:
    gr.Markdown("""
# Welcome to Watchface Extractor
## You can easily extract Wathface previews!
    """)

    with gr.Tab('Single Watchface Output'):
        with gr.Row():
            with gr.Column():
                input_file = gr.File(label="Watchface file", file_count='single', file_types=['watchface'])

                single_button = gr.Button("Extract", variant="primary")

            with gr.Column():
                output_1 = gr.Image(label='Bordered Output', type='filepath')
                output_2 = gr.Image(label='Borderless Output', type='filepath')
    with gr.Tab('Multiple Watchfaces Output'):
        with gr.Row():
            with gr.Column():
                input_files = gr.File(label="Watchface file", file_count='multiple', file_types=['watchface'])
                output_kind_choice = gr.Dropdown(choices=output_kind, value=output_kind[0], label="Output kind for Watchface")

                button = gr.Button("Extract", variant="primary")

            output_gallery = gr.Gallery(label='Output Images')

    single_button.click(
        extract_snapshots,
        inputs=[input_files],
        outputs=[output_1, output_2]
    )

    button.click(
        extract_snapshots_from,
        inputs=[input_files, output_kind_choice],
        outputs=[output_gallery]
    )
    
    input_file.upload(
        fn=extract_snapshots,
        inputs=[input_file],
        outputs=[output_1, output_2]
    )

    input_files.upload(
        fn=extract_snapshots_from,
        inputs=[input_files, output_kind_choice],
        outputs=[output_gallery]
    )

if __name__ == '__main__':
    interface.launch()