import gradio as gr
from gradio import networking
from huggingface_hub import hf_hub_download
import numpy as np
import re
import copy
from init_nemo_models import (
    ASR_LOCAL_MODEL_EXTENSIONS,
    DIAR_LOCAL_MODEL_EXTENSIONS,
    MultitalkerTranscriptionConfig,
    StreamingSession,
    resolve_model_reference,
)
from gradient_html import SPEAKER_COLORS, array_to_gradient_html
import time
import argparse
import os
import glob
import asyncio
import html
import threading
import traceback

try:
    import spaces
except ImportError:
    class _SpacesNoop:
        def GPU(self, *args, **kwargs):
            def decorator(fn):
                return fn
            return decorator

    spaces = _SpacesNoop()

# Global variables for controlling the flow of the program
MAX_LINES_DISPLAYED = 10
MAX_NUM_OF_SPEAKERS = 8
CHECK_INTERVAL = 0.05  # 100ms
DOUBLE_LOADING_STOPPER_DELTA_THRESHOLD = 20
RUDE_LOOP_STOP: bool = False
FILE_PROCESSING_STOPPED = threading.Event()
FILE_PROCESSING_STOPPED.set()
MICROPHONE_PROCESSING_ENABLED: bool = False
MICROPHONE_PROCESSING_LOCK = threading.Lock()
GPU_DURATION_SECONDS = int(os.environ.get("SPACES_GPU_DURATION", "300"))
CONFIG_RENEWED_FLAG: bool = False
START_TIME_OFFSET_EXAMPLES: float = -1.6 # seconds
START_TIME_OFFSET_FILE_UPLOAD: float = -0.8 # seconds
DEFAULT_DIAR_MODEL = "nvidia/Nemotron-3-Diarization"
DEFAULT_ASR_MODEL = "nvidia/nemotron-3.5-asr-streaming-0.6b"
SPACE_CHECKPOINT_FILENAMES = {
    DEFAULT_DIAR_MODEL: "Nemotron-3-Diarization.nemo",
    DEFAULT_ASR_MODEL: "nemotron-3.5-asr-streaming-0.6b.nemo",
}
DEFAULT_TARGET_LANG = "auto"
LEFT_CONTEXT_SIZE = 56
RIGHT_CONTEXT_SIZE = 13
ANSI_ESCAPE_PATTERN = re.compile(r"\x1b\[[0-9;]*m")
TIMESTAMP_PATTERN = re.compile(
    r"\[(?P<start>(?:\d{2}:)?\d{2}:\d{2}\.\d{2})-"
    r"(?P<end>(?:\d{2}:)?\d{2}:\d{2}\.\d{2})\]\s*"
)
SPEAKER_TAG_PATTERN = re.compile(r"\[ Speaker-(\d+) \]")

def enable_share_tunnel_retries(max_attempts=10, retry_delay_sec=5.0):
    """Retry transient Gradio share-tunnel failures without reloading the models."""
    setup_tunnel = networking.setup_tunnel

    def setup_tunnel_with_retries(*args, **kwargs):
        for attempt in range(1, max_attempts + 1):
            try:
                return setup_tunnel(*args, **kwargs)
            except Exception as error:
                if attempt == max_attempts:
                    raise
                print(
                    f"[WARN] Gradio share tunnel attempt {attempt}/{max_attempts} failed: {error}. "
                    f"Retrying in {retry_delay_sec:g} seconds..."
                )
                time.sleep(retry_delay_sec)

    networking.setup_tunnel = setup_tunnel_with_retries

def parse_arguments():
    """Parse model references, defaulting to Hugging Face repositories."""
    parser = argparse.ArgumentParser(
        description="NeMo Multispeaker Cache Aware Streaming Demo",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--diar_model",
        type=str,
        default=DEFAULT_DIAR_MODEL,
        help="Local diarization checkpoint (.nemo/.ckpt) or Hugging Face repository ID"
    )
    parser.add_argument(
        "--asr_model", 
        type=str,
        default=DEFAULT_ASR_MODEL,
        help="Local ASR checkpoint (.nemo) or Hugging Face repository ID"
    )
    parser.add_argument(
        "--target_lang",
        type=str,
        default=DEFAULT_TARGET_LANG,
        help="Language-locale prompt for prompt-conditioned ASR models, e.g. auto, en-US, ko-KR",
    )
    return parser.parse_args()


def prefetch_space_checkpoint(model_reference):
    """Download large Hub checkpoints before a ZeroGPU request acquires a GPU."""
    checkpoint_filename = SPACE_CHECKPOINT_FILENAMES.get(model_reference)
    if not os.environ.get("SPACE_ID") or checkpoint_filename is None:
        return model_reference

    print(
        f"[MODEL PREFETCH] Downloading {model_reference}/{checkpoint_filename} during startup...",
        flush=True,
    )
    checkpoint_path = hf_hub_download(
        repo_id=model_reference,
        filename=checkpoint_filename,
    )
    print(f"[MODEL PREFETCH] Ready: {checkpoint_path}", flush=True)
    return checkpoint_path

# Parse command line arguments
args = parse_arguments()

# Resolve local files while preserving pretrained model identifiers.
args.diar_model = resolve_model_reference(
    args.diar_model,
    model_label="Diarization",
    local_extensions=DIAR_LOCAL_MODEL_EXTENSIONS,
)
args.asr_model = resolve_model_reference(
    args.asr_model,
    model_label="ASR",
    local_extensions=ASR_LOCAL_MODEL_EXTENSIONS,
)
args.diar_model = prefetch_space_checkpoint(args.diar_model)
args.asr_model = prefetch_space_checkpoint(args.asr_model)

ATT_CONTEXT_SIZE = [LEFT_CONTEXT_SIZE, RIGHT_CONTEXT_SIZE]

# Configuration for the streaming session
config_mic = MultitalkerTranscriptionConfig(
    diar_model=args.diar_model,
    asr_model=args.asr_model,
    sent_break_sec=1.0,
    displayed_transcript_lines=10,
    att_context_size=ATT_CONTEXT_SIZE,
    target_lang=args.target_lang,
    masked_asr=True,
    cache_gating=True,
    cache_gating_buffer_size=2,
    sample_rate=16000,
    max_num_of_spks=8,
    binary_diar_preds=True,
    spkcache_update_period=222,
    fifo_len=264,
    deploy_mode=True,
    microphone_streaming=True,
    audio_pre_loading=True,
)

config_file = MultitalkerTranscriptionConfig(
    diar_model=args.diar_model,
    asr_model=args.asr_model,
    sent_break_sec=1.0,
    displayed_transcript_lines=10,
    att_context_size=ATT_CONTEXT_SIZE,
    target_lang=args.target_lang,
    masked_asr=True,
    cache_gating=True,
    cache_gating_buffer_size=2,
    sample_rate=16000,
    max_num_of_spks=8,
    binary_diar_preds=True,
    spkcache_update_period=222,
    fifo_len=264,
    deploy_mode=True,
    audio_pre_loading=True,
)

# Initialize streaming session containers. Models are loaded lazily so ZeroGPU
# Spaces can start the app before CUDA is allocated to a request.
session_mic = StreamingSession()
session_file = StreamingSession()
MODEL_INIT_LOCK = threading.Lock()
MODELS_INITIALIZED = False
MIC_SESSION_INITIALIZED = False
FILE_SESSION_INITIALIZED = False
MAX_NUM_OF_SPEAKERS = min(config_mic.max_num_of_spks, config_file.max_num_of_spks)


def ensure_sessions_initialized(reload=False, load_mic=True, load_file=True):
    """Load or reload ASR/diarization sessions on demand."""
    global MODELS_INITIALIZED, MIC_SESSION_INITIALIZED, FILE_SESSION_INITIALIZED
    with MODEL_INIT_LOCK:
        if load_mic and (reload or not MIC_SESSION_INITIALIZED):
            session_mic.setup_streaming_session(config_mic, reload=MIC_SESSION_INITIALIZED)
            MIC_SESSION_INITIALIZED = True
        if load_file and (reload or not FILE_SESSION_INITIALIZED):
            session_file.setup_streaming_session(config_file, reload=FILE_SESSION_INITIALIZED)
            FILE_SESSION_INITIALIZED = True
        MODELS_INITIALIZED = MIC_SESSION_INITIALIZED or FILE_SESSION_INITIALIZED

