from . import modules
from pathlib import Path
from . import scheduler
from .interface import Interface
from .modules.transformer import VampNet


__version__ = "0.0.1"

ROOT = Path(__file__).parent.parent
MODELS_DIR = ROOT / "models" / "vampnet"

from huggingface_hub import hf_hub_download, HfFileSystem
DEFAULT_HF_MODEL_REPO_DIR = ROOT / "DEFAULT_HF_MODEL_REPO"
DEFAULT_HF_MODEL_REPO = DEFAULT_HF_MODEL_REPO_DIR.read_text().strip()
# DEFAULT_HF_MODEL_REPO = "hugggof/vampnet"
# The weights' revision, which every download below names, so that the models cannot change
# under the Space
REVISION = "69e9c7dba72b54a45f93122ef91aa738d2e803ba"
FS = HfFileSystem()

def download_codec():
    # from dac.model.dac import DAC
    from lac.model.lac import LAC as DAC
    repo_id = DEFAULT_HF_MODEL_REPO
    filename = "codec.pth"
    codec_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        subfolder=None, 
        local_dir=MODELS_DIR,
        revision=REVISION,
    )
    return codec_path
    

# hf_hub_download writes each checkpoint beside its final path and moves it there only once its
# size is confirmed, so an interrupted download cannot leave a truncated file that a later start
# would take for a complete one
def download_default():
    filenames = ["coarse.pth", "c2f.pth", "wavebeat.pth"]
    paths = [
        hf_hub_download(repo_id=DEFAULT_HF_MODEL_REPO, filename=filename, local_dir=MODELS_DIR,
                        revision=REVISION)
        for filename in filenames
    ]

    # load the models
    return paths[0], paths[1]


def download_finetuned(name, repo_id=DEFAULT_HF_MODEL_REPO):
    filenames = ["coarse.pth", "c2f.pth"]
    paths = [
        hf_hub_download(repo_id=repo_id, filename=f"loras/{name}/{filename}", local_dir=MODELS_DIR,
                        revision=REVISION)
        for filename in filenames
    ]

    # load the models
    return paths[0], paths[1]
    
def list_finetuned(repo_id=DEFAULT_HF_MODEL_REPO):
    diritems = FS.listdir(f"{repo_id}@{REVISION}/loras")
    # iterate through all the names
    valid_diritems = []
    for item in diritems:
        model_file_items = FS.listdir(item["name"])
        item_names = [item["name"].split("/")[-1] for item in model_file_items]
        # check that theres a "c2f.pth" and "coarse.pth" in the items
        c2f_exists = "c2f.pth" in item_names
        coarse_exists = "coarse.pth" in item_names
        if c2f_exists and coarse_exists:
            valid_diritems.append(item)

    # get the names of the valid items
    names = [item["name"].split("/")[-1] for item in valid_diritems]
    return names