#!/usr/bin/env python

"""
Scan the models directory and print out a new models.yaml
"""

import os
import sys
import argparse

from pathlib import Path
from omegaconf import OmegaConf


def main():
    parser = argparse.ArgumentParser(description="Model directory scanner")
    parser.add_argument("models_directory")
    parser.add_argument(
        "--all-models",
        default=False,
        action="store_true",
        help="If true, then generates stanzas for all models; otherwise just diffusers",
    )

    args = parser.parse_args()
    directory = args.models_directory

    conf = OmegaConf.create()
    conf["_version"] = "3.0.0"

    for root, dirs, files in os.walk(directory):
        parents = root.split("/")
        subpaths = parents[parents.index("models") + 1 :]
        if len(subpaths) < 2:
            continue
        base, model_type, *_ = subpaths

        if args.all_models or model_type == "diffusers":
            for d in dirs:
                conf[f"{base}/{model_type}/{d}"] = dict(
                    path=os.path.join(root, d),
                    description=f"{model_type} model {d}",
                    format="folder",
                    base=base,
                )

            for f in files:
                basename = Path(f).stem
                format = Path(f).suffix[1:]
                conf[f"{base}/{model_type}/{basename}"] = dict(
                    path=os.path.join(root, f),
                    description=f"{model_type} model {basename}",
                    format=format,
                    base=base,
                )

    OmegaConf.save(config=dict(sorted(conf.items())), f=sys.stdout)


if __name__ == "__main__":
    main()