ui/events.py

"""
Event handlers and UI logic for subtitle merging in DuoSubs Hugging Face Web UI.

This module contains functions to handle merging of subtitles, and update UI elements 
based on user interactions.
"""

from pathlib import Path

import gradio as gr
from duosubs import (
    LoadSubsError,
    MergeArgs,
    MergeSubsError,
    MergingMode,
    OmitFile,
    SaveSubsError,
    SubtitleFormat,
    load_subtitles,
    merge_subtitles,
    save_subtitles_in_zip,
)
from sentence_transformers import SentenceTransformer


def start_merging(
        model_pool: dict[str, SentenceTransformer],
        primary_subtitles: str,
        secondary_subtitles: str,
        model_name: str,
        batch_size: int,
        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:
    """
    The main function to handle the merging process of subtitles, which starts from 
    loading subtitles, merging subtitles, and saving the output in a ZIP file.

    Args:
        model_pool (dict[str, SentenceTransformer]): Mapping of model names to shared 
            SentenceTransformer instances.
        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.
        batch_size (int): Batch size for inference.
        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).
        cancel_state (list[bool]): List tracking cancellation state.
        progress (gradio.Progress) : Gradio progress object (optional). 
            Defaults to None.

    Returns:
        str | None: Path to the output ZIP file, or None if cancelled.

    Raises:
        gradio.Error: If any error occurs during loading, merging, or saving subtitles.
    """
    if progress is None:
        progress = gr.Progress()

    args = MergeArgs(
        primary=primary_subtitles,
        secondary=secondary_subtitles,
        batch_size=int(batch_size),
        merging_mode=MergingMode(merging_mode.lower()),
        retain_newline=retain_newline,
        secondary_above=secondary_above_primary,
        omit=[OmitFile.EDIT],
        format_combined=SubtitleFormat(combined_format),
        format_primary=SubtitleFormat(primary_format),
        format_secondary=SubtitleFormat(secondary_format)
    )

    if "Combined" in omit_subtitles:
        args.omit.append(OmitFile.COMBINED)
    if "Primary" in omit_subtitles:
        args.omit.append(OmitFile.PRIMARY)
    if "Secondary" in omit_subtitles:
        args.omit.append(OmitFile.SECONDARY)
    
    zip_name_with_path: str | None = None

    if len(args.omit) == 4:
        gr.Warning(
            (
                "Nothing to merge — Please adjust "
                "<strong><em>Excluded Subtitle Files</em></strong> options "
                "in <strong><em>File Exports</em></strong>"
            ),
            duration=7
        )
        return zip_name_with_path

    try:
        if not cancel_state[0]:
            progress(progress=0, desc= "Stage 1 → Loading subtitles", total=1)
            primary_subs_data, secondary_subs_data = load_subtitles(args)
            progress(progress=1, desc= "Stage 1 → Loading subtitles", total=1)

        if not cancel_state[0]:
            def update_progress(current: int) -> None:
                progress(
                    progress=current/100,
                    desc= (
                        f"Stage 2 → Merging subtitles "
                        f"({args.merging_mode.value.capitalize()} mode, "
                        f"using {model_name})"
                    ),
                    total=100
                )
            merged_subs = merge_subtitles(
                args,
                model_pool[model_name],
                primary_subs_data,
                secondary_subs_data,
                cancel_state,
                progress_callback=update_progress
            )

        if not cancel_state[0]:
            full_zip_path_without_ext = str(Path(args.primary).with_suffix(""))
            zip_name_with_path = f"{full_zip_path_without_ext}.zip"
            zip_name = Path(zip_name_with_path).name
            progress(
                progress=0,
                desc= f"Stage 3 → Compressing files into {zip_name}",
                total=1
            )
            save_subtitles_in_zip(
                args,
                merged_subs,
                primary_subs_data.styles,
                secondary_subs_data.styles
            )
            progress(
                progress=1,
                desc= f"Stage 3 → Compressing files into {zip_name}",
                total=1
            )

        if cancel_state[0]:
            gr.Info("The merging process is stopped.", duration=7)

    except LoadSubsError as e1:
        raise gr.Error(str(e1)) from e1
    except MergeSubsError as e2:
        raise gr.Error(str(e2)) from e2
    except SaveSubsError as e3:
        raise gr.Error(str(e3)) from e3

    return zip_name_with_path

def cancel_merging(cancel_state: list[bool]) -> gr.Button:
    """
    Cancels the merging process and updates the UI state.

    Args:
        cancel_state (list[bool]): List tracking cancellation state.

    Returns:
        gradio.Button: Cancel button with updated interactivity.
    """
    cancel_state[0] = True
    gr.Info("Cancelling merge process...", duration=7)
    return gr.Button("Cancel", interactive=False)

def states_during_merging(cancel_state: list[bool]) -> tuple[gr.Button, gr.Button]:
    """
    Sets UI state for buttons during the merging process, which disables the Merge 
    button and enables the Cancel button.

    Args:
        cancel_state (list[bool]): List tracking cancellation state.

    Returns:
        tuple[gradio.Button, gradio.Button]: Updated Merge and Cancel buttons.
    """
    cancel_state[0] = False
    return gr.Button("Merge", interactive=False), gr.Button("Cancel", interactive=True)

def states_after_merging(cancel_state: list[bool]) -> tuple[gr.Button, gr.Button]:
    """
    Sets UI state for buttons after the merging process, which enables the Merge button
    and disables the Cancel button.

    Args:
        cancel_state (list[bool]): List tracking cancellation state.

    Returns:
        tuple[gradio.Button, gradio.Button]: Updated Merge and Cancel buttons.
    """
    cancel_state[0] = False
    return gr.Button("Merge", interactive=True), gr.Button("Cancel", interactive=False)

def validate_excluded_subtitle_file(selected: list[str]) -> None:
    """
    Validates the selected options for excluded subtitle files.
    If all options are selected, it shows a warning message.

    Args:
        selected (list[str]): List of selected options.
    """
    if len(selected) == 3:
        gr.Warning(
            (
                "Nothing to merge — Please adjust "
                "<strong><em>Excluded Subtitle Files</em></strong> options "
                "in <strong><em>File Exports</em></strong>"
            ),
            duration=7
        )