src/ultimate_rvc/web/common.py

"""
Module defining common utility functions and classes for the
web application of the Ultimate RVC project.
"""

from typing import Any, Concatenate

from collections.abc import Callable, Sequence

import gradio as gr

from ultimate_rvc.core.exceptions import NotProvidedError
from ultimate_rvc.core.generate.song_cover import (
    get_named_song_dirs,
    get_song_cover_name,
)
from ultimate_rvc.core.manage.audio import get_saved_output_audio
from ultimate_rvc.web.typing_extra import (
    ComponentVisibilityKwArgs,
    DropdownChoices,
    DropdownValue,
    TextBoxKwArgs,
    UpdateDropdownKwArgs,
)

PROGRESS_BAR = gr.Progress()


def exception_harness[T, **P](
    fn: Callable[P, T],
    info_msg: str | None = None,
) -> Callable[P, T]:
    """
    Wrap a function in a harness that catches exceptions and re-raises
    them as instances of `gradio.Error`.

    Parameters
    ----------
    fn : Callable[P, T]
        The function to wrap.

    info_msg : str, optional
        Message to display in an info-box pop-up after the function
        executes successfully.

    Returns
    -------
    Callable[P, T]
        The wrapped function.

    """

    def _wrapped_fn(*args: P.args, **kwargs: P.kwargs) -> T:
        try:
            res = fn(*args, **kwargs)
        except gr.Error:
            raise
        except NotProvidedError as e:
            msg = e.ui_msg or e
            raise gr.Error(str(msg)) from None
        except Exception as e:
            raise gr.Error(str(e)) from e
        else:
            if info_msg:
                gr.Info(info_msg, duration=0.5)
            return res

    return _wrapped_fn


def confirmation_harness[T, **P](
    fn: Callable[P, T],
) -> Callable[Concatenate[bool, P], T]:
    """
    Wrap a function in a harness that requires a confirmation before
    executing and catches exceptions, re-raising them as instances of
    `gradio.Error`.

    Parameters
    ----------
    fn : Callable[P, T]
        The function to wrap.

    Returns
    -------
    Callable[Concatenate[bool, P], T]
        The wrapped function.

    """

    def _wrapped_fn(confirm: bool, *args: P.args, **kwargs: P.kwargs) -> T:
        if confirm:
            return exception_harness(fn)(*args, **kwargs)
        err_msg = "Confirmation missing!"
        raise gr.Error(err_msg)

    return _wrapped_fn


def render_msg(
    template: str,
    *args: str,
    display_info: bool = False,
    **kwargs: str,
) -> str:
    """
    Render a message template with the provided arguments.

    Parameters
    ----------
    template : str
        Message template to render.
    args : str
        Positional arguments to pass to the template.
    display_info : bool, default=False
        Whether to display the rendered message as an info message
        in addition to returning it.
    kwargs : str
        Keyword arguments to pass to the template.

    Returns
    -------
    str
        Rendered message.

    """
    msg = template.format(*args, **kwargs)
    if display_info:
        gr.Info(msg)
    return msg


def confirm_box_js(msg: str) -> str:
    """
    Generate a JavaScript code snippet which:
      * defines an anonymous function that takes one named parameter and
      zero or more unnamed parameters
      * renders a confirmation box
      * returns the choice selected by the user in that confirmation
      box in addition to any unnamed parameters passed to the function.

    Parameters
    ----------
    msg : str
        Message to display in the confirmation box rendered by the
        JavaScript code snippet.

    Returns
    -------
    str
        The JavaScript code snippet.

    """
    return f"(x, ...args) => [confirm('{msg}'), ...args]"


def update_value(x: str) -> dict[str, Any]:
    """
    Update the value of a component.

    Parameters
    ----------
    x : str
        New value for the component.

    Returns
    -------
    dict[str, Any]
        Dictionary which updates the value of the component.

    """
    return gr.update(value=x)


def update_dropdowns[**P](
    fn: Callable[P, DropdownChoices],
    num_components: int,
    value: DropdownValue = None,
    value_indices: Sequence[int] = [],
    *args: P.args,
    **kwargs: P.kwargs,
) -> gr.Dropdown | tuple[gr.Dropdown, ...]:
    """
    Update the choices and optionally the value of one or more dropdown
    components.

    Parameters
    ----------
    fn : Callable[P, DropdownChoices]
        Function to get updated choices for the dropdown components.
    num_components : int
        Number of dropdown components to update.
    value : DropdownValue, optional
        New value for dropdown components.
    value_indices : Sequence[int], default=[]
        Indices of dropdown components to update the value for.
    args : P.args
        Positional arguments to pass to the function used to update
        dropdown choices.
    kwargs : P.kwargs
        Keyword arguments to pass to the function used to update
        dropdown choices.

    Returns
    -------
    gr.Dropdown | tuple[gr.Dropdown,...]
        Updated dropdown component or components.

    Raises
    ------
    ValueError
        If not all provided indices are unique or if an index exceeds
        or is equal to the number of dropdown components.

    """
    if len(value_indices) != len(set(value_indices)):
        err_msg = "Value indices must be unique."
        raise ValueError(err_msg)
    if value_indices and max(value_indices) >= num_components:
        err_msg = (
            "Index of a dropdown component to update the value for exceeds the number"
            " of dropdown components to update."
        )
        raise ValueError(err_msg)
    updated_choices = fn(*args, **kwargs)
    update_args_list: list[UpdateDropdownKwArgs] = [
        {"choices": updated_choices} for _ in range(num_components)
    ]
    for index in value_indices:
        update_args_list[index]["value"] = value

    match update_args_list:
        case [update_args]:
            # NOTE This is a workaround as gradio does not support
            # singleton tuples for components.
            return gr.Dropdown(**update_args)
        case _:
            return tuple(gr.Dropdown(**update_args) for update_args in update_args_list)


