# main.py

import os
from pathlib import Path
from gui import get_gui
from utils import download_manager
from mdx_core import MDX
import json

MDX_DOWNLOAD_LINK = "https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/"
UVR_MODELS = [
    "UVR-MDX-NET-Voc_FT.onnx",
    "UVR_MDXNET_KARA_2.onnx",
    "Reverb_HQ_By_FoxJoy.onnx",
    "UVR-MDX-NET-Inst_HQ_4.onnx",
]

BASE_DIR = Path(__file__).resolve().parent
mdxnet_models_dir = BASE_DIR / "mdx_models"

# Ensure mdx_models is a directory
if mdxnet_models_dir.exists() and not mdxnet_models_dir.is_dir():
    mdxnet_models_dir.unlink()
mdxnet_models_dir.mkdir(parents=True, exist_ok=True)


# 🔧 Auto-generate fallback data.json if needed
def generate_data_json(mdxnet_models_dir):
    presets = {
        "UVR-MDX-NET-Voc_FT.onnx": {
            "mdx_dim_f_set": 2048,
            "mdx_dim_t_set": 3,
            "mdx_n_fft_scale_set": 6144,
            "primary_stem": "Vocals",
            "compensate": 1.035
        },
        "UVR_MDXNET_KARA_2.onnx": {
            "mdx_dim_f_set": 1024,
            "mdx_dim_t_set": 3,
            "mdx_n_fft_scale_set": 4096,
            "primary_stem": "Main",
            "compensate": 1.035
        },
        "UVR-MDX-NET-Inst_HQ_4.onnx": {
            "mdx_dim_f_set": 2048,
            "mdx_dim_t_set": 3,
            "mdx_n_fft_scale_set": 6144,
            "primary_stem": "Instrumental",
            "compensate": 1.0
        },
        "Reverb_HQ_By_FoxJoy.onnx": {
            "mdx_dim_f_set": 2048,
            "mdx_dim_t_set": 3,
            "mdx_n_fft_scale_set": 6144,
            "primary_stem": "Vocals",
            "compensate": 1.035
        },
    }

    data = {}
    for filename, params in presets.items():
        path = mdxnet_models_dir / filename
        if path.exists():
            model_hash = MDX.get_hash(str(path))
            data[model_hash] = params

    with open(mdxnet_models_dir / "data.json", "w") as f:
        json.dump(data, f, indent=2)
    print("✅ Auto-generated data.json with", len(data), "entries.")


# Check for data.json and create if missing
json_path = mdxnet_models_dir / "data.json"
print("🔍 Checking for data.json at:", json_path)
print("📁 mdx_models contents:", list(mdxnet_models_dir.iterdir()))
print("📂 Current working directory:", os.getcwd())

if not json_path.exists():
    print("⚠️ data.json missing, generating fallback.")
    generate_data_json(mdxnet_models_dir)
else:
    print("✅ data.json found.")


def download_models():
    for model in UVR_MODELS:
        url = os.path.join(MDX_DOWNLOAD_LINK, model)
        download_manager(url, str(mdxnet_models_dir))


if __name__ == "__main__":
    download_models()
    theme = "NoCrypt/miku"
    app = get_gui(theme)
    app.queue(default_concurrency_limit=40)
    app.launch(
        max_threads=40,
        share=False,
        show_error=True,
        quiet=False,
        debug=False,
    )