import os

from transformers import AutoTokenizer, AutoModel

from suno_utils.tasks.ditto_v2 import preload_models as preload_ditto_models
from suno_utils.tasks.dac_vae_fixed_25hz import preload_models as preload_vae_models_
from suno_utils.tasks.mert_25 import preload_models as preload_mert_models_
from suno_utils.tasks.musicfm_v3 import preload_models as preload_musicfm_models_

_mmbert_tokenizer = None
_mmbert_encoder = None
_hoot_tokenizer = None
_hoot_encoder = None
_midi_model = None
_ditto_model_loaded = False
_vae_model_loaded = False
_semantic_model_loaded = False


def load_mmbert_tokenizer_and_encoder():
    """Load mmBERT tokenizer and encoder"""
    global _mmbert_tokenizer, _mmbert_encoder
    if _mmbert_tokenizer is None or _mmbert_encoder is None:
        _mmbert_tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/mmBERT-base")
        _mmbert_encoder = AutoModel.from_pretrained("jhu-clsp/mmBERT-base")
    return _mmbert_tokenizer, _mmbert_encoder


def load_hoot_tokenizer_and_encoder():
    from suno_utils.tasks.hoot import load_model_list

    global _hoot_tokenizer, _hoot_encoder
    if _hoot_tokenizer is None or _hoot_encoder is None:
        model_list = load_model_list(
            checkpoint_filepath=os.path.join("/app/suno/models", "hoot_v3.pt"),
            tokenizer_filepath="/app/suno/models/hoot_v3_tokenizer.model",
            n_gpus=1,
        )
        model = model_list[0]["model"]
        model.to("cuda")
        tokenizer = model_list[0]["tokenizer"]
        _hoot_tokenizer = tokenizer
        _hoot_encoder = model
    return _hoot_tokenizer, _hoot_encoder


def load_midi_model():
    """Load MIDI transcription model"""
    from suno_utils.gpt.generation import load_model
    from suno_utils.utils.s3 import download_s3_file_if_needed

    global _midi_model
    if _midi_model is None:
        model_container = load_model(
            ckpt_path=download_s3_file_if_needed(
                "s3://suno-data/m4burns/midi_transcription_v1_sft_new_3_408k_lean.pt"
            ),
            tokenizer_path=download_s3_file_if_needed(
                "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json"
            ),
        )
        _midi_model = model_container["model"]
        _midi_model.to("cuda")
        _midi_model.eval()
    return _midi_model


def load_ditto_model():
    global _ditto_model_loaded
    if not _ditto_model_loaded:
        preload_ditto_models()
        _ditto_model_loaded = True


def load_vae_model():
    global _vae_model_loaded
    if not _vae_model_loaded:
        preload_vae_models_(checkpoint_filepath="s3://suno-data/minz/models/dac_vae_tuned_25hz.pth")
        _vae_model_loaded = True


def preload_semantic_models(semantic_type):
    global _semantic_model_loaded
    if not _semantic_model_loaded:
        if semantic_type.startswith("mert"):
            preload_mert_models_(
                checkpoint_filepath="s3://suno-data/georg/models/semantic/mert_25.pt",
                centroids_filepath="s3://suno-data/minz/models/mert_64residual_kmeans_centroids.npy",
            )
        elif semantic_type == "musicfm_kmeans":
            preload_musicfm_models_(
                checkpoint_filepath="/app2/suno/checkpoints/2025-10-23_23-53-15/last_ckpt_infer.pt",
                centroids_filepath="/home/minz/temp/musicfm3_1024d_nov17_centroids_4000_50.npy",
            )
        elif semantic_type == "musicfm_internal":
            preload_musicfm_models_(
                checkpoint_filepath="/app2/suno/checkpoints/2025-10-23_23-53-15/last_ckpt_infer.pt"
            )
        _semantic_model_loaded = True