app.py

from app.logger_config import logger as logging
import numpy as np
import gradio as gr
import asyncio
from fastrtc.webrtc import WebRTC
from pydub import AudioSegment
import time
import os 
from gradio.utils import get_space

from app.logger_config import logger as logging
from app.utils import (
    generate_coturn_config
)

EXAMPLE_FILES = ["data/bonjour.wav", "data/bonjour2.wav"]
DEFAULT_FILE = EXAMPLE_FILES[0] 
# Utilisé pour signaler l'arrêt du streaming à l'intérieur du générateur
stop_stream_state = gr.State(value=False)

def read_and_stream_audio(filepath_to_stream: str):
    """
    Un générateur synchrone qui lit un fichier audio (via filepath_to_stream)
    et le streame chunk par chunk d'1 seconde.
    """
    
    if not filepath_to_stream or not os.path.exists(filepath_to_stream):
        logging.error(f"Fichier audio non trouvé ou non spécifié : {filepath_to_stream}")
        # Tenter d'utiliser le fichier par défaut en cas de problème
        if os.path.exists(DEFAULT_FILE):
            logging.warning(f"Utilisation du fichier par défaut : {DEFAULT_FILE}")
            filepath_to_stream = DEFAULT_FILE
        else:
            logging.error("Fichier par défaut non trouvé. Arrêt du stream.")
            return

    logging.info(f"Préparation du segment audio depuis : {filepath_to_stream}")
    # Réinitialiser le signal d'arrêt à chaque lancement
    stop_stream_state.value = False
    
    try:
        segment = AudioSegment.from_file(filepath_to_stream)
        chunk_duree_ms = 1000
        logging.info(f"Début du streaming en chunks de {chunk_duree_ms}ms...")

        for i, chunk in enumerate(segment[::chunk_duree_ms]):
            iter_start_time = time.perf_counter()
            logging.info(f"Envoi du chunk {i+1}...")

            if stop_stream_state.value:
                logging.info("Signal d'arrêt reçu, arrêt de la boucle.")
                break

            output_chunk = (
                chunk.frame_rate,
                np.array(chunk.get_array_of_samples()).reshape(1, -1),
            )

            yield output_chunk

            iter_end_time = time.perf_counter()
            processing_duration_ms = (iter_end_time - iter_start_time) * 1000
            
            sleep_duration = (chunk_duree_ms / 1000.0) - (processing_duration_ms / 1000.0) - 0.1
            if sleep_duration < 0:
                sleep_duration = 0.01 # Éviter un temps de sommeil négatif
                
            logging.debug(f"Temps de traitement: {processing_duration_ms:.2f}ms, Sommeil: {sleep_duration:.2f}s")

            elapsed = 0.0
            interval = 0.05
            while elapsed < sleep_duration:
                if stop_stream_state.value:
                    logging.info("Signal d'arrêt reçu pendant l'attente.")
                    break
                wait_chunk = min(interval, sleep_duration - elapsed)
                time.sleep(wait_chunk)
                elapsed += wait_chunk
            if stop_stream_state.value:
                break

        logging.info("Streaming terminé.")

    except asyncio.CancelledError:
        logging.info("Stream arrêté par l'utilisateur (CancelledError).")
        raise
    except FileNotFoundError:
        logging.error(f"Erreur critique : Fichier non trouvé : {filepath_to_stream}")
    except Exception as e:
        logging.error(f"Erreur pendant le stream: {e}", exc_info=True)
        raise
    finally:
        stop_stream_state.value = False
        logging.info("Signal d'arrêt nettoyé.")


def stop_streaming():
    """Active le signal d'arrêt pour le générateur."""
    logging.info("Bouton Stop cliqué: envoi du signal d'arrêt.")
    stop_stream_state.value = True
    return None

# --- Interface Gradio ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "## Application 'Streamer' WebRTC (Serveur -> Client)\n"
        "Utilisez l'exemple fourni, uploadez un fichier ou enregistrez depuis votre micro, "
        "puis cliquez sur 'Start' pour écouter le stream."
    )

    # 1. État pour stocker le chemin du fichier à lire
    active_filepath = gr.State(value=DEFAULT_FILE)

    with gr.Row():
        with gr.Column():
            main_audio = gr.Audio(
                label="Source Audio",
                sources=["upload", "microphone"], # Combine les deux sources
                type="filepath",
                value=DEFAULT_FILE, # Défaut au premier exemple
            )
        with gr.Column():
            webrtc_stream = WebRTC(
                label="Stream Audio",
                mode="receive",
                modality="audio",
                rtc_configuration=generate_coturn_config(),
                visible=True, # Caché par défaut
                height = 200,
            )
    # 4. Boutons de contrôle
    with gr.Row():
        with gr.Column():
            start_button = gr.Button("Start Streaming", variant="primary")
            stop_button = gr.Button("Stop Streaming", variant="stop", interactive=False)
        with gr.Column():
            gr.Text()

    def set_new_file(filepath):
        """Met à jour l'état avec le nouveau chemin, ou revient au défaut si None."""
        if filepath is None:
            logging.info("Audio effacé, retour au fichier d'exemple par défaut.")
            new_path = DEFAULT_FILE
        else:
            logging.info(f"Nouvelle source audio sélectionnée : {filepath}")
            new_path = filepath
        # Retourne la valeur à mettre dans le gr.State
        return new_path

    # Mettre à jour le chemin si l'utilisateur upload, efface, ou change le fichier
    main_audio.change(
        fn=set_new_file, 
        inputs=[main_audio], 
        outputs=[active_filepath]
    )
    
    # Mettre à jour le chemin si l'utilisateur termine un enregistrement
    main_audio.stop_recording(
        fn=set_new_file, 
        inputs=[main_audio], 
        outputs=[active_filepath]
    )


    # Fonctions pour mettre à jour l'état de l'interface
    def start_streaming_ui():
        logging.info("UI : Démarrage du streaming. Désactivation des contrôles.")
        return {
            start_button: gr.Button(interactive=False),
            stop_button: gr.Button(interactive=True),
            main_audio: gr.Audio(visible=False), 
        }

    def stop_streaming_ui():
        logging.info("UI : Arrêt du streaming. Réactivation des contrôles.")
        return {
            start_button: gr.Button(interactive=True),
            stop_button: gr.Button(interactive=False),
            main_audio: gr.Audio(
                label="Source Audio",
                sources=["upload", "microphone"], # Combine les deux sources
                type="filepath",
                value=active_filepath.value,
                visible=True
        ), 
            }


    ui_components = [
        start_button, stop_button, 
        main_audio, 
    ]

    stream_event = webrtc_stream.stream(
        fn=read_and_stream_audio,
        inputs=[active_filepath],  
        outputs=[webrtc_stream],
        trigger=start_button.click,
        concurrency_id="audio_stream", # ID de concurrence
        concurrency_limit=10
    )

    # Mettre à jour l'interface au clic sur START
    start_button.click(
        fn=start_streaming_ui,
        outputs=ui_components
    )

    # Correction : S'assurer que le stream est bien annulé
    stop_button.click(
        fn=stop_streaming, 
        outputs=[webrtc_stream], 
    ).then(
        fn=stop_streaming_ui, # ENSUITE, mettre à jour l'interface
        inputs=None,
        outputs=ui_components
    )


if __name__ == "__main__":
    demo.queue(max_size=10, api_open=False).launch(show_api=False, debug=True)