import os
import torch

from fairseq import checkpoint_utils


def get_index_path_from_model(sid):
    return next(
        (
            f
            for f in [
                os.path.join(root, name)
                for root, _, files in os.walk(os.getenv("index_root"), topdown=False)
                for name in files
                if name.endswith(".index") and "trained" not in name
            ]
            if str(sid).split(".")[0] in f
        ),
        "",
    )


def load_hubert(config, hubert_path: str):
    # PyTorch 2.6+ changed weights_only default to True, which breaks fairseq checkpoints
    # Monkey-patch torch.load to use weights_only=False for fairseq
    original_torch_load = torch.load

    def patched_torch_load(f, map_location=None, *args, **kwargs):
        kwargs.setdefault('weights_only', False)
        return original_torch_load(f, map_location=map_location, *args, **kwargs)

    torch.load = patched_torch_load
    try:
        models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
            [hubert_path],
            suffix="",
        )
    finally:
        torch.load = original_torch_load

    hubert_model = models[0]
    hubert_model = hubert_model.to(config.device)
    hubert_model = hubert_model.half() if config.is_half else hubert_model.float()
    return hubert_model.eval()