def render_diarization_state(diar_state, config):
    """Render Sortformer's speaker cache, FIFO, and current chunk predictions."""
    num_speakers = int(config.max_num_of_spks)
    default_chunk_len = config.att_context_size[1] + 1
    default_spkcache_len = config.spkcache_len or config.spkcache_update_period or 1
    default_fifo_len = config.fifo_len or 1
    capacities = (
        diar_state.get("capacities", {})
        if diar_state is not None
        else {"spkcache": default_spkcache_len, "fifo": default_fifo_len, "chunk": default_chunk_len}
    )
    valid_frames = diar_state.get("valid_frames", {}) if diar_state is not None else {}

    def to_speaker_frame_array(name):
        values = None if diar_state is None else diar_state.get(name)
        if values is None:
            return np.zeros((num_speakers, 0), dtype=np.float32)
        return values[0].float().cpu().numpy().transpose()

    partitions = [
        {
            "name": "SPKCACHE · STORED",
            "values": to_speaker_frame_array("spkcache"),
            "capacity": capacities["spkcache"],
            "valid_frames": valid_frames.get("spkcache", 0),
        },
        {
            "name": "FIFO · FRESH",
            "values": to_speaker_frame_array("fifo"),
            "capacity": capacities["fifo"],
            "valid_frames": valid_frames.get("fifo", 0),
        },
        {
            "name": "CHUNK · FRESH",
            "values": to_speaker_frame_array("chunk"),
            "capacity": capacities["chunk"],
            "valid_frames": valid_frames.get("chunk", 0),
        },
    ]
    return array_to_gradient_html(num_of_speakers=num_speakers, partitions=partitions)


def format_transcript(text):
    """Formats the transcript with speaker colors for HTML display."""
    html_output = '<div style="color: #FFFFFF !important;">'
    if not text:
        return ""

    single_speaker_mode = False
    if MIC_SESSION_INITIALIZED and session_mic.multispk_asr_streamer is not None:
        single_speaker_mode = session_mic.multispk_asr_streamer._single_speaker_mode
    if FILE_SESSION_INITIALIZED and session_file.multispk_asr_streamer is not None:
        single_speaker_mode = (
            single_speaker_mode or session_file.multispk_asr_streamer._single_speaker_mode
        )
    current_color = "#76B900" if single_speaker_mode else "#A8A8A8"

    for raw_line in (line for line in text.strip().splitlines() if line.strip()):
        line = ANSI_ESCAPE_PATTERN.sub("", raw_line)

        timestamp_match = TIMESTAMP_PATTERN.search(line)
        if timestamp_match:
            timestamp = timestamp_match.group(0).strip()
            content = line[timestamp_match.end():]
        else:
            timestamp = ""
            content = line

        speaker_match = SPEAKER_TAG_PATTERN.search(content)
        if single_speaker_mode:
            content = SPEAKER_TAG_PATTERN.sub("", content).strip()
            current_color = "#76B900"
        elif speaker_match:
            speaker_index = int(speaker_match.group(1))
            current_color = SPEAKER_COLORS[speaker_index % len(SPEAKER_COLORS)]

        timestamp_html = (
            f'<span style="color: #E0E0E0;">{html.escape(timestamp)}</span> '
            if timestamp
            else ""
        )
        html_output += (
            f'{timestamp_html}<span style="color: {current_color}; font-weight: bold;">'
            f'{html.escape(content.strip())}</span><br>'
        )
    
    html_output += '</div>'
    return html_output

def update_config_for_file(evt: gr.SelectData, state=None): 
    """
    Update the configuration for the file based on the filename.
    Directly modifies the global config_file object and updates the status display.
    """
    print(f"[DEBUG] update_config_for_file() state[is_playing]: {state.get('is_playing', False)}")
    global RUDE_LOOP_STOP, CONFIG_RENEWED_FLAG
    RUDE_LOOP_STOP = True
    state["should_stop"] = False
    
    # Extract full filepath using the index from the SelectData event
    # evt.index gives us the position in the examples list
    filepath = examples[evt.index] if evt.index is not None else evt.value
    print(f"[DEBUG] update_config_for_file() full filepath: {filepath}")
    
    # Extract just the filename for config lookup
    filename = filepath.split('/')[-1] if isinstance(filepath, str) else filepath
    if isinstance(filename, list):
        filename = filename[0]
    print(f"[DEBUG] update_config_for_file() filename for config: {filename}")
    
    file_configs = {
        "4spk_example_16k.wav": {
            "audio_pre_loading": config_file.audio_pre_loading,
            "num_speakers": 4,
            "sent_break": 5.0,
            "binary": True,
            "chunk_size": 13,
            "single_spk": False,
            "cache_gate": True,
        },
        "fully_overlapped_speech.wav": {
            "audio_pre_loading": config_file.audio_pre_loading,
            "num_speakers": 2,
            "sent_break": 1.0,
            "binary": False,
            "chunk_size": 6,
            "single_spk": False,
            "cache_gate": True,
        },
        "pizza_debate.wav": {   
            "audio_pre_loading": config_file.audio_pre_loading,
            "num_speakers": 2,
            "sent_break": 1.0,
            "binary": False,
            "chunk_size": 13,
            "single_spk": False,
            "cache_gate": True,
        },
        "short_2spk_example.wav": {
            "audio_pre_loading": config_file.audio_pre_loading,
            "num_speakers": 2,
            "sent_break": 5.0,
            "binary": False,
            "chunk_size": 6,
            "single_spk": False,
            "cache_gate": True,
        },
        "tongue_twisters.wav": {
            "audio_pre_loading": config_file.audio_pre_loading,
            "num_speakers": 4,
            "sent_break": 8.0,
            "binary": False,
            "chunk_size": 13,
            "single_spk": False,
            "cache_gate": True,
        },
    }
    
    # Get config for this file using the filename, or use defaults
    config = file_configs.get(filename, {
        "audio_pre_loading": config_file.audio_pre_loading,
        "num_speakers": 8, 
        "sent_break": 1.0,
        "binary": True, 
        "chunk_size": 13, 
        "single_spk": False, 
        "cache_gate": True,
    })
    config["num_speakers"] = min(config["num_speakers"], MAX_NUM_OF_SPEAKERS)
    
    # Directly modify the global config_file object
    config_file.att_context_size = [LEFT_CONTEXT_SIZE, config["chunk_size"]]
    config_file.single_speaker_mode = config["single_spk"]
    config_file.max_num_of_spks = config["num_speakers"]
    config_file.cache_gating = config["cache_gate"]
    config_file.binary_diar_preds = config["binary"]
    config_file.sent_break_sec = config["sent_break"]
    CONFIG_RENEWED_FLAG = True

    # Generate status HTML
    status_html = get_text_box_html(
        audio_pre_loading=config["audio_pre_loading"],
        streaming_chunk_size=config["chunk_size"],
        single_speaker_mode="ON" if config["single_spk"] else "OFF",
        number_of_speakers=config["num_speakers"],
        cache_gating="ON" if config["cache_gate"] else "OFF",
        binary_diar="ON" if config["binary"] else "OFF",
        text_break_silence=config["sent_break"],
    )
    
    print(f"[DEBUG] update_config_for_file(): Configuration updated, state reset, returning filepath: {filepath}")
    # Return values to update all UI components and reset state
    return (
        state,                     # state (with should_stop reset)
        filepath,                  # filepath from the selected example
        config["chunk_size"],      # streaming_chunk_size_slider
        config["num_speakers"],    # num_speakers_slider
        config["sent_break"],      # sent_break_sec_slider
        config["audio_pre_loading"],      # audio_pre_loading_checkbox
        config["single_spk"],      # single_speaker_mode_checkbox
        config["cache_gate"],      # cache_gating_checkbox
        config["binary"],          # binary_diar_checkbox
        status_html,               # status_textbox
    )

def file_processing(filepath, state, file_upload=False, start_time_offset=0.0):
    """Handles file uploads and processes the audio in chunks with real-time simulation."""
    ensure_sessions_initialized(load_mic=False, load_file=True)
    print(f"[DEBUG] file_processing() FIRST LINE: filepath: {filepath}, state[is_playing]: {state.get('is_playing', False)}, file_upload: {file_upload}")
    if not filepath or state.get("should_stop", False):
        yield state, state["diar_html"], format_transcript(state.get("full_transcript", ""))
        return
    print(f"[DEBUG] fp1 Past the first if statement in file_processing()") 
    
    full_transcript = ""
    hop_duration_sec = (session_file.multispk_asr_streamer._att_context_size[1]+1) * 0.08
    audio_time = 0.0
    state["end_of_file"] = False
    start_time = time.time() + start_time_offset
    # Extra sleep to prevent transcription goes ahead of the audio file playing.
    time.sleep(0.4)
    try:    
        while not RUDE_LOOP_STOP:
            try:
                chunk_audio, chunk_lengths = next(session_file.streaming_buffer_iter)
            except StopIteration:
                break
            diar_state = session_file.process_audio_chunk(chunk_audio, chunk_lengths)
            full_transcript = session_file.transcribed_speaker_texts
            state["full_transcript"] = full_transcript
            html_diar_out = render_diarization_state(diar_state, config_file)
            state["diar_html"] = html_diar_out
            # Simulate real-time by waiting for hop duration
            # Adjust for any paused time to maintain real-time simulation
            if not state.get("is_playing", True):
                # Mark when pause started
                if state.get("pause_start_time") is None:
                    state["pause_start_time"] = time.time()
                
                while not state['is_playing'] and not state['end_of_file']:
                    time.sleep(CHECK_INTERVAL)
                    if RUDE_LOOP_STOP:
                        break
    
                # Calculate paused duration when resuming
                if state.get("pause_start_time") is not None:
                    paused_duration = time.time() - state["pause_start_time"]
                    state["total_paused_time"] = state.get("total_paused_time", 0.0) + paused_duration
                    state["pause_start_time"] = None
            audio_time += hop_duration_sec
            total_paused = state.get("total_paused_time", 0.0)
            adjusted_start = start_time + total_paused
            sleep_time = (adjusted_start + audio_time) - time.time()
            if sleep_time > 0:
                time.sleep(sleep_time)
            yield state, html_diar_out, format_transcript(full_transcript)
        if RUDE_LOOP_STOP:
            print(f"[DEBUG] file_processing(): RUDE_LOOP_STOP is True, ABORTING... filepath: {filepath}")
            state["should_stop"] = False
            state["last_processed_file"] = filepath
            yield state, state["diar_html_initial"], f"<p style='text-align: center; font-size:30px;'><b>Processing has been interrupted. <br> Please wait until the new audio file is loaded. </b></p>"
            return
    except GeneratorExit:
        # Allow generator to be closed properly
        raise
    except Exception:
        print("Streaming has been interrupted at file_processing()")
        traceback.print_exc()
        state["mode"] = "idle"
        if file_upload:
            state["should_stop"] = False
        else:
            state["should_stop"] = True
        yield state, state["diar_html"], format_transcript(full_transcript)
        return  # Exit the function immediately
    
    state["mode"] = "idle"
    state["should_stop"] = False
    yield state, html_diar_out, format_transcript(full_transcript)
    return 

