ui/layout.py

"""
Defines the main Gradio UI layout and configuration for DuoSubs subtitle merging 
Hugging Face web app.

This module sets up the UI components, event handlers, and manages the model loading
and merging process. It includes device selection, model configuration, and subtitle 
file handling.
"""

import gradio as gr

from .constants import (
    DEFAULT_SUB_EXT,
    MERGING_MODE_INFO,
    MERGING_MODE_LIST,
    MODEL_NAME_LIST,
    SUB_EXT_LIST,
    SUB_EXT_LIST_WITH_DOT,
    TITLE_HTML,
)
from .events import (
    cancel_merging,
    start_merging,
    states_after_merging,
    states_during_merging,
    validate_excluded_subtitle_file,
)
from .utils import create_model_pools, open_html

model_pool = create_model_pools(MODEL_NAME_LIST)

def create_main_gr_blocks_ui(
        cache_delete_frequency: int = 3600,
        cache_delete_age: int = 14400,
        merging_concurrency_limit: int = 5
    ) -> gr.Blocks:
    """
    Builds and returns the main Gradio Blocks UI for DuoSubs.

    Args:
        cache_delete_frequency (int): How often to delete cache (seconds).
            Default is 1 hour or 3600 seconds.
        cache_delete_age (int): Age threshold for cache deletion (seconds).
            Default is 4 hour or 14400 seconds.
        merging_concurrency_limit (int): Maximum number of concurrent merges allowed.
            Default is 5.

    Returns:
        gradio.Blocks: The constructed Gradio UI.
    """
    main_block = gr.Blocks(
        title="DuoSubs",
        theme=gr.themes.Ocean(),
        delete_cache=(cache_delete_frequency, cache_delete_age),
        analytics_enabled=False
    )
    ui: gr.Blocks
    with main_block as ui:
        global model_pool
        cancel_state = gr.State([False])

        title_content = open_html(TITLE_HTML)
        gr.HTML(title_content)
        with gr.Row():
            with gr.Column(scale=9):
                (
                    primary_file,
                    secondary_file,
                    merged_file,
                    merge_button,
                    cancel_button
                ) = _create_subtitles_io_block()
            with gr.Column(scale=11):
                gr.Markdown("### ⚙️ Configurations")
                with gr.Tab("Model Selection"):
                    model_name =  _create_model_configurations_block()
                with gr.Tab("Alignment Behavior"):
                    merging_mode = _create_alignment_behaviour_block()
                with gr.Tab("Output Styling"):
                    (
                        retain_newline,
                        secondary_above_primary
                    ) = _create_output_styling_block()
                with gr.Tab("File Exports"):
                    (
                        omit_subtitles,
                        combined_format,
                        primary_format,
                        secondary_format
                    ) = _create_file_exports_block()        

        omit_subtitles.change(
            fn=validate_excluded_subtitle_file,
            inputs=omit_subtitles
        )
        
        merge_button.click(
            fn=states_during_merging,
            inputs=cancel_state,
            outputs=[merge_button, cancel_button]
        ).then(
            fn=wrapped_start_merging,
            inputs=[
                primary_file,
                secondary_file,
                model_name,
                merging_mode, 
                retain_newline,
                secondary_above_primary,
                omit_subtitles,
                combined_format, 
                primary_format,
                secondary_format,
                cancel_state
            ],
            outputs=merged_file,
            concurrency_limit=merging_concurrency_limit
        ).then(
            fn=states_after_merging,
            inputs=cancel_state,
            outputs=[merge_button, cancel_button]
        )

        cancel_button.click(
            fn=cancel_merging,
            inputs=cancel_state,
            outputs=cancel_button,
            concurrency_limit=None
        )
        
    return ui
    
def _create_subtitles_io_block(
    ) -> tuple[gr.File, gr.File, gr.File, gr.Button, gr.Button]:
    """
    Creates subtitle file input/output UI components.

    This function sets up the UI for uploading primary and secondary subtitle files,
    buttons to initiate and cancel the merging process, and creates the merged output 
    file.

    Returns:
        tuple[gradio.File, gradio.File, gradio.File, gradio.Button, gradio.Button]:
        - primary_file
        - secondary_file
        - merged_file
        - merge_button
        - cancel_button
    """
    gr.Markdown("### 📄 Input Subtitles")
    with gr.Row():
        primary_file = gr.File(
            label="Primary Subtitle File",
            file_types=SUB_EXT_LIST_WITH_DOT
        )
        secondary_file = gr.File(
            label="Secondary Subtitle File",
            file_types=SUB_EXT_LIST_WITH_DOT
        )

    gr.Markdown("### 📦 Output Zip")
    merged_file = gr.File(label="Processed Subtitles (in zip)")

    with gr.Row():
        merge_button = gr.Button("Merge")
        cancel_button = gr.Button("Cancel", interactive=False)
    
    return primary_file, secondary_file, merged_file, merge_button, cancel_button