def update_cached_songs(
    num_components: int,
    value: DropdownValue = None,
    value_indices: Sequence[int] = [],
) -> gr.Dropdown | tuple[gr.Dropdown, ...]:
    """
    Update the choices of one or more dropdown components to the set of
    currently cached songs.

    Optionally update the default value of one or more of these
    components.

    Parameters
    ----------
    num_components : int
        Number of dropdown components to update.
    value : DropdownValue, optional
        New value for the dropdown components.
    value_indices : Sequence[int], default=[]
        Indices of dropdown components to update the value for.

    Returns
    -------
    gr.Dropdown | tuple[gr.Dropdown,...]
        Updated dropdown component or components.

    """
    return update_dropdowns(get_named_song_dirs, num_components, value, value_indices)


def update_output_audio(
    num_components: int,
    value: DropdownValue = None,
    value_indices: Sequence[int] = [],
) -> gr.Dropdown | tuple[gr.Dropdown, ...]:
    """
    Update the choices of one or more dropdown components to the set of
    currently saved output audio files.

    Optionally update the default value of one or more of these
    components.

    Parameters
    ----------
    num_components : int
        Number of dropdown components to update.
    value : DropdownValue, optional
        New value for dropdown components.
    value_indices : Sequence[int], default=[]
        Indices of dropdown components to update the value for.

    Returns
    -------
    gr.Dropdown | tuple[gr.Dropdown,...]
        Updated dropdown component or components.

    """
    return update_dropdowns(
        get_saved_output_audio,
        num_components,
        value,
        value_indices,
    )


def toggle_visible_component(
    num_components: int,
    visible_index: int,
) -> dict[str, Any] | tuple[dict[str, Any], ...]:
    """
    Reveal a single component from a set of components. All other
    components are hidden.

    Parameters
    ----------
    num_components : int
        Number of components to set visibility for.
    visible_index : int
        Index of the component to reveal.

    Returns
    -------
    dict[str, Any] | tuple[dict[str, Any], ...]
        A single dictionary or a tuple of dictionaries that update the
        visibility of the components.

    Raises
    ------
    ValueError
        If the visible index exceeds or is equal to the number of
        components to set visibility for.

    """
    if visible_index >= num_components:
        err_msg = (
            "Visible index must be less than the number of components to set visibility"
            " for."
        )
        raise ValueError(err_msg)
    update_args_list: list[ComponentVisibilityKwArgs] = [
        {"visible": False, "value": None} for _ in range(num_components)
    ]
    update_args_list[visible_index]["visible"] = True
    match update_args_list:
        case [update_args]:
            return gr.update(**update_args)
        case _:
            return tuple(gr.update(**update_args) for update_args in update_args_list)


def update_song_cover_name(
    effected_vocals_track: str | None = None,
    song_dir: str | None = None,
    model_name: str | None = None,
    update_placeholder: bool = False,
) -> gr.Textbox:
    """
    Update a textbox component so that it displays a suitable name for a
    cover of a given song.

    If the path of an existing song directory is provided, the name of
    the song is inferred from that directory. If the name of a voice
    model is not provided but the path of an existing song directory
    and the path of an effected vocals track in that directory are
    provided, then the voice model is inferred from the effected vocals
    track.


    Parameters
    ----------
    effected_vocals_track : str, optional
        The path to an effected vocals track.
    song_dir : str, optional
        The path to a song directory.
    model_name : str, optional
        The name of a voice model.
    update_placeholder : bool, default=False
        Whether to update the placeholder text instead of the value of
        the textbox component.

    Returns
    -------
    gr.Textbox
        Textbox component with updated value or placeholder text.

    """
    update_args: TextBoxKwArgs = {}
    update_key = "placeholder" if update_placeholder else "value"
    if effected_vocals_track or song_dir or model_name:
        song_cover_name = get_song_cover_name(
            effected_vocals_track,
            song_dir,
            model_name,
        )
        update_args[update_key] = song_cover_name
    else:
        update_args[update_key] = None
    return gr.Textbox(**update_args)