# model_manager.py

import torch
from TTS.api import TTS
from voice.clearvoice import ClearVoice
import os

os.environ["COQUI_TOS_AGREED"] = "1"

_model_cache = {}


def load_model():
    global _model_cache
    """Load and cache all necessary models."""
    device = "cuda" if torch.cuda.is_available() else "cpu"

    if "xtts_v2" not in _model_cache:
        _model_cache["xtts_v2"] = TTS(
            "tts_models/multilingual/multi-dataset/xtts_v2"
        ).to(device)

    if "FRCRN_SE_16K" not in _model_cache:
        _model_cache["FRCRN_SE_16K"] = ClearVoice(
            task="speech_enhancement", model_names=["FRCRN_SE_16K"]
        )

    if "MossFormer2_SE_48K" not in _model_cache:
        _model_cache["MossFormer2_SE_48K"] = ClearVoice(
            task="speech_enhancement", model_names=["MossFormer2_SE_48K"]
        )

    if "MossFormer2_SS_16K" not in _model_cache:
        _model_cache["MossFormer2_SS_16K"] = ClearVoice(
            task="speech_separation", model_names=["MossFormer2_SS_16K"]
        )

    if "AV_MossFormer2_TSE_16K" not in _model_cache:
        _model_cache["AV_MossFormer2_TSE_16K"] = ClearVoice(
            task="target_speaker_extraction", model_names=["AV_MossFormer2_TSE_16K"]
        )

    if "MossFormer2_SR_48K" not in _model_cache:
        _model_cache["MossFormer2_SR_48K"] = ClearVoice(
            task="speech_super_resolution", model_names=["MossFormer2_SR_48K"]
        )

    return _model_cache


def unload_model():
    """Clear cached models and free up memory."""
    _model_cache.clear()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def predict(task: str, **kwargs):
    """
    Central prediction dispatcher based on task type.

    Args:
        task (str): Task type key.
        kwargs: Task-specific arguments.

    Returns:
        Result of the model inference.
    """
    if not _model_cache:
        load_model()

    if task == "clone_voice":
        model = _model_cache["xtts_v2"]
        return model.tts_to_file(
            text=kwargs["text"],
            speaker_wav=kwargs["speaker_wav"],
            language=kwargs.get("language", "en"),
            file_path=kwargs["file_path"],
        )

    elif task == "speech_enhancement_16k":
        model = _model_cache["FRCRN_SE_16K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=kwargs.get("online_write", False),
        )

    elif task == "speech_enhancement_48k":
        model = _model_cache["MossFormer2_SE_48K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=kwargs.get("online_write", False),
            output_path=kwargs.get("output_path"),
        )

    elif task == "speech_separation":
        model = _model_cache["MossFormer2_SS_16K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=kwargs.get("online_write", False),
        )

    elif task == "target_speaker_extraction":
        model = _model_cache["AV_MossFormer2_TSE_16K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=True,
            output_path=kwargs["output_path"],
        )

    elif task == "speech_super_resolution":
        model = _model_cache["MossFormer2_SR_48K"]
        return model(
            input_path=kwargs["input_path"],
            online_write=kwargs.get("online_write", False),
        )

    else:
        raise ValueError(f"Unknown task type: {task}")


# from model_manager import predict

# # Assume you have speaker and output audio in bytes
# with open("speaker.wav", "rb") as f:
#     speaker_bytes = f.read()

# text_to_speak = "Hello, this is a cloned voice."
# output_path = "cloned_voice.wav"

# predict(
#     task="clone_voice",
#     text=text_to_speak,
#     speaker_wav=speaker_bytes,
#     language="en",
#     file_path=output_path
# )


# from model_manager import predict

# with open("noisy_16k.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="speech_enhancement_16k",
#     input_audio=audio_bytes
# )


# from model_manager import predict

# with open("noisy_48k.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="speech_enhancement_48k",
#     input_audio=audio_bytes,
#     output_path="enhanced_48k.wav"
# )


# from model_manager import predict

# with open("overlapping_speakers.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="speech_separation",
#     input_audio=audio_bytes
# )


# from model_manager import predict

# with open("mixture_with_target.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="target_speaker_extraction",
#     input_audio=audio_bytes,
#     output_path="target_extracted.wav"
# )


# from model_manager import predict

# with open("low_quality_48k.wav", "rb") as f:
#     audio_bytes = f.read()

# predict(
#     task="speech_super_resolution",
#     input_audio=audio_bytes
# )