def _create_model_configurations_block() -> gr.Dropdown:
    """
    Creates model selection configuration UI components.

    This function sets up the UI for selecting the model name.

    Returns:
        gradio.Dropdown: model name
    """
    with gr.Column():
        model_name = gr.Dropdown(
            choices=MODEL_NAME_LIST,
            label="Sentence Transformer Model",
            value=MODEL_NAME_LIST[0],
            info=(
                "Model for computing subtitle sentence similarity, "
                "with float32 precision"
            )
        )
    return model_name

def _create_alignment_behaviour_block() -> gr.Radio:
    """
    Creates alignment behavior UI components.

    This function sets up a radio buttons for selecting the subtitles merging mode.
    
    Returns:
        gradio.Radio: Radio buttons for merging mode.
    """
    mode_content = open_html(MERGING_MODE_INFO)
    with gr.Column():
        merging_mode = gr.Radio(
            label="Merging Mode",
            choices=MERGING_MODE_LIST,
            value=MERGING_MODE_LIST[0],
            info="Please refer to **ℹ️ Info** below for more information" # noqa: RUF001
        )
        with gr.Accordion("ℹ️ Info"): # noqa: RUF001
            gr.HTML(mode_content)
    return merging_mode

def _create_output_styling_block() -> tuple[gr.Checkbox, gr.Checkbox]:
    """
    Creates output styling UI components.

    This function sets up checkboxes for retaining newlines in the original subtitles 
    and placing secondary subtitles above primary subtitles in the merged output.

    Returns:
        tuple[gradio.Checkbox, gradio.Checkbox]:
        - retain_newline
        - secondary_above_primary
    """
    with gr.Column():
        retain_newline = gr.Checkbox(
            label="Retain Newlines",
            value=False,
            info="**Retain line breaks** from the original subtitles"
        )
        secondary_above_primary = gr.Checkbox(
            label="Secondary subtitle above primary subtitle",
            value=False,
            info="Place **secondary** subtitle **above** the **primary**"
        )
    
    return retain_newline, secondary_above_primary

def _create_file_exports_block(
    ) -> tuple[gr.CheckboxGroup, gr.Dropdown, gr.Dropdown, gr.Dropdown]:
    """
    Creates file export UI components.

    This function sets up checkboxes for excluding certain subtitle files from the ZIP 
    output, and dropdowns for selecting the format of combined, primary, and secondary 
    subtitles.

    Returns:
        tuple[gradio.CheckboxGroup, gradio.Dropdown, gradio.Dropdown, gradio.Dropdown]:
        - omit_subtitles
        - combined_format
        - primary_format
        - secondary_format
    """
    with gr.Column():
        omit_subtitles = gr.CheckboxGroup(
            ["Combined", "Primary", "Secondary"],
            type="value",
            label="Excluded Subtitle Files from ZIP"
        )
    with gr.Column():
        gr.Markdown("Subtitle Format")
        combined_format = gr.Dropdown(
            choices=SUB_EXT_LIST,
            value=DEFAULT_SUB_EXT,
            label="Combined"
        )
        primary_format = gr.Dropdown(
            choices=SUB_EXT_LIST,
            value=DEFAULT_SUB_EXT,
            label="Primary"
        )
        secondary_format = gr.Dropdown(
            choices=SUB_EXT_LIST,
            value=DEFAULT_SUB_EXT,
            label="Secondary"
        )
    return omit_subtitles, combined_format, primary_format, secondary_format

def wrapped_start_merging(
        primary_subtitles: str,
        secondary_subtitles: str,
        model_name: str,
        merging_mode: str,
        retain_newline: bool,
        secondary_above_primary: bool,
        omit_subtitles: list[str],
        combined_format: str,
        primary_format:str,
        secondary_format: str,
        cancel_state: list[bool],
        progress: gr.Progress | None = None
    ) -> str | None:
    """
    Wrapper for starting the merging process with all required arguments.

    Args:
        primary_subtitles (str): Path to primary subtitle file.
        secondary_subtitles (str): Path to secondary subtitle file.
        model_name (str): Name of the model to use.
        merging_mode (str): Subttitle merging mode (based on the value of MergingMode).
        retain_newline (bool): Whether to retain newlines in output.
        secondary_above_primary (bool): Whether to place secondary subtitle above 
            primary.
        omit_subtitles (list[str]): List of subtitle types to omit from output (based 
            on the value of OmitFile).
        combined_format (str): Format for combined subtitles (based on the value of 
            SubtitleFormat).
        primary_format (str): Format for primary subtitles (based on the value of 
            SubtitleFormat).
        secondary_format (str): Format for secondary subtitles (based on the value of 
            SubtitleFormat).
        gpu_list (list[str]): List of available GPU names.
        loaded_model_device (list[str]): List tracking loaded model device.
        loaded_model_name (list[str]): List tracking loaded model name.
        cancel_state (list[bool]): List tracking cancellation state.
        request (gradio.Request): Gradio request object.
        progress (gradio.Progress) : Gradio progress object (optional).
            Defaults to None.

    Returns:
        str | None: Path to the output ZIP file, or None if cancelled.
    """
    return start_merging(
        model_pool,
        primary_subtitles,
        secondary_subtitles,
        model_name,
        256,
        merging_mode,
        retain_newline,
        secondary_above_primary,
        omit_subtitles,
        combined_format,
        primary_format,
        secondary_format,
        cancel_state,
        progress
    )