import os
import logging
import functools
import torch

logger = logging.getLogger(__name__)

try:
    from fairseq import checkpoint_utils
    FAIRSEQ_AVAILABLE = True
except Exception:
    FAIRSEQ_AVAILABLE = False


def _patch_torch_load():
    """Patch torch.load to default weights_only=False for fairseq compatibility (PyTorch 2.6+)."""
    _original = torch.load

    @functools.wraps(_original)
    def _patched(*args, **kwargs):
        kwargs.setdefault("weights_only", False)
        return _original(*args, **kwargs)

    return _original, _patched


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 sid.split(".")[0] in f
        ),
        "",
    )


def load_hubert(config):
    if FAIRSEQ_AVAILABLE:
        _original, _patched = _patch_torch_load()
        torch.load = _patched
        try:
            models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
                ["assets/hubert/hubert_base.pt"],
                suffix="",
            )
        finally:
            torch.load = _original
        hubert_model = models[0]
        hubert_model = hubert_model.to(config.device)
        if config.is_half:
            hubert_model = hubert_model.half()
        else:
            hubert_model = hubert_model.float()
        return hubert_model.eval()

    try:
        import torchaudio

        class HubertWrapper:
            def __init__(self, model):
                self.model = model
                self.final_proj = getattr(model, "final_proj", torch.nn.Identity())

            def extract_features(self, source, padding_mask=None, output_layer=None):
                feats, _ = self.model.extract_features(source)
                if output_layer is None:
                    idx = -1
                else:
                    idx = min(output_layer - 1, len(feats) - 1)
                return (feats[idx], None)

            def to(self, device):
                self.model = self.model.to(device)
                return self

            def half(self):
                self.model = self.model.half()
                return self

            def float(self):
                self.model = self.model.float()
                return self

            def eval(self):
                self.model.eval()
                return self

        model = torchaudio.pipelines.HUBERT_BASE.get_model()
        hubert_model = HubertWrapper(model).to(config.device)
        if config.is_half:
            hubert_model = hubert_model.half()
        else:
            hubert_model = hubert_model.float()
        return hubert_model.eval()
    except Exception as e:
        raise RuntimeError(
            "HuBERT 模型加载失败，请检查 fairseq 和 torchaudio 是否已安装"
        ) from e