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

TEYVAT = ["Mondstadt", "Liyue", "Inazuma", "Sumeru", "Fontaine", "Natlan"]
WEIGHTS_PATH = (
    snapshot_download("Genius-Society/hoyoMusic", cache_dir="./__pycache__")
    + "/weights.pth"
)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TEMP_DIR = "./__pycache__/tmp"
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
PATCH_SAMPLING_BATCH_SIZE = 0  # Batchsize for patch during training, 0=full context
SHARE_WEIGHTS = False  # Whether to share weights between the encoder and decoder


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/Genius-Society/piano_trans/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 = os.getenv("mscore")