# forked from maimaiDX-Api
import json
import xml.dom.minidom as minidom
from pathlib import Path

ONLY_REMOVED = True


def makeMusicDBJson():
    """
    从 HDD 的文件来生成 music_db.json
    推荐的是如果要国服用 那就用国际服的文件来生成
    免得国服每次更新还要重新生成太麻烦
    """
    # 记得改
    A000_DIR = Path(
        "C:/MaimaiDX/SDEZ-1.56-B/Standard/Package/Sinmai_Data/StreamingAssets/A000"
    )
    OPTION_DIR = Path("C:/MaimaiDX/SDGA-1.50-G/NoMovieData/StreamingAssets")

    music_db: list[dict[str, str | int | list[dict[str, str | int]]]] = []
    DEST_PATH = Path("./musicDB.json")

    dup_count = 0
    music_ids = set()

    music_folders = [f for f in (A000_DIR / "music").iterdir() if f.is_dir()]
    for option_dir in OPTION_DIR.iterdir():
        # only removed ones
        if ONLY_REMOVED and option_dir.name != "A100":
            continue

        if (option_dir / "music").exists():
            music_folders.extend(
                [f for f in (option_dir / "music").iterdir() if f.is_dir()]
            )

    for folder in music_folders:
        xml_path = folder / "Music.xml"
        if xml_path.exists():
            xml = minidom.parse(xml_path.as_posix())
            data = xml.getElementsByTagName("MusicData")[0]
            music_id = int(
                data.getElementsByTagName("name")[0]
                .getElementsByTagName("id")[0]
                .firstChild.data
            )
            music_name = (
                data.getElementsByTagName("name")[0]
                .getElementsByTagName("str")[0]
                .firstChild.data
            )
            music_version = (
                data.getElementsByTagName("AddVersion")[0]
                .getElementsByTagName("id")[0]
                .firstChild.data
            )

            def handle_note(note: minidom.Element):
                if (
                    "false"
                    == note.getElementsByTagName("isEnable")
                    .pop()
                    .firstChild.data.lower()
                ):
                    return

                if music_id >= 100000:
                    level = 5
                else:
                    level = int(
                        note.getElementsByTagName("file")
                        .pop()
                        .getElementsByTagName("path")
                        .pop()
                        .firstChild.data[-5]
                    )
                difficulty_int = (
                    note.getElementsByTagName("level").pop().firstChild.data
                )
                difficulty_dec = (
                    note.getElementsByTagName("levelDecimal").pop().firstChild.data
                )
                difficulty = f"{difficulty_int}.{difficulty_dec}"
                return level, difficulty

            music_levels = [
                {"level": level, "difficulty": difficulty}
                for level, difficulty in filter(
                    lambda d: d is not None,
                    (
                        handle_note(note)
                        for note in data.getElementsByTagName("notesData")[
                            0
                        ].getElementsByTagName("Notes")
                    ),
                )
            ]

            if music_id not in music_ids:
                music_ids.add(music_id)
                music_db.append(
                    {
                        "id": music_id,
                        "name": music_name,
                        "version": int(music_version),
                        "levels": music_levels,
                    }
                )
            else:
                # e.g. SDEZ-only song
                dup_count += 1

    print(f"Found {len(music_db)} music data")
    print(f"Found {dup_count} duplications")

    music_db = sorted(
        music_db,
        key=lambda m: m["id"],
    )

    with open(DEST_PATH, "w", encoding="utf-8") as f:
        json.dump(music_db, f, ensure_ascii=False, indent=4)


if __name__ == "__main__":
    makeMusicDBJson()
    print("Done.")