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


def generate_fase(audio_path, model_version, civil_channel, fase):
    if fase == 1:
        text = generate_fase_1(
            audio_path,
            model_version=model_version,
            civil_channel=civil_channel
        )
        return text, None, None, None, None, None

    elif fase == 2:
        text, sex, age, silence_event, shout_event, meteo_event = generate_fase_2(
            audio_path,
            model_version=model_version,
            civil_channel=civil_channel
        )
        return text, sex, age, silence_event, shout_event, meteo_event

    else:
        raise ValueError("Invalid fase. Must be 1 or 2.")

@spaces.GPU
def transcribe(inputs: str, model_version: str, civil_channel: str, fase: int):
    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(inputs, model_version, civil_channel, fase)


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"):

            gr.Markdown(
                "### 🎧 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."
            )

            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,
                        inputs=[input_1, model_version_1, civil_channel_1, gr.Number(value=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"):

            gr.Markdown(
                "### 🧠 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ó d'informació detectada."
            )

            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ó meteo", lines=2)

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

                    submit_btn2.click(
                        fn=transcribe,
                        inputs=[input_2, model_version_2, civil_channel_2, gr.Number(value=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()