#!/usr/bin/env python3
"""
dav_loader.py — import MiniMaxMusic3DAV without executing SimpleTuner's
package __init__ chain.

`simpletuner/__init__.py` pulls in the whole training stack, which monkeypatches
diffusers attention on import. Against a pinned diffusers commit that patching
can fail (e.g. TemplatedRingAnythingAttention not present), taking the DAV
class down with it.

vocoder.py itself imports only math / torch / diffusers, so it loads fine
standalone. This registers stub parent packages and execs just that file.

    from dav_loader import get_dav_class
    MiniMaxMusic3DAV = get_dav_class(Path("~/gpu_work/m3").expanduser())
"""

from __future__ import annotations

import importlib.util
import sys
import types
from pathlib import Path

_CACHE = {}

_PKG_CHAIN = [
    "simpletuner",
    "simpletuner.helpers",
    "simpletuner.helpers.models",
    "simpletuner.helpers.models.minimaxmusic",
]
_MODNAME = "simpletuner.helpers.models.minimaxmusic.vocoder"


def get_dav_class(workdir: Path | str, simpletuner_dir: Path | str | None = None):
    """Return the MiniMaxMusic3DAV class, loaded in isolation."""
    key = str(workdir)
    if key in _CACHE:
        return _CACHE[key]

    root = Path(simpletuner_dir) if simpletuner_dir else Path(workdir) / "SimpleTuner"
    vpath = root / "simpletuner" / "helpers" / "models" / "minimaxmusic" / "vocoder.py"
    if not vpath.is_file():
        raise FileNotFoundError(
            f"vocoder.py not found at {vpath}\n"
            f"Re-fetch it with:\n"
            f"  git -C {root} fetch --depth 1 origin refs/pull/3074/head && "
            f"git -C {root} checkout -q FETCH_HEAD"
        )

    # Stub the parent packages so the dotted module name resolves without
    # any real __init__.py being executed.
    for name in _PKG_CHAIN:
        if name not in sys.modules:
            stub = types.ModuleType(name)
            stub.__path__ = []          # mark as a package
            sys.modules[name] = stub

    spec = importlib.util.spec_from_file_location(_MODNAME, str(vpath))
    if spec is None or spec.loader is None:
        raise ImportError(f"could not build a spec for {vpath}")
    mod = importlib.util.module_from_spec(spec)
    sys.modules[_MODNAME] = mod
    spec.loader.exec_module(mod)

    cls = getattr(mod, "MiniMaxMusic3DAV", None)
    if cls is None:
        raise ImportError(f"MiniMaxMusic3DAV not found in {vpath}")

    _CACHE[key] = cls
    return cls


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser()
    ap.add_argument("--workdir", type=Path, required=True)
    a = ap.parse_args()

    cls = get_dav_class(a.workdir.expanduser().resolve())
    print(f"loaded {cls.__name__} OK")

    ckpt = a.workdir.expanduser().resolve() / "minimax_music3" / "dav.pth"
    if ckpt.exists():
        dav = cls.from_original_dav(str(ckpt))
        n = sum(p.numel() for p in dav.parameters())
        print(f"weights OK | {n/1e6:.1f}M params | hop_length={dav.hop_length} "
              f"| latent_channels={dav.config.latent_channels}")
    else:
        print(f"(checkpoint not at {ckpt}; class import verified only)")