import torch
import spaces

import gradio as gr
from transformers import pipeline

import os

model_ids = {
"Fast": "openai/whisper-small",
"Balanced": "openai/whisper-medium",
"Accurate": "openai/whisper-large-v3-turbo"
}
device = 0 if torch.cuda.is_available() else "cpu"

@spaces.GPU
def transcribe(model_name, inputs):
    if inputs is None:
        raise gr.Error("No audio file submitted!")
    
    pipe = pipeline(
        task="automatic-speech-recognition",
        model=model_ids[model_name],
        chunk_length_s=30,
        device=device,
    )
    text = pipe(inputs, batch_size=8, generate_kwargs={"task": "transcribe" , "language": "en"}, return_timestamps=True)["text"] 
    return  text

def read_template_files(folder_path):

    template_file_contents = []
    file_list = sorted(os.listdir(folder_path))

    for file_name in file_list:
        with open(folder_path+file_name, 'r', encoding='utf-8') as file:
            content = file.read()
            template_file_contents.append(content)
    return template_file_contents

templates_content = read_template_files("report_templates/")

templates = {
    "" : "",
    "Mauro Cervical Spine Levels -- 3MCS" : templates_content[0],
    "Mauro - CT Chest (With Contrastor Non-Contrast) -- MCSMCTC" : templates_content[1],
    "Houman - CT Cervical Nerve -- HECTPSI" : templates_content[2],
    "Houman - Result Consult -- HERESCON" : templates_content[3],
}

demo = gr.Blocks(theme=gr.themes.Ocean())

file_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Radio(list(model_ids.keys()), label="Step 2. Select your model⬇️", value="Fast"),
        gr.Audio(sources="upload", type="filepath",label="Step 3. Upload your audio file, and click the submit button⬇️")
    ],
    outputs="text",
    flagging_mode="never",
)

mf_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Radio(list(model_ids.keys()), label="Step 2. Select your model⬇️", value="Fast"),
        gr.Audio(sources="microphone", type="filepath",label="Step 3. Record your audio, and click the submit button⬇️")
    ],
    outputs="text",
    flagging_mode="never",
)
def show_template(tn):
    return templates[tn]
with demo:
    with gr.Row():
        gr.Markdown("<div style='text-align: center;'><h2>Automated transcription of voice comments</h2></div>")
    with gr.Row():
        with gr.Column():
            dropdown = gr.Dropdown(choices=list(templates.keys()), label="Step 1. Select your report template⬇️ (in this version, templates are only for guidance)")
            output = gr.TextArea()
            dropdown.change(fn=show_template, inputs=dropdown, outputs=output)
        with gr.Column():
            gr.TabbedInterface([file_transcribe, mf_transcribe], ["Audio File","Microphone"])

demo.queue().launch(ssr_mode=False)