import os
import sys
import time
import torch
import requests
import subprocess
from tqdm import tqdm
from modelscope import snapshot_download

TEYVAT = {
    "蒙德 Mondstadt": "Mondstadt",
    "璃月 Liyue": "Liyue",
    "稻妻 Inazuma": "Inazuma",
    "须弥 Sumeru": "Sumeru",
    "枫丹 Fontaine": "Fontaine",
}
WEIGHTS_PATH = (
    snapshot_download("MuGemSt/hoyoMusic", cache_dir="./__pycache__")
    + "/weights.pth"
)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TEMP_DIR = "./flagged"
PATCH_LENGTH = 128  # Patch Length
PATCH_SIZE = 32  # Patch Size
PATCH_NUM_LAYERS = 9  # Number of layers in the encoder
CHAR_NUM_LAYERS = 3  # Number of layers in the decoder
# Batch size for patch during training, 0 for full context
PATCH_SAMPLING_BATCH_SIZE = 0
# Whether to share weights between the encoder and decoder
SHARE_WEIGHTS = False


def download(filename: str, url: str):
    try:
        response = requests.get(url, stream=True)
        total_size = int(response.headers.get("content-length", 0))
        chunk_size = 1024
        with open(filename, "wb") as file, tqdm(
            desc=f"Downloading {filename} from '{url}'...",
            total=total_size,
            unit="B",
            unit_scale=True,
            unit_divisor=1024,
        ) as bar:
            for data in response.iter_content(chunk_size=chunk_size):
                size = file.write(data)
                bar.update(size)

    except Exception as e:
        print(f"Error: {e}, retrying...")
        time.sleep(10)
        download(filename, url)


if sys.platform.startswith("linux"):
    apkname = "MuseScore.AppImage"
    extra_dir = "squashfs-root"
    download(
        filename=apkname,
        url="https://www.modelscope.cn/studio/MuGemSt/piano_transcription/resolve/master/MuseScore.AppImage",
    )
    if not os.path.exists(extra_dir):
        subprocess.run(["chmod", "+x", f"./{apkname}"])
        subprocess.run([f"./{apkname}", "--appimage-extract"])

    MSCORE = f"./{extra_dir}/AppRun"
    os.environ["QT_QPA_PLATFORM"] = "offscreen"

else:
    MSCORE = "D:/Program Files/MuseScore 3/bin/MuseScore3.exe"