import gradio as gr
from StreamVAD import StreamVAD
from dataclasses import dataclass, field


vad = StreamVAD(
    'ax650',
    sensitivity=0.5,
    silence_ms=200
)

@dataclass
class AppState:
    history: list = field(default_factory=list)


def process_audio(audio, state):
    # print(audio)
    # audio is a tuple of (sample_rate, numpy int16 array)
    sr, audio_data = audio
    for result in vad.run(audio_data, sr):
        if result:
            state.history.append(
                gr.ChatMessage(role='user', content=gr.Audio(
                    label=f"{result['start_ts']} - {result['end_ts']}",
                    value=(sr, result['audio']),
                    waveform_options=gr.WaveformOptions(show_recording_waveform=False),
                    editable=False
                    )
                ),
            )

    return state.history


with gr.Blocks() as demo:
    state = gr.State(value=AppState())

    with gr.Row():
        chatbot = gr.Chatbot(type='messages')
    
    with gr.Row():
        input_audio = gr.Audio(sources=['microphone'], type='numpy', streaming=True)

    # streaming process
    input_audio.stream(fn=process_audio, inputs=[input_audio, state], outputs=[chatbot])

demo.launch(debug=True)