def get_progress_bar_html(step, num_steps, sleep_per_step, size_mb, total_sleep_time): 
    """
    Create HTML string with progress bar.
    """
    # Create HTML progress bar
    progress_percent = int((step / num_steps) * 100)
    elapsed_time = step * sleep_per_step

    if step == num_steps:
        progress_bar_html = f"""
        <div style="width: 100%; background-color: #f0f0f0; border-radius: 10px; padding: 3px;">
            <div style="width: {progress_percent}%; background-color: #4CAF50; height: 30px; border-radius: 4px; 
                        text-align: center; line-height: 30px; color: white; font-weight: bold; font-size: 16px;">
                {progress_percent}%
            </div>
        </div>
        <p style="text-align: center; margin-top: 30px; font-size: 32px;">
            <b>Processing file ({size_mb:.1f} MB)... Done.</b><br>
            <b style="color: #76B900;">Click ▶ To Start Processing.</b>
        </p>
        """
    else:
        progress_bar_html = f"""
        <div style="width: 100%; background-color: #f0f0f0; border-radius: 10px; padding: 3px;">
            <div style="width: {progress_percent}%; background-color: #4CAF50; height: 30px; border-radius: 4px; 
                        text-align: center; line-height: 30px; color: white; font-weight: bold; font-size: 16px;">
                {progress_percent}%
            </div>
        </div>
        <p style="text-align: center; margin-top: 30px; font-size: 32px;">
            <b>Processing file ({size_mb:.1f} MB)... Do not start and please wait.</b><br>
            <span style="color: #666; font-size: 32px;">Time elapsed: {elapsed_time:.2f}s / {total_sleep_time:.2f}s</span>
        </p>
        """
    return progress_bar_html


async def show_progress_bar_with_sleep(filepath, num_steps=500, seconds_per_mb=0.5, power_factor=1.2):
    """
    Sleep for approximately 1 second per MB while showing a progress bar.
    This is needed to prevent the duplicated play signal from Gradio Audio component.
    If Gradio Audio gets updated with the fix for this, this function can be removed.
    
    Args:
        filepath: Path to the audio file
        num_steps: Number of progress updates to show
        seconds_per_mb: Number of seconds to sleep per MB
        
    Yields:
        HTML string with progress bar
    """
    size_bytes = os.path.getsize(filepath)
    size_mb = size_bytes / (1024 * 1024)
    print(f"[DEBUG] show_progress_bar_with_sleep(): size_mb: {size_mb}, seconds_per_mb: {seconds_per_mb}, bytes: {size_bytes}")
    total_sleep_time = (size_mb ** power_factor) * seconds_per_mb  
    sleep_per_step = total_sleep_time / num_steps
    
    for step in range(num_steps + 1):
        yield get_progress_bar_html(step=step, 
                                    num_steps=num_steps, 
                                    sleep_per_step=sleep_per_step, 
                                    size_mb=size_mb, 
                                    total_sleep_time=total_sleep_time
                                )
        
        if step < num_steps:
            await asyncio.sleep(sleep_per_step)

async def file_processing_example_click_async(filepath, state, start_time_offset=0.0):
    """Async wrapper for file_processing with pause/resume support (doesn't update audio component)."""
    ensure_sessions_initialized(load_mic=False, load_file=True)
    state["should_stop"] = False
    state["is_playing"] = False  
    state["pause_start_time"] = None
    state["total_paused_time"] = 0.0
    state["diar_html_initial"] = render_diarization_state(None, config_file)
    state["diar_html"] = copy.deepcopy(state["diar_html_initial"])
    state["diar_visualization_state"] = None
    state["full_transcript"] = ""
    state["stream_step_num"] = 0
    global RUDE_LOOP_STOP
    RUDE_LOOP_STOP = False

    ###########################################################
    # ### Wait for audio to start playing
    state["mode"], state["full_transcript"] = "file", ""
    state["last_processed_file"] = filepath
    print(f"[DEBUG] fpa1 Past the reset and setup in file_processing_example_click_async()") 
    
    yield state, state["diar_html_initial"], "<b>Loading the audio file...</b>"
    if RUDE_LOOP_STOP:
        return

    # Reset buffer and setup streaming session with current config
    session_file.reset_buffer()
    session_file.setup_streaming_session(config_file, reload=True)

    session_file.add_file_to_streaming_buffer(filepath)

    print(f"[DEBUG] fpa2 File added to streaming buffer, starting a loop processing... file_processing_example_click_async()")
    print(f"[DEBUG] fpa3 Final Waiting for PLAY >>>>>>> file_processing_upload_event_async(): is_playing: {state['is_playing']}")

    
    # Show progress bar while sleeping proportionally to file size
    if config_file.audio_pre_loading:
        async for progress_html in show_progress_bar_with_sleep(filepath):
            if RUDE_LOOP_STOP:
                return
            yield state, state["diar_html_initial"], progress_html
    else:
        async for progress_html in show_progress_bar_with_sleep(filepath, num_steps=2, seconds_per_mb=0.0, power_factor=1.0):
            if RUDE_LOOP_STOP:
                return
            yield state, state["diar_html_initial"], progress_html
    
    while not state['is_playing']:
        await asyncio.sleep(CHECK_INTERVAL)
        if RUDE_LOOP_STOP:
            return

    # print(f"[DEBUG] PLAY STARTED ><><><>< file_processing_upload_event_async(): is_playing: {state['is_playing']}")
    gen = file_processing(filepath, state, start_time_offset=start_time_offset)
    for s, html_diar, transcript in gen:
        # Wait while paused
        if not state.get("is_playing", True):
            # Mark when pause started
            if state.get("pause_start_time") is None:
                state["pause_start_time"] = time.time()
            
            while not state['is_playing']:
                await asyncio.sleep(CHECK_INTERVAL)
                if RUDE_LOOP_STOP:
                    break
 
            # Calculate paused duration when resuming
            if state.get("pause_start_time") is not None:
                paused_duration = time.time() - state["pause_start_time"]
                state["total_paused_time"] = state.get("total_paused_time", 0.0) + paused_duration
                state["pause_start_time"] = None
        
        if state.get("should_stop", False):
            return
        
        yield s, html_diar, transcript

def file_processing_example_click(filepath, state, start_time_offset=0.0):
    """Wrapper that runs the async function synchronously for Gradio compatibility."""
    # Create a new event loop for this generator

    # Initialize finished_file_list if not present
    if "finished_file_list" not in state:
        state["finished_file_list"] = {}
    
    if filepath in state["finished_file_list"]:
        delta_time = time.time() - state["finished_file_list"][filepath]["time_offset"]
        if delta_time < DOUBLE_LOADING_STOPPER_DELTA_THRESHOLD:
            print(f"[DEBUG] file_processing_example_click(): filepath {filepath} and time delta {delta_time} is less than {DOUBLE_LOADING_STOPPER_DELTA_THRESHOLD}, ABORTING...")
            yield state, state["diar_html"], format_transcript(state.get("full_transcript", ""))
            return
        else:
            print(f"[DEBUG] file_processing_example_click(): filepath {filepath} and time delta {delta_time} is greater than {DOUBLE_LOADING_STOPPER_DELTA_THRESHOLD}, continuing... ")
            # Add to list IMMEDIATELY to prevent duplicate processing from rapid .change() calls
            state["finished_file_list"][filepath] = {"time_offset": time.time()}
    else:
        print(f"[DEBUG] file_processing_example_click(): filepath {filepath} is not in the finished_file_list, adding to list... ")
        # Add to list IMMEDIATELY to prevent duplicate processing from rapid .change() calls
        state["finished_file_list"][filepath] = {"time_offset": time.time()}

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    FILE_PROCESSING_STOPPED.clear()
    
    try:
        # Convert async generator to sync generator
        async_gen = file_processing_example_click_async(filepath, state, start_time_offset=start_time_offset)
        
        while True:
            try:
                result = loop.run_until_complete(async_gen.__anext__())
                yield result
            except StopAsyncIteration:
                break
    finally:
        loop.close()
        FILE_PROCESSING_STOPPED.set()

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def file_processing_example_click_upload(filepath, state):
    """Wrapper for file upload with upload offset."""
    # file_processing_example_click is a generator, so yield from it
    for result in file_processing_example_click(filepath, state, start_time_offset=START_TIME_OFFSET_FILE_UPLOAD):
        yield result

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def file_processing_example_click_examples(filepath, state):
    """Wrapper for examples with examples offset."""
    # file_processing_example_click is a generator, so yield from it
    for result in file_processing_example_click(filepath, state, start_time_offset=START_TIME_OFFSET_EXAMPLES):
        yield result

