import os
import gradio as gr
import spaces

from whisper_cs_fase_1 import generate_fase_1
from whisper_cs_fase_2 import generate_fase_2
from AinaTheme import theme

@spaces.GPU()
def transcribe_fase_1(inputs: str, model_version: str, civil_channel: str):
    if inputs is None:
        raise gr.Error("Cap fitxer d'àudio introduit! Si us plau pengeu un fitxer o enregistreu un àudio abans d'enviar la vostra sol·licitud")
    return generate_fase_1(audio_path=inputs, model_version=model_version, civil_channel=civil_channel)

@spaces.GPU()
def transcribe_fase_2_display(inputs: str, model_version: str, civil_channel: str):
    if inputs is None:
        raise gr.Error("Cap fitxer d'àudio introduit! Si us plau pengeu un fitxer o enregistreu un àudio abans d'enviar la vostra sol·licitud")
    return generate_fase_2(audio_path=inputs, model_version=model_version, civil_channel=civil_channel)


def clear_fase_1(model_version, civil_channel):
    return None, model_version, civil_channel

def clear_fase_2(model_version, civil_channel):
    return None, model_version, civil_channel, "", "", "", "", "", ""


with gr.Blocks(theme=theme) as demo:
    gr.Markdown("## 🗣️ Transcripció automàtica d'àudio — Mode amb dues fases")

    with gr.Tabs():
        with gr.Tab("Fase 1"):
            description_string = (
                "### 🎧 Transcripció de trucades multilingüe de bona qualitat per a transcripció fiable\n"
                "- **v2_fast**: Inclou separació de canals i inferència ràpida.\n"
                "- **v1.0**: Inclou inferència moderada sense separació de canals."
            )
            gr.Markdown(description_string)

            with gr.Row():
                with gr.Column(scale=1):
                    model_version_1 = gr.Dropdown(
                        label="Model Version",
                        choices=["v2_fast", "v1.0"],
                        value="v2_fast",
                        elem_id="fase1-model-version",
                    )
                    civil_channel_1 = gr.Dropdown(
                        label="Canal del Civil (persona que truca)",
                        choices=["Left", "Right"],
                        value="Left",
                    )
                    input_1 = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
                with gr.Column(scale=1):
                    output_1 = gr.Textbox(label="Output", lines=8)

            with gr.Row(variant="panel"):
                clear_btn = gr.Button("Clear")
                submit_btn = gr.Button("Submit", variant="primary")

            submit_btn.click(fn=transcribe_fase_1, inputs=[input_1, model_version_1, civil_channel_1], outputs=[output_1])
            clear_btn.click(fn=clear_fase_1, inputs=[model_version_1, civil_channel_1], outputs=[input_1, model_version_1, civil_channel_1], queue=False)

        with gr.Tab("Fase 2"):
            description_string = (
                "### 🧠 Transcripció de trucades multilingüe de bona qualitat per a anàlisi d'informe\n"
                "- **v2_fast_and_detection_v1**: Inclou inferència ràpida, separació de parlants i explotació de nova informació per processos analítics i informes avançats."
            )
            gr.Markdown(description_string)

            with gr.Row():
                with gr.Column(scale=1):
                    model_version_2 = gr.Dropdown(
                        label="Model Version",
                        choices=["v2_fast_and_detection_v1"],
                        value="v2_fast_and_detection_v1",
                        elem_id="fase2-model-version",
                    )
                    civil_channel_2 = gr.Dropdown(
                        label="Canal del Civil (persona que truca)",
                        choices=["Left", "Right"],
                        value="Left",
                    )
                    input_2 = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
                with gr.Column(scale=1):
                    output_text = gr.Textbox(label="Transcripció ASR", lines=8)
                    output_sex = gr.Textbox(label="Gènere", lines=1)
                    output_age = gr.Textbox(label="Edat", lines=1)
                    output_silence = gr.Textbox(label="Detecció de silenci", lines=2)
                    output_shout = gr.Textbox(label="Detecció de crits", lines=2)
                    output_meteo = gr.Textbox(label="Detecció d'esdeveniment meteorològic", lines=2)

            with gr.Row(variant="panel"):
                clear_btn2 = gr.Button("Clear")
                submit_btn2 = gr.Button("Submit", variant="primary")

            submit_btn2.click(
                fn=transcribe_fase_2_display,
                inputs=[input_2, model_version_2, civil_channel_2],
                outputs=[output_text, output_sex, output_age, output_silence, output_shout, output_meteo]
            )

            clear_btn2.click(fn=clear_fase_2, inputs=[model_version_2, civil_channel_2], outputs=[input_2, model_version_2, civil_channel_2, output_text, output_sex, output_age, output_silence, output_shout, output_meteo], queue=False)


if __name__ == "__main__":
    demo.launch()