import gradio as gr
import numpy as np
from gradio_scripts.pdf_handler import make_pdf
import os

js_update_tab_labels = """
    async () => {
        let el_list = document.getElementById("tab_labeler").getElementsByClassName("svelte-1kcgrqr")
        let idx = (el_list[1].value === "LOADING") ? 1 : parseInt(el_list[1].value)
        console.log(idx)
        style_sheet = document.getElementById("tab_style")
        style_sheet.innerHTML = ""
        for (let i = 1; i <= idx; i++) {
            style_sheet.innerHTML += "#result_tabs button.svelte-kqij2n:nth-child(" + i + "):before {content: 'Result " + i + "';}"
        }
    }
"""

table_headers = ["TOTAL", "START_FRAME", "END_FRAME", "DETECTION_DROPOUT", "DIR", "R", "THETA", "L", "TRAVEL"]
info_headers = [
    "TOTAL_TIME", "DATE", "START", "END", "FRAME_RATE", "",
    "TOTAL_FISH", "UPSTREAM_FISH", "DOWNSTREAM_FISH", "NONDIRECTIONAL_FISH", "",
    "UPSTREAM_MOTION", "INTENSITY", "THRESHOLD", "WATER_TEMP", "",
]
max_tabs = 10

tabs = []
tab_parent = None
zip_out = None

def update_result(i, state, result, inference_handler):
    print("loading result")
    # If index is larger than max_tabs, only add file to zip list
    if i >= max_tabs:
        return {
            zip_out: gr.update(value=result["path_zip"])
        }

    # Check if inference is done
    not_done = state['index'] < state['total']

    annotation_avaliable = not (result["aris_input"][i] == None)

    make_pdf(state['index']-1, state, result, table_headers)

    # Check if files exist
    video_path = result["path_video"][i]
    if not os.path.exists(video_path): video_path = None

    # Send update to UI, and to inference_handler to start next file inference

    print("loaded result")
    return {
        zip_out: gr.update(value=result["path_zip"]),
        tabs[i]['tab']: gr.update(),
        tabs[i]['video']: gr.update(value=video_path, visible=video_path is not None),
        tabs[i]['metadata']: gr.update(value=result["fish_info"][i], visible=True),
        tabs[i]['table']: gr.update(value=result["fish_table"][i], visible=True),
        tabs[i]['annotation_btn']: gr.update(visible=annotation_avaliable),
        tab_parent: gr.update(selected=i),
        inference_handler: gr.update(value = str(np.random.rand()), visible=not_done)
    }




def Result_Gradio(prepare_annotation, components):
    global tabs, tab_parent, zip_out

    # Dummy element to call inference events, this also displays the inference progress
    components['inference_handler'] = gr.Text(value=str(np.random.rand()), visible=False)

    # Dummy element to call UI events
    components['result_handler'] = gr.Text(value="LOADING", visible=False)

    # Dummy element for updating tab titles
    components['tab_labeler'] = gr.Text(value="", visible=False, elem_id="tab_labeler")
    components['tab_labeler'].change(lambda x: x, None, None, _js=js_update_tab_labels)

    with gr.Row():
        components['cancelBtn'] = gr.Button("Cancel Inference", visible=False)
        components['skipBtn'] = gr.Button("Skip this file", visible=False)

    # List of all UI components that will recieve outputs from the result_handler
    visual_components = []

    # Zip file output
    zip_out = gr.File(label="ZIP Output", interactive=False)
    visual_components.append(zip_out)


    # Create result tabs
    tabs = []
    with gr.Tabs(elem_id="result_tabs") as tab_parent:
        visual_components.append(tab_parent)

        # Create 'max_tab' tabs for showing result
        for i in range(max_tabs):
            with gr.Tab(label="", id=i, elem_id="result_tab"+str(i)) as tab:

                with gr.Row():
                    # List of clip info (date, time, number of fish, temperature, etc.) 
                    metadata_out = gr.Matrix(label="Info", interactive=False, headers=["Field", "Value"], datatype="markdown", visible=False, elem_id="marking_json")
                    
                    # Annotated video
                    video_out = gr.Video(label='Annotated Video', interactive=False, visible=False)


                # Table of found fish
                table_out = gr.Matrix(label='Indentified Fish', headers=table_headers, interactive=False, visible=False)

                # Button for opening result in annotation editor
                annotation_btn = gr.Button("Edit Annotation", visible=False)
                annotation_btn.click(prepare_annotation, annotation_btn, [components['annotation_progress'], components['master_tabs']], _js="() => " + str(i))
                
                # Add components to tab dict for easy access later on
                tabs.append({
                    'tab': tab,
                    'metadata': metadata_out,
                    'video': video_out,
                    'table': table_out,
                    'annotation_btn': annotation_btn
                })

                # Add all components to list of visualization outputs
                visual_components.extend([tab, metadata_out, video_out, table_out, annotation_btn])

    components['result_tabs'] = tab_parent

    return visual_components