async def file_processing_upload_event_async(filepath, state):
    """Async wrapper for file_processing that syncs with audio playback (play/pause)."""
    
    state["should_stop"] = False
    state["is_playing"] = False  # Start paused, wait for audio to play
    state["pause_start_time"] = None
    state["total_paused_time"] = 0.0
    global RUDE_LOOP_STOP
    RUDE_LOOP_STOP = False

    # Wait for audio to start playing
    print(f"[DEBUG] Waiting for PLAY >>>>>>> file_processing_upload_event_async(): is_playing: {state['is_playing']}")
    while not state['is_playing']:
        await asyncio.sleep(CHECK_INTERVAL)
        if RUDE_LOOP_STOP:
            break

    print(f"[DEBUG] PLAY STARTED ><><><>< file_processing_upload_event_async(): is_playing: {state['is_playing']}")
    gen = file_processing(filepath, state, file_upload=True, start_time_offset=START_TIME_OFFSET_FILE_UPLOAD)
    for s, html_diar, transcript in gen:
        # Wait while paused
        if not state['is_playing']:
            # Mark when pause started
            if state.get("pause_start_time") is None:
                state["pause_start_time"] = time.time()
            
            while not state['is_playing']:
                await asyncio.sleep(CHECK_INTERVAL)
                if RUDE_LOOP_STOP:
                    break
            
            # Calculate paused duration when resuming
            if state.get("pause_start_time") is not None:
                paused_duration = time.time() - state["pause_start_time"]
                state["total_paused_time"] = state.get("total_paused_time", 0.0) + paused_duration
                state["pause_start_time"] = None
        
        yield s, html_diar, transcript

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def file_processing_upload_event(filepath, state): 
    """Wrapper that runs the async function synchronously for Gradio compatibility."""
    print(f"[DEBUG] First line of file_processing_upload_event(): filepath: {filepath}, state[is_playing]: {state.get('is_playing', False)}")
    # Initialize finished_file_list if not present
    if "finished_file_list" not in state:
        state["finished_file_list"] = []
    
    if filepath in state["finished_file_list"]:
        print(f"[DEBUG] file_processing_upload_event(): filepath {filepath} is already in the finished_file_list, ABORTING...")
        yield state, state["diar_html"], format_transcript(state.get("full_transcript", ""))
        return
    else:
        print(f"[DEBUG] file_processing_upload_event(): filepath {filepath} is not in the finished_file_list, continuing... ")
        # Add to list IMMEDIATELY to prevent duplicate processing from rapid .change() calls
        state["finished_file_list"].append(filepath)


    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    
    try:
        # Convert async generator to sync generator
        async_gen = file_processing_upload_event_async(filepath, state)
        
        while True:
            try:
                result = loop.run_until_complete(async_gen.__anext__())
                yield result
            except StopAsyncIteration:
                break
    finally:
        loop.close()

def turn_state_to_resume(state):
    """Called when audio playback is resumed."""
    print(f"[DEBUG] ++++++++++ turn_state_to_resume(): is_playing set to True.")
    state["is_playing"] = True
    return state

def turn_state_to_pause(state):
    """Called when audio playback is paused."""
    if state["end_of_file"]:
        print(f"[DEBUG] ========== turn_state_to_pause(): end_of_file is {state['end_of_file']}, NOT Pausing...")
        state["is_playing"] = True
    else:
        print(f"[DEBUG] ========== turn_state_to_pause(): end_of_file is {state['end_of_file']}, Pausing...")
        state["is_playing"] = False
    return state

def example_selection_stop_file_processing(state, filepath):
    """Stops file processing when audio is cleared and prepares for new file."""
    state["mode"] = "idle"
    state["last_processed_file"] = None
    global RUDE_LOOP_STOP
    RUDE_LOOP_STOP = True
    print(f"[DEBUG] example_selection_stop_file_processing(): RUDE_LOOP_STOP set to {RUDE_LOOP_STOP}, filepath: {filepath}")
    time.sleep(1.7)  # Give time for loop to stop
    RUDE_LOOP_STOP = False  # Reset for next processing
    print(f"[DEBUG] example_selection_stop_file_processing(): RUDE_LOOP_STOP set to {RUDE_LOOP_STOP}, filepath: {filepath}")
    state["should_stop"] = False
    state["finished_file_list"] = {}
    return state, filepath

def file_upload_stop_file_processing(state):
    """Stops file processing when audio is cleared and prepares for new file."""
    state["mode"] = "idle"
    state["last_processed_file"] = None
    state["is_playing"] = False
    print(f"[DEBUG] file_upload_stop_file_processing(): RUDE_LOOP_STOP set to True")
    global RUDE_LOOP_STOP
    RUDE_LOOP_STOP = True
    ### if session_file is not None, reset the buffer and delete the session
    if FILE_SESSION_INITIALIZED:
        session_file.reset_buffer()
    time.sleep(1.7)  # Give time for loop to stop
    RUDE_LOOP_STOP = False  # Reset for next processing
    state["should_stop"] = False
    return state, None  # Return None for the audio component to clear it properly

def stop_file_processing_example(state):
    """Stops file processing when audio is cleared."""
    state["mode"] = "idle"
    state["last_processed_file"] = None
    print(f"[DEBUG] global RUDE_LOOP_STOP is set to True")
    global RUDE_LOOP_STOP
    RUDE_LOOP_STOP = True
    processing_stopped = FILE_PROCESSING_STOPPED.wait(timeout=5.0)
    if processing_stopped:
        if FILE_SESSION_INITIALIZED:
            session_file.reset_buffer()
        RUDE_LOOP_STOP = False
    else:
        print("[WARN] Timed out waiting for file processing to stop; buffer reset was skipped.")
    fresh_state, diar_html_output, _, _ = load_and_reset()
    state.clear()
    state.update(fresh_state)
    return state, None, diar_html_output, ""

def close_file_processing(state):
    """Clear the loaded file and hide the nearby Close button."""
    state, filepath, diar_html_output, transcript = stop_file_processing_example(state)
    return (
        state,
        filepath,
        diar_html_output,
        transcript,
        gr.Button(value="✕", visible=False),
    )

def update_file_close_button(filepath):
    """Show the Close button only while a file is loaded."""
    return gr.Button(value="✕", visible=bool(filepath))

def toggle_examples_panel(is_visible):
    """Show or hide the example picker in the File Upload tab."""
    is_visible = not bool(is_visible)
    label = "Hide Examples" if is_visible else "Choose Example"
    return is_visible, gr.update(visible=is_visible), label

def hide_examples_panel():
    """Collapse the example picker after an example is selected."""
    return False, gr.update(visible=False), "Choose Example"

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def reset_on_tab_switch(state):
    """Stop active processing and recreate both streaming instances."""
    global MICROPHONE_PROCESSING_ENABLED, RUDE_LOOP_STOP
    RUDE_LOOP_STOP = True
    MICROPHONE_PROCESSING_ENABLED = False
    processing_stopped = FILE_PROCESSING_STOPPED.wait(timeout=5.0)
    if not processing_stopped:
        print("[WARN] Timed out waiting for file processing during tab switch.")

    if MIC_SESSION_INITIALIZED:
        with MICROPHONE_PROCESSING_LOCK:
            session_mic.setup_streaming_session(config_mic, reload=True)
    if processing_stopped:
        if FILE_SESSION_INITIALIZED:
            session_file.setup_streaming_session(config_file, reload=True)
        RUDE_LOOP_STOP = False

    fresh_state, diar_html_output, transcript_html, _ = load_and_reset()
    state.clear()
    state.update(fresh_state)
    return (
        state,
        None,
        gr.update(recording=False),
        diar_html_output,
        transcript_html,
        gr.Button(value="Pause", visible=False),
        gr.Button(value="✕", visible=False),
        False,
        gr.update(visible=False),
        "Choose Example",
    )

def stop_file_processing_upload(state):
    """Stops file processing when audio is cleared."""
    state["should_stop"] = False
    state["mode"] = "idle"
    state["last_processed_file"] = None
    file_path = None
    global RUDE_LOOP_STOP
    RUDE_LOOP_STOP = True
    time.sleep(1.3)  # Give time for loop to stop
    RUDE_LOOP_STOP = False  # Reset for next processing
    state["should_stop"] = False
    return state, file_path

