import gradio as gr 
from utilities.setup import get_files
#import spaces
from services.diarization import Diarizer
from services.asr import Transcriber

#@spaces.GPU
def process_meeting(audio_input, num_speakers, speaker_names):
    """
    audio_input: filepath --> str
    num_speakers: number --> int
    speaker_names: dataset --> np.array
    """
    
    print(audio_input) 
    # first, pass it through a diarization stage.
    diarization_result, label_file = diarizer.run(audio_input, num_speakers)
    # Next, pass it through a transctiption stage
    transcriber.run(audio_input)
    # After, match the diarization with the transcription

    # Finally, Clean up the docs.

    return diarization_result, label_file



def main(conf):

    with gr.Blocks(theme=gr.themes.Soft(text_size="lg")) as demo:
    
        with gr.TabItem(conf["layout"]["page_names"][0]):
            
            gr.Markdown("# 🎤 Non-Video Meeting Transcription and Speaker Diarization")
            gr.Markdown("![](file/microphone_pen_and_paper.png)")
            gr.Markdown(get_files.load_markdown_file(conf["layout"]["about"]))
        

        with gr.TabItem(conf["layout"]["page_names"][1]):
            gr.Markdown("# 🔊 Upload or record your meeting")
            audio_input = gr.Audio(type="filepath", label="Upload Audio File")
            num_speakers = gr.Dropdown(list(range(conf["session"]["min_speakers"], 
                                                  conf["session"]["max_speakers"])),
                                       label="Number of Speakers", 
                                       value=conf["session"]["min_speakers"])
            speaker_names = gr.Dataframe(
                label="Type your names and details. Your actual entries will be limited to the speakers you selected above.",
                headers=["Name", "Supporting Details"],
                datatype=["str", "str"],
                row_count=(5,"fixed"),
                col_count=(2, "fixed"),
                type="pandas"
            )

            process_button = gr.Button("Process")

            
        with gr.TabItem(conf["layout"]["page_names"][2]):
            gr.Markdown("# 🔊 View and download your meeting transcript")
            diarization_output = gr.Textbox(label="Diarization Output")
            label_file_link = gr.File(label="Download DAW Labels")


        # Process
        
        process_button.click(
            fn=process_meeting,
            inputs=[audio_input, num_speakers, speaker_names],
            outputs=[diarization_output, label_file_link]
        )

    demo.launch(server_name="0.0.0.0", allowed_paths=["/"])




if __name__ == "__main__":
    # get config
    conf = get_files.json_cfg()
    # initialize diarizer
    diarizer = Diarizer(conf)
    transcriber = Transcriber(conf)
    main(conf)