import huggingface_hub
import torch
from typing import Tuple

from audclas.tortoise_audio_classifier import TortoiseAudioClassifier
from shttst.audio.utils import save_audio
from shttst.models import ShmartSileroVad, ShmartUltimateVocalRemover

initialized_models= {}

def initialize_model(type):
    model = None
    
    if type == 'classifier':
        model = TortoiseAudioClassifier()
    elif type == 'denoiser':
        model = ShmartUltimateVocalRemover()
    elif type == 'titanet':
        from nemo.collections.asr.models import EncDecSpeakerLabelModel
        model = EncDecSpeakerLabelModel.from_pretrained("nvidia/speakerverification_en_titanet_large", map_location='cpu')
        model.cuda().eval()
    elif type == 'resemblyzer':
        from resemblyzer import VoiceEncoder
        model = VoiceEncoder('cuda', verbose=False)
        model.eval()
    elif type == 'vad':
        model = ShmartSileroVad()
    elif type == 'whisper':
        from faster_whisper import WhisperModel
        model_path = huggingface_hub.snapshot_download("shmart/shmisper-medium-PL")
        model = WhisperModel(model_path, device="cuda", compute_type="float16")
    
    initialized_models[type] = model


def get_model(type):
    if type not in initialized_models:
        initialize_model(type)

    return initialized_models[type]

def classify_audio_file(file_path: str) -> str:
    model: TortoiseAudioClassifier = get_model('classifier')
    return model(file_path)

def split_by_vad(file_path: str, max_duration: int, min_silence: int) -> str:
    model: ShmartSileroVad = get_model('vad')
    return model.vad_process(file_path, min_silence_duration=min_silence, desired_max_speech_duration=max_duration)

def clean_audio_noise(file_path: str) -> str:
    vr_model: ShmartUltimateVocalRemover = get_model('denoiser')
    vocals, instruments = vr_model(file_path, tta=True)
    save_audio(file_path, vocals)

def get_custom_embedding(file_path, type='titanet'):
    if type == 'titanet':
        from nemo.collections.asr.models import EncDecSpeakerLabelModel
        titanet_model: EncDecSpeakerLabelModel = get_model(type)
        return titanet_model.get_embedding(file_path).squeeze(0).cpu()
    elif type == 'resemblyzer':
        from resemblyzer import VoiceEncoder, preprocess_wav
        wav = preprocess_wav(file_path)
        resemblyzer_model: VoiceEncoder = get_model(type)
        return torch.from_numpy(resemblyzer_model.embed_utterance(wav))

def get_transcription(audio_path, options={
    'language': "pl",
    'beam_size': 5,
    'without_timestamps': True,
    'suppress_tokens': [],
    'log_prob_threshold': None,
    'no_speech_threshold': 0.06
}) -> Tuple[str, float]:
    from faster_whisper import WhisperModel
    model: WhisperModel = get_model('whisper')
    result, info = model.transcribe(audio_path, **options)
    try:
      return ' '.join([r.text for r in result]), info.duration
    except:
      return '', -1
    
if __name__ == '__main__':
    get_custom_embedding('\\content\\cyber\\2077\\wavs\\55.wav')