def trigger_end_of_file_flag(state):
    """Triggers the end of file flag."""
    state["end_of_file"] = True
    print(f"[DEBUG] END OF FILE TRIGGERED trigger_end_of_file_flag(): end_of_file set to {state['end_of_file']}")
    return state

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def start_stream_processing(state):
    """Initiates stream processing mode."""
    ensure_sessions_initialized(load_mic=True, load_file=False)
    global MICROPHONE_PROCESSING_ENABLED
    print(f"[DEBUG] start_stream_processing() CALLED - mode before: {state.get('mode', 'None')}")
    MICROPHONE_PROCESSING_ENABLED = True
    if state.get("mode") == "stream":
        return state, gr.Button(value="Pause", visible=True)

    state["stream_step_num"] = 0
    state["mode"] = "stream"
    state["microphone_paused"] = False
    state["full_transcript"] = ""
    state["diar_visualization_state"] = None

    print(f"[DEBUG] start_stream_processing() RETURNING - mode after: '{state['mode']}'")
    return state, gr.Button(value="Pause", visible=True)

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def toggle_microphone_pause(state):
    """Pause or resume microphone transcription while recording continues."""
    ensure_sessions_initialized(load_mic=True, load_file=False)
    if state.get("mode") != "stream":
        return (
            state,
            gr.Button(value="Pause", visible=False),
            state.get("diar_html", state.get("diar_html_initial", "")),
            format_transcript(state.get("full_transcript", "")),
        )

    pausing = not state.get("microphone_paused", False)
    state["microphone_paused"] = pausing
    if pausing:
        with MICROPHONE_PROCESSING_LOCK:
            previous_step = session_mic.step_num
            diar_state = session_mic.flush_microphone_buffer()
        state["stream_step_num"] += session_mic.step_num - previous_step
        state["full_transcript"] = session_mic.transcribed_speaker_texts
        if diar_state is not None:
            state["diar_visualization_state"] = diar_state
            state["diar_html"] = render_diarization_state(diar_state, config_mic)

    label = "Resume" if pausing else "Pause"
    return (
        state,
        gr.Button(value=label, visible=True),
        state["diar_html"],
        format_transcript(state.get("full_transcript", "")),
    )

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def stop_stream_processing(state):
    """Stops stream processing mode."""
    ensure_sessions_initialized(load_mic=True, load_file=False)
    global MICROPHONE_PROCESSING_ENABLED
    print(f"[DEBUG] stop_stream_processing() CALLED - mode before: {state.get('mode', 'None')}")
    MICROPHONE_PROCESSING_ENABLED = False
    
    # Reset the session
    with MICROPHONE_PROCESSING_LOCK:
        session_mic.reset_buffer()
        session_mic.setup_streaming_session(config_mic, reload=True)
    
    # Get fresh initial state from load_and_reset() to ensure complete cleanup
    fresh_state, diar_html_output, transcript_textbox_html, _ = load_and_reset()
    
    # Update the state with fresh values
    state.update(fresh_state)
    state["mode"] = "idle"
    state["microphone_paused"] = False
    state["full_transcript"] = ""
    state["stream_step_num"] = 0
    print(f"[DEBUG] stop_stream_processing() RETURNING - mode after: '{state['mode']}'")

    return state, diar_html_output, transcript_textbox_html, gr.Button(value="Pause", visible=False)

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def stream_processing(state, audio_data):
    """Handles microphone streaming."""
    ensure_sessions_initialized(load_mic=True, load_file=False)
    # Initialize state if it's None or doesn't have required fields
    if state is None:
        state = {"mode": "stream", "full_transcript": "", "diar_html_initial": ""}
    
    # Ensure diar_html_initial exists
    if "diar_html_initial" not in state or state["diar_html_initial"] == "":
        state["diar_html_initial"] = render_diarization_state(None, config_mic)
    
    # Check if we're in streaming mode
    current_mode = state.get("mode", "idle")
    if current_mode != "stream" or not MICROPHONE_PROCESSING_ENABLED:
        return gr.skip(), gr.skip(), gr.skip()

    print(f"[DEBUG] stream_processing(): current mode = '{current_mode}', audio_data type: {type(audio_data)}")

    if state.get("microphone_paused", False):
        return gr.skip(), gr.skip(), gr.skip()
    
    if audio_data is None: 
        print(f"[DEBUG] stream_processing() L2: audio_data is None, Aborting...")
        return gr.skip(), gr.skip(), gr.skip()
    
    sr, chunk_audio = audio_data
    if chunk_audio is None: 
        print(f"[DEBUG] stream_processing() L3: chunk_audio is None, Aborting...")
        return gr.skip(), gr.skip(), gr.skip()

    with MICROPHONE_PROCESSING_LOCK:
        previous_step = session_mic.step_num
        diar_state = session_mic.process_microphone_chunk(chunk_audio, sr)
    processed_steps = session_mic.step_num - previous_step
    if processed_steps == 0:
        return gr.skip(), gr.skip(), gr.skip()

    if state["stream_step_num"] == 0 and not session_mic.transcribed_speaker_texts:
        state["full_transcript"] += (
            '<div style="text-align: center; font-size: 20px; font-weight: bold;">Streaming Transcription Started...</div>'
        )
    else:
        state["full_transcript"] = session_mic.transcribed_speaker_texts

    state["stream_step_num"] += processed_steps
    print(f"[DEBUG] stream_processing() L4: state['stream_step_num'] = {state['stream_step_num']}")
    # Display the last 50 strings ofr session_mic.transcribed_speaker_texts
    
    if diar_state is not None:
        state["diar_visualization_state"] = diar_state
    else:
        diar_state = state.get("diar_visualization_state", None)

    if diar_state is not None:
        html_diar_out = render_diarization_state(diar_state, config_mic)
    else:
        html_diar_out = state["diar_html_initial"]
    state["diar_html"] = html_diar_out
    
    return state, html_diar_out, format_transcript(state["full_transcript"])

def reload_models_with_slider_update(
    audio_pre_loading,
    streaming_chunk_size,
    single_speaker_mode,
    num_speakers,
    cache_gating,
    binary_diar,
    sent_break_sec,
):
    """Wrapper that reloads models and updates the slider when single speaker mode changes."""
    if single_speaker_mode:
        num_speakers = 1
    
    html_string = reload_models(
        audio_pre_loading,
        streaming_chunk_size,
        single_speaker_mode,
        num_speakers,
        cache_gating,
        binary_diar,
        sent_break_sec,
    )
    return html_string, num_speakers

def reload_models_with_checkbox_update(
    audio_pre_loading,
    streaming_chunk_size,
    single_speaker_mode,
    num_speakers,
    cache_gating,
    binary_diar,
    sent_break_sec,
):
    """Wrapper that reloads models and updates the checkbox when num_speakers slider changes."""
    if num_speakers > 1:
        single_speaker_mode = False

    html_string = reload_models(
        audio_pre_loading,
        streaming_chunk_size,
        single_speaker_mode,
        num_speakers,
        cache_gating,
        binary_diar,
        sent_break_sec,
    )
    return html_string, single_speaker_mode

@spaces.GPU(duration=GPU_DURATION_SECONDS)
def reload_models(
    audio_pre_loading,
    streaming_chunk_size,
    single_speaker_mode,
    num_speakers,
    cache_gating,
    binary_diar,
    sent_break_sec,
):
    """Reloads the streaming session with new configurations."""
    try:
        num_speakers = min(int(num_speakers), MAX_NUM_OF_SPEAKERS)
        if single_speaker_mode:
            num_speakers = 1
        
        # Update Microphone and File config with new paths and settings
        config_mic.diar_model = args.diar_model
        config_mic.asr_model = args.asr_model
        config_mic.att_context_size = [LEFT_CONTEXT_SIZE, streaming_chunk_size]
        config_mic.target_lang = args.target_lang
        config_mic.masked_asr = True
        config_mic.single_speaker_mode = single_speaker_mode
        config_mic.sent_break_sec = sent_break_sec
        config_mic.max_num_of_spks = num_speakers
        config_mic.audio_pre_loading = audio_pre_loading
        config_mic.cache_gating = cache_gating
        config_mic.binary_diar_preds = binary_diar

        # Update File Streaming config with new paths and settings
        config_file.diar_model = args.diar_model
        config_file.asr_model = args.asr_model
        config_file.att_context_size = [LEFT_CONTEXT_SIZE, streaming_chunk_size]
        config_file.target_lang = args.target_lang
        config_file.masked_asr = True
        config_file.single_speaker_mode = single_speaker_mode
        config_file.sent_break_sec = sent_break_sec
        config_file.max_num_of_spks = num_speakers
        config_file.audio_pre_loading = audio_pre_loading
        config_file.cache_gating = cache_gating
        config_file.binary_diar_preds = binary_diar

        # Reload only sessions that are already resident. Config changes made
        # before first use are picked up by the corresponding lazy load.
        if MIC_SESSION_INITIALIZED:
            with MICROPHONE_PROCESSING_LOCK:
                session_mic.setup_streaming_session(config_mic, reload=True)
        if FILE_SESSION_INITIALIZED:
            session_file.setup_streaming_session(config_file, reload=True)

        # Assert all these parameters are in session_file.multispk_asr_streamer
        for session, initialized in (
            (session_file, FILE_SESSION_INITIALIZED),
            (session_mic, MIC_SESSION_INITIALIZED),
        ):
            if not initialized:
                continue
            streamer = session.multispk_asr_streamer
            assert streamer._att_context_size == [LEFT_CONTEXT_SIZE, streaming_chunk_size]
            assert streamer._single_speaker_mode == single_speaker_mode
            assert streamer._max_num_of_spks == num_speakers
            assert streamer._sent_break_sec == sent_break_sec
            assert streamer._cache_gating == cache_gating
            assert streamer._binary_diar_preds == binary_diar

        _, _, _, status_textbox_html = load_and_reset()
        return status_textbox_html
    except Exception as e:
        return f"<span style='color: red;'>Error reloading models: {str(e)}</span>"

