from typing import Callable
import torch
import torchaudio
import os
from typing import Optional, Literal, List, Dict, Callable
import torchaudio.functional as F
import torchaudio.transforms as T
from vox_bench.utils import _filter_kwargs, get_label, PRIM_AUD_FUNC_REGISTRY, load_audio

def default_loader(
    audio_path: str,
    cls: str,
    loader_name: str,
    prob_dist: bool = False,
    **kwargs
):
    waveform, sample_rate = load_audio(
        audio_path,
        max_secs=kwargs.get("max_secs", None)
    )

    if loader_name not in LOADER_REGISTRY:
        raise ValueError(f"{loader_name} not registered!")

    feature_func = PRIM_LOADER_REGISTRY[loader_name]

    kwargs["sample_rate"] = sample_rate
    filtered_kwargs = _filter_kwargs(feature_func, kwargs)

    features = feature_func(waveform, **filtered_kwargs)
    label = get_label(cls, prob_dist)

    return {
        "X": {"x1": features},
        "Y": {"y1": label}
    }

LOADER_REGISTRY : Dict[str, Callable] = {
    "default": default_loader
}

def register_loader(name : str):
    def decorator(func: Callable):
        LOADER_REGISTRY[name] = func
        return func
    return decorator
    
def get_loader_func(name: str) -> Callable:
    return LOADER_REGISTRY[name]
