import os
import subprocess
import sys

import spaces

ZeroGPU = bool(os.getenv("SPACES_ZERO_GPU"))

# Install latest git main code
ZeroGPU and subprocess.run((
    sys.executable, "-m", "pip", "install", "--upgrade",
    "git+https://github.com/BrokenSource/BrokenSource@main",
    "git+https://github.com/BrokenSource/ShaderFlow@main",
    "git+https://github.com/BrokenSource/DepthFlow@main",
))

# -----------------------------------------------|

from DepthFlow.Webui import ESTIMATORS, UPSCALERS

from Broken.Externals.Depthmap import DepthAnythingV2

# Remove legacy upscalers
UPSCALERS.pop("Real-ESRGAN")
UPSCALERS.pop("Waifu2x")

# Remove unsuitable models
ESTIMATORS.clear()
ESTIMATORS.update({
    # Base and Large models are CC BY-NC 4.0, better safe than sorry
    "DepthAnything2": DepthAnythingV2(model=DepthAnythingV2.Model.Small),
})

# -----------------------------------------------|
# Create Vulkan and EGL ICD loaders files

from pathlib import Path

if ZeroGPU:
    egl_vendor = Path("/usr/share/glvnd/egl_vendor.d/10_nvidia.json")
    egl_vendor.parent.mkdir(parents=True, exist_ok=True)
    egl_vendor.write_text('{"file_format_version":"1.0.0","ICD":{"library_path":"libEGL_nvidia.so.0"}}')

    vulkan_icd = Path("/usr/share/vulkan/icd.d/nvidia_icd.json")
    vulkan_icd.parent.mkdir(parents=True, exist_ok=True)
    vulkan_icd.write_text('{"file_format_version":"1.0.0","ICD":{"library_path":"libGLX_nvidia.so.0","api_version":"1.3"}}')

# -----------------------------------------------|
from DepthFlow.Webui import DepthGradio

DepthGradio().launch(
    workers=4,
    nvenc=True,
    turbo=True,
    pwa=True,
)