def get_text_box_html(
    audio_pre_loading,
    streaming_chunk_size,
    single_speaker_mode,
    number_of_speakers,
    cache_gating,
    binary_diar,
    text_break_silence,
):
    # Use joiin to make line changes
    # Bold the parameter values
    ITEM_NAME_COLOR = "#F5F5F5"
    return "\n".join([
        f'<span style="color: #76B900; font-weight: bold;">Configs Applied Successfully!</span> <br> <br>',
        f'<span style="color: #F5F5F5;">Audio Pre-loading :</span> <span style="font-weight: bold;">{audio_pre_loading}</span> <br>',
        f'<span style="color: #F5F5F5;">Streaming Chunk Size :</span> <span style="font-weight: bold;">{streaming_chunk_size}</span> token(s) <br>',
        f'<span style="color: #F5F5F5;">Single Speaker Mode :</span> <span style="font-weight: bold;">{single_speaker_mode}</span> <br>',
        f'<span style="color: #F5F5F5;">Number of Speakers :</span> <span style="font-weight: bold;">{number_of_speakers}</span> <br>',
        f'<span style="color: #F5F5F5;">Cache Gating :</span> <span style="font-weight: bold;">{cache_gating}</span> <br>',
        f'<span style="color: #F5F5F5;">Binary Diarization Decision :</span> <span style="font-weight: bold;">{binary_diar}</span> <br>',
        f'<span style="color: #F5F5F5;">Text Break Silence :</span> <span style="font-weight: bold;">{text_break_silence}</span> seconds <br>',
    ])

def load_and_reset():
    """Resets the UI state on page load."""
    initialized_session = None
    if FILE_SESSION_INITIALIZED:
        initialized_session = session_file
    elif MIC_SESSION_INITIALIZED:
        initialized_session = session_mic

    if initialized_session is not None:
        streamer = initialized_session.multispk_asr_streamer
        AUDIO_PRE_LOADING = initialized_session._audio_pre_loading
        STREAMING_CHUNK_SIZE = streamer._att_context_size[1]
        SINGLE_SPEAKER_MODE = 'ON' if streamer._single_speaker_mode else 'OFF'
        NUMBER_OF_SPEAKERS = f"{streamer._max_num_of_spks}"
        TEXT_BREAK_SILENCE = f"{streamer._sent_break_sec} seconds"
        CACHE_GATING_STATUS = 'ON' if streamer._cache_gating else 'OFF'
        BINARY_DIAR_PREDICTION_STATUS = 'ON' if streamer._binary_diar_preds else 'OFF'
    else:
        AUDIO_PRE_LOADING = config_file.audio_pre_loading
        STREAMING_CHUNK_SIZE = config_file.att_context_size[1]
        SINGLE_SPEAKER_MODE = 'ON' if config_file.single_speaker_mode else 'OFF'
        NUMBER_OF_SPEAKERS = f"{config_file.max_num_of_spks}"
        TEXT_BREAK_SILENCE = f"{config_file.sent_break_sec} seconds"
        CACHE_GATING_STATUS = 'ON' if config_file.cache_gating else 'OFF'
        BINARY_DIAR_PREDICTION_STATUS = 'ON' if config_file.binary_diar_preds else 'OFF'

    status_textbox_html= get_text_box_html(
        audio_pre_loading=AUDIO_PRE_LOADING,
        streaming_chunk_size=STREAMING_CHUNK_SIZE,
        single_speaker_mode=SINGLE_SPEAKER_MODE,
        number_of_speakers=NUMBER_OF_SPEAKERS, 
        cache_gating=CACHE_GATING_STATUS, 
        binary_diar=BINARY_DIAR_PREDICTION_STATUS, 
        text_break_silence=TEXT_BREAK_SILENCE,
        )
    diar_html_output = render_diarization_state(None, config_file)
    transcript_textbox = "<p style='text-align: center; font-size:20px;'><b>No transcription yet. Models load on first transcription request.</b></p>"
    return {"should_stop": False,
            "stream_step_num": 0,
            "microphone_paused": False,
            "diar_visualization_state": None,
            "diar_html": copy.deepcopy(diar_html_output),
            "diar_html_initial": copy.deepcopy(diar_html_output),
            "end_of_file": False,
            "mode": "idle", 
            "full_transcript": "", 
            "finished_file_list": {}}, diar_html_output, transcript_textbox, status_textbox_html 

SUPPORTED_AUDIO_EXTENSIONS = {".flac", ".m4a", ".mp3", ".ogg", ".wav"}
examples = sorted(
    path
    for path in glob.glob("audio_data/**/*", recursive=True)
    if os.path.isfile(path) and os.path.splitext(path)[1].lower() in SUPPORTED_AUDIO_EXTENSIONS
)

# Custom CSS for NVIDIA Sans font and styling
custom_css = """
@font-face {
    font-family: 'NVIDIA Sans';
    src: url('file/fonts/NVIDIASans-Light.ttf') format('truetype');
    font-weight: 300;
    font-style: normal;
}

@font-face {
    font-family: 'NVIDIA Sans';
    src: url('file/fonts/NVIDIASans-Regular.ttf') format('truetype');
    font-weight: 400;
    font-style: normal;
}

@font-face {
    font-family: 'NVIDIA Sans';
    src: url('file/fonts/NVIDIASans-Medium.ttf') format('truetype');
    font-weight: 500;
    font-style: normal;
}

@font-face {
    font-family: 'NVIDIA Sans';
    src: url('file/fonts/NVIDIASans-Bold.ttf') format('truetype');
    font-weight: 700;
    font-style: normal;
}

/* Apply NVIDIA Sans to most elements */
body, .gradio-container, button, input, textarea, select, label, p, span, div {
    font-family: 'NVIDIA Sans', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
    font-size: 14px !important;
}

html,
body {
    margin: 0 !important;
    padding: 0 !important;
}

.gradio-container {
    width: 100% !important;
    max-width: 100% !important;
    min-height: 100vh !important;
    margin: 0 !important;
    padding: 2px 4px !important;
}

#main_layout {
    align-items: flex-start !important;
    width: 100% !important;
    margin: 0 !important;
    padding: 0 !important;
    gap: 4px !important;
}

#controls_panel {
    min-width: 270px !important;
    max-width: 340px !important;
}

#results_panel {
    min-width: 0 !important;
    gap: 4px !important;
}

#mode_tabs {
    position: relative !important;
    z-index: 1 !important;
    margin-top: -30px !important;
    margin-bottom: 0 !important;
}

#mode_tabs .tab-wrapper {
    padding-bottom: 0 !important;
    margin-bottom: 0 !important;
}

#app_header {
    min-height: 0 !important;
    padding: 0 8px 2px 0 !important;
    margin: 0 !important;
    border: 0 !important;
    background: transparent !important;
    text-align: right !important;
    position: relative !important;
    z-index: 2 !important;
    pointer-events: none !important;
}

#app_header > div,
#app_header .wrap,
#app_header .prose {
    min-height: 0 !important;
    padding: 0 !important;
    margin: 0 !important;
    gap: 0 !important;
}

#app_header h1 {
    font-size: clamp(1.2rem, 1.6vw, 1.65rem) !important;
    line-height: 1.15 !important;
    margin: 0 !important;
    text-align: right !important;
}

#app_header p {
    font-size: clamp(0.74rem, 0.85vw, 0.9rem) !important;
    line-height: 1.25 !important;
    margin: 3px 0 !important;
    text-align: right !important;
}

#app_header a {
    position: relative !important;
    z-index: 3 !important;
    pointer-events: auto !important;
    cursor: pointer !important;
}

#controls_panel .block {
    padding-top: 5px !important;
    padding-bottom: 5px !important;
}

#controls_panel label,
#controls_panel span,
#controls_panel input,
#controls_panel button {
    font-size: 12px !important;
}

#audio_player {
    height: 160px !important;
    min-height: 160px !important;
    max-height: 160px !important;
    overflow: visible !important;
}

#audio_player > div,
#audio_player .wrap,
#audio_player .upload-container {
    height: 100% !important;
    min-height: 0 !important;
    max-height: 100% !important;
}

#audio_player audio {
    height: 32px !important;
}

#audio_player [class*="waveform"],
#audio_player canvas {
    max-height: 42px !important;
    min-height: 42px !important;
}

#audio_player [class*="controls"] {
    display: flex !important;
    min-height: 34px !important;
    visibility: visible !important;
}

#file_audio_panel {
    position: relative !important;
}

#audio_player .rewind,
#audio_player .skip {
    display: none !important;
}

#file_close_button {
    position: absolute !important;
    left: calc(50% + 48px) !important;
    bottom: 9px !important;
    width: 32px !important;
    min-width: 32px !important;
    z-index: 10 !important;
}

#file_close_button button {
    height: 32px !important;
    min-height: 32px !important;
    font-size: 18px !important;
    padding: 0 !important;
}

#choose_examples_button {
    height: 24px !important;
    min-height: 24px !important;
    max-height: 24px !important;
    margin-top: -14px !important;
    padding: 0 !important;
    flex: none !important;
}

#choose_examples_button button {
    height: 24px !important;
    min-height: 24px !important;
    max-height: 24px !important;
    padding: 0 10px !important;
}

#microphone_panel {
    position: relative !important;
}

#mic_pause_resume_button {
    position: absolute !important;
    left: 110px !important;
    bottom: 14px !important;
    width: 88px !important;
    min-width: 88px !important;
    z-index: 10 !important;
}

#mic_pause_resume_button button {
    height: 32px !important;
    min-height: 32px !important;
}

#diarization_viz {
    width: 100% !important;
    min-width: 0 !important;
    border: 1px solid #76B900 !important;
    border-radius: 6px !important;
    box-sizing: border-box !important;
    padding: 6px 10px !important;
    margin: 0 !important;
}

#diarization_viz .diar-map {
    width: 100% !important;
    box-sizing: border-box !important;
    border: 0 !important;
    border-radius: 0 !important;
    background: #000 !important;
    padding: 0 !important;
    display: flex !important;
    flex-direction: column !important;
    gap: 3px !important;
}

#diarization_viz .diar-partition-labels {
    display: flex !important;
    justify-content: center !important;
    gap: 18px !important;
    width: 100% !important;
    margin-bottom: 5px !important;
}

#diarization_viz .diar-partition-label {
    color: #A8A8A8 !important;
    font-size: 10px !important;
    font-weight: 600 !important;
    line-height: 1.1 !important;
    text-align: center !important;
    letter-spacing: 0.04em !important;
}

#diarization_viz .diar-row {
    width: 100% !important;
    height: clamp(10px, 1.45vh, 16px) !important;
    display: grid !important;
    gap: 0 !important;
    overflow: hidden !important;
}

#diarization_viz .diar-partition {
    display: grid !important;
    min-width: 0 !important;
    height: 100% !important;
    gap: 0 !important;
}

#diarization_viz .diar-partition + .diar-partition {
    border-left: 2px solid #5A5A5A !important;
}

#diarization_viz .diar-cell {
    display: block !important;
    width: 100% !important;
    height: 100% !important;
}

#transcription_viz {
    width: 100% !important;
    height: clamp(220px, calc(26vh + 30px), 310px) !important;
    min-height: 220px !important;
    max-height: 310px !important;
    overflow-y: auto !important;
    overflow-x: hidden !important;
    scroll-behavior: auto;
    padding: 6px 10px !important;
    border: 1px solid #76B900 !important;
    border-radius: 6px !important;
    box-sizing: border-box !important;
    contain: layout paint !important;
    margin: 0 !important;
}

#transcription_viz,
#transcription_viz * {
    max-width: 100% !important;
    box-sizing: border-box !important;
    overflow-wrap: anywhere !important;
}

#transcription_viz .prose {
    height: 100% !important;
    overflow-y: auto !important;
    padding: 0 !important;
    margin: 0 !important;
}

#transcription_viz,
#transcription_viz span,
#transcription_viz p,
#transcription_viz div {
    font-size: 13px !important;
    line-height: 1.25 !important;
}

@media (max-width: 1100px) {
    #controls_panel {
        max-width: none !important;
    }
    #main_layout {
        flex-wrap: wrap !important;
    }
}

footer {
    visibility: hidden;
}
"""

# Create custom dark theme with #76B900 as primary color and black background
custom_theme = gr.themes.Default(
    primary_hue=gr.themes.Color(
        c50="#76B900",
        c100="#dcf4c7",
        c200="#c8ee9f",
        c300="#b4e877",
        c400="#9fe24f",
        c500="#76B900",  # Main color
        c600="#5e9400",
        c700="#466f00",
        c800="#2e4a00",
        c900="#162500",
        c950="#0b1200",
    ),
    neutral_hue=gr.themes.Color(
        c50="#0a0a0a",
        c100="#b2b2b2",
        c200="#b2b2b2",
        c300="#b2b2b2", # Countours of Checkboxes.
        c400="#c2c2c2", # Small letters in sliders and audio players.
        c500="#5a5a5a",
        c600="#6a6a6a",
        c700="#7a7a7a",
        c800="#8a8a8a",
        c900="#9a9a9a",
        c950="#aaaaaa",
    ),
    text_size=gr.themes.sizes.text_md,
    font=("NVIDIA Sans", "ui-sans-serif", "sans-serif"),
    font_mono=("ui-monospace", "monospace"),
).set(
    # Background colors
    body_background_fill="#000000",
    body_background_fill_dark="#000000",
    background_fill_primary="#1a1a1a",
    background_fill_primary_dark="#1a1a1a",
    background_fill_secondary="#0a0a0a",
    background_fill_secondary_dark="#0a0a0a",
    
    # Block/Panel colors
    block_background_fill="#1a1a1a",
    block_background_fill_dark="#1a1a1a",
    block_border_color="#2a2a2a",
    block_border_color_dark="#2a2a2a",
    
    # Input/Component colors
    input_background_fill="#2a2a2a",
    input_background_fill_dark="#2a2a2a",
    input_border_color="#9a9a9a",
    input_border_color_dark="#9a9a9a",
    
    # Button colors
    button_primary_background_fill="#a0a0a0",
    button_primary_background_fill_hover="#76B900",
    button_primary_background_fill_dark="#a0a0a0",
    button_primary_background_fill_hover_dark="#a0a0a0",
    button_primary_text_color="white",
    button_primary_text_color_dark="white",
    button_secondary_background_fill="#a0a0a0",
    button_secondary_background_fill_hover="#76B900",
    button_secondary_background_fill_dark="#a0a0a0",
    button_secondary_background_fill_hover_dark="#a0a0a0",
    button_secondary_text_color="white",
    button_secondary_text_color_dark="white",
    
    # Slider colors
    slider_color="#76B900",
    slider_color_dark="#76B900",
    
    # Text colors
    body_text_color="#e0e0e0",
    body_text_color_dark="#e0e0e0",
    block_title_text_color="#ffffff",
    block_title_text_color_dark="#ffffff",
    block_label_text_color="#e0e0e0",
    block_label_text_color_dark="#e0e0e0",
    
    # Border and shadow
    block_shadow="0 0 0 1px rgba(118, 185, 0, 0.1)",
    block_shadow_dark="0 0 0 1px rgba(118, 185, 0, 0.1)",
)

auto_scroll_js = """
() => {
    const scrollToBottom = () => {
        const transcript = document.getElementById("transcription_viz");
        if (!transcript) return;
        window.requestAnimationFrame(() => {
            transcript.scrollTop = transcript.scrollHeight;
            const content = transcript.querySelector(".prose");
            if (content) content.scrollTop = content.scrollHeight;
        });
    };
    if (window.transcriptAutoScrollObserver) {
        window.transcriptAutoScrollObserver.disconnect();
    }
    window.transcriptAutoScrollObserver = new MutationObserver(scrollToBottom);
    window.transcriptAutoScrollObserver.observe(document.body, {
        childList: true,
        subtree: true,
        characterData: true
    });
    scrollToBottom();
}
"""

reset_audio_js = """
() => {
    window.clearTimeout(window.demoAudioResetTimer);
    let attempts = 0;

    const resetAudio = () => {
        const waveform = document.querySelector("#audio_player #waveform > div");
        const media = waveform?.shadowRoot?.querySelector("audio");
        if (!media) {
            if (++attempts < 40) {
                window.demoAudioResetTimer = window.setTimeout(resetAudio, 50);
            }
            return;
        }

        const rewind = () => {
            media.pause();
            media.currentTime = 0;
            media.dispatchEvent(new Event("timeupdate"));
        };
        media.addEventListener("loadedmetadata", rewind, { once: true });
        window.requestAnimationFrame(rewind);
    };

    resetAudio();
}
"""

stop_microphone_recording_js = """
(state) => {
    const root = document.querySelector("#microphone_input");
    const stopButton = Array.from(root?.querySelectorAll("button") || []).find((button) => {
        const label = [
            button.getAttribute("aria-label"),
            button.getAttribute("title"),
            button.textContent,
        ].filter(Boolean).join(" ").trim().toLowerCase();
        return label.includes("stop") && window.getComputedStyle(button).display !== "none";
    });
    stopButton?.click();
    return state;
}
"""

# Load article from HTML file
with open("space_article.html", "r") as f:
    article = f.read()

with gr.Blocks(
    title="NeMo Multispeaker Cache Aware Streaming Demo",
    theme=custom_theme,
    css=custom_css,
    js=auto_scroll_js,
) as demo:
    # with gr.Row():
    #     with gr.Column(scale=4):
            # gr.HTML(article)
            # state = gr.State({"mode": "idle", "full_transcript": "", "finished_file_list": {}})

    with gr.Row(elem_id="main_layout"):
        with gr.Column(scale=1, min_width=270, elem_id="controls_panel"):
            audio_pre_loading_checkbox = gr.Checkbox(value=True, label="Audio Pre-loading (If network is slow)")
            streaming_chunk_size_slider = gr.Slider(minimum=1, maximum=13, step=1, value=13, label="Streaming Chunk Size (tokens)")
            num_speakers_slider = gr.Slider(
                minimum=1,
                maximum=MAX_NUM_OF_SPEAKERS,
                step=1,
                value=MAX_NUM_OF_SPEAKERS,
                label="Number of Speakers",
            )
            sent_break_sec_slider = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=1.0, label="Text Break Silence (sec)")
            single_speaker_mode_checkbox = gr.Checkbox(value=False, label="Single Speaker Mode")
            cache_gating_checkbox = gr.Checkbox(value=True, label="Cache Gating")
            binary_diar_checkbox = gr.Checkbox(value=True, label="Binarize Speaker Activity")
            load_model_button = gr.Button("Reset", variant="primary")
            status_textbox = gr.HTML(label="Status")
        
        with gr.Column(scale=4, min_width=640, elem_id="results_panel"):
            gr.HTML(article, elem_id="app_header")
            state = gr.State({"mode": "idle", "full_transcript": "", "finished_file_list": {}})
            examples_visible_state = gr.State(False)
            with gr.Tabs(elem_id="mode_tabs"):
                with gr.TabItem("Microphone") as microphone_tab:
                    with gr.Column(elem_id="microphone_panel"):
                        mic_input = gr.Audio(
                            sources=["microphone"],
                            streaming=True,
                            type="numpy",
                            label="Record from Microphone",
                            elem_id="microphone_input",
                        )
                        mic_pause_resume_button = gr.Button(
                            "Pause",
                            size="sm",
                            visible=False,
                            elem_id="mic_pause_resume_button",
                        )
                with gr.TabItem("File Upload") as file_upload_tab:
                    with gr.Column(elem_id="file_audio_panel"):
                        file_input_example = gr.Audio(
                            sources=["upload"],
                            type="filepath",
                            label="Audio Playback",
                            autoplay=False,
                            elem_id="audio_player",
                        )
                        file_close_button = gr.Button(
                            "✕",
                            size="sm",
                            visible=False,
                            elem_id="file_close_button",
                        )
                    choose_examples_button = gr.Button(
                        "Choose Example",
                        size="sm",
                        elem_id="choose_examples_button",
                    )
                    with gr.Column(visible=False) as examples_panel:
                        gr_examples = gr.Examples(
                            examples=examples,
                            inputs=[file_input_example],
                            label="Click an example to load it",
                        )
            # Create default empty visualization
            default_diar_html = render_diarization_state(None, config_file)
            diarization_logit_box = gr.HTML(value=default_diar_html, label="Speaker Activity Maps", elem_id="diarization_viz")
            transcription_text_box = gr.HTML(
                label="Transcription",
                elem_id="transcription_viz",
                js_on_load="""
                const scrollToBottom = () => requestAnimationFrame(() => {
                    const scroller = element.closest("#transcription_viz") || element;
                    scroller.scrollTop = scroller.scrollHeight;
                });
                watch("value", scrollToBottom);
                scrollToBottom();
                """,
            )
    reload_inputs = [
        audio_pre_loading_checkbox,
        streaming_chunk_size_slider,
        single_speaker_mode_checkbox,
        num_speakers_slider,
        cache_gating_checkbox,
        binary_diar_checkbox,
        sent_break_sec_slider,
    ]
    file_input_example.change(
        fn=None,
        js=reset_audio_js,
        queue=False,
    )
    file_input_example.change(
        fn=update_file_close_button,
        inputs=[file_input_example],
        outputs=[file_close_button],
        queue=False,
    )
    choose_examples_button.click(
        fn=toggle_examples_panel,
        inputs=[examples_visible_state],
        outputs=[examples_visible_state, examples_panel, choose_examples_button],
        queue=False,
    )

    tab_reset_outputs = [
        state,
        file_input_example,
        mic_input,
        diarization_logit_box,
        transcription_text_box,
        mic_pause_resume_button,
        file_close_button,
        examples_visible_state,
        examples_panel,
        choose_examples_button,
    ]
    microphone_tab.select(
        fn=reset_on_tab_switch,
        inputs=[state],
        outputs=tab_reset_outputs,
        queue=False,
    )
    file_upload_tab.select(
        fn=reset_on_tab_switch,
        inputs=[state],
        outputs=tab_reset_outputs,
        js=stop_microphone_recording_js,
        queue=False,
    )

    load_model_button.click(
        fn=reload_models,
        inputs=reload_inputs,
        outputs=[status_textbox],
    )
    # Auto-reload when any setting changes
    streaming_chunk_size_slider.change(
        fn=reload_models,
        inputs=reload_inputs,
        outputs=[status_textbox],
    )
    single_speaker_mode_checkbox.change(
        fn=reload_models_with_slider_update,
        inputs=reload_inputs,
        outputs=[status_textbox, num_speakers_slider],
    )
    num_speakers_slider.change(
        fn=reload_models_with_checkbox_update,
        inputs=reload_inputs,
        outputs=[status_textbox, single_speaker_mode_checkbox],
    )
    sent_break_sec_slider.change(
        fn=reload_models,
        inputs=reload_inputs,
        outputs=[status_textbox],
    )
    audio_pre_loading_checkbox.change(
        fn=reload_models,
        inputs=reload_inputs,
        outputs=[status_textbox],
    )
    cache_gating_checkbox.change(
        fn=reload_models,
        inputs=reload_inputs,
        outputs=[status_textbox],
    )
    binary_diar_checkbox.change(
        fn=reload_models,
        inputs=reload_inputs,
        outputs=[status_textbox],
    )

    mic_pause_resume_button.click(
        fn=toggle_microphone_pause,
        inputs=[state],
        outputs=[state, mic_pause_resume_button, diarization_logit_box, transcription_text_box],
        queue=False,
    )
    file_close_button.click(
        fn=close_file_processing,
        inputs=[state],
        outputs=[
            state,
            file_input_example,
            diarization_logit_box,
            transcription_text_box,
            file_close_button,
        ],
    )
    
    mic_input.start_recording(
        fn=start_stream_processing,
        inputs=[state],
        outputs=[state, mic_pause_resume_button],
        queue=False,
    )
    mic_input.stream(
        fn=stream_processing,
        inputs=[state, mic_input],
        outputs=[state, diarization_logit_box, transcription_text_box],
        show_progress="hidden",
        stream_every=0.2,
    )
    mic_input.stop_recording(
        fn=stop_stream_processing,
        inputs=[state],
        outputs=[state, diarization_logit_box, transcription_text_box, mic_pause_resume_button],
    )
    
    file_input_example.upload(fn=file_processing_example_click_upload, 
                              inputs=[file_input_example, state], 
                              outputs=[state, diarization_logit_box, 
                              transcription_text_box,])
    file_input_example.stop(fn=trigger_end_of_file_flag, inputs=[state], outputs=[state])
    file_input_example.play(fn=turn_state_to_resume, inputs=[state], outputs=[state])
    file_input_example.pause(fn=turn_state_to_pause, inputs=[state], outputs=[state])
    file_input_example.clear(
        fn=close_file_processing,
        inputs=[state],
        outputs=[
            state,
            file_input_example,
            diarization_logit_box,
            transcription_text_box,
            file_close_button,
        ],
    )
    
    selected_filepath = gr.State()  # Hidden state to hold the selected filepath
    
    gr_examples.dataset.select(
        fn=update_config_for_file, 
        inputs=[state], 
        outputs=[
            state,
            selected_filepath,         # filepath extracted from evt.value
            streaming_chunk_size_slider,
            num_speakers_slider,
            sent_break_sec_slider,
            audio_pre_loading_checkbox,
            single_speaker_mode_checkbox,
            cache_gating_checkbox,
            binary_diar_checkbox,
            status_textbox,
        ]
    ).then(
        fn=example_selection_stop_file_processing, 
        inputs=[state, selected_filepath],
        outputs=[state, selected_filepath],
    ).then(
        fn=hide_examples_panel,
        outputs=[examples_visible_state, examples_panel, choose_examples_button],
    ).then(
        fn=file_processing_example_click_examples, 
        inputs=[selected_filepath, state], 
        outputs=[state, diarization_logit_box, transcription_text_box]
    )

    demo.load(fn=load_and_reset, outputs=[state, diarization_logit_box, transcription_text_box, status_textbox])

if __name__ == "__main__":
    running_on_space = bool(os.environ.get("SPACE_ID"))
    if not running_on_space:
        enable_share_tunnel_retries()
    server_port = int(os.environ.get("PORT", os.environ.get("GRADIO_SERVER_PORT", "7860")))
    demo.launch(server_name="0.0.0.0", server_port=server_port, share=not running_on_space)