from __future__ import annotations

import os
import site
import sys
from pathlib import Path

os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "False")
os.environ.setdefault("HF_HOME", "/home/user/.cache/huggingface")
os.environ.setdefault("TRANSFORMERS_CACHE", "/home/user/.cache/huggingface/transformers")
os.environ.setdefault("PYTHONUNBUFFERED", "1")
os.environ.setdefault("XFORMERS_DISABLED", "1")
os.environ.setdefault("ATTN_BACKEND", "sdpa")
os.environ.setdefault("SPARSE_ATTN_BACKEND", "flash_attn")
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "8.6")


def _prepare_cuda_runtime_path() -> None:
    if os.environ.get("ISCENE_LD_REEXEC") == "1":
        return

    lib_dirs = []
    for site_dir in site.getsitepackages():
        root = Path(site_dir)
        lib_dirs.extend(root.glob("nvidia/*/lib"))
        lib_dirs.append(root / "torch" / "lib")

    compat_dir = Path("/tmp/iscene_cuda_compat")
    compat_dir.mkdir(parents=True, exist_ok=True)
    for lib_dir in lib_dirs:
        for cudart in lib_dir.glob("libcudart.so.12*"):
            target = compat_dir / "libcudart.so.13"
            if not target.exists():
                target.symlink_to(cudart)

    paths = [str(compat_dir)] + [str(path) for path in lib_dirs if path.exists()]
    os.environ["LD_LIBRARY_PATH"] = ":".join(paths + [os.environ.get("LD_LIBRARY_PATH", "")])
    os.environ["ISCENE_LD_REEXEC"] = "1"
    os.execvpe(sys.executable, [sys.executable, *sys.argv], os.environ)


_prepare_cuda_runtime_path()

import spaces
import gradio as gr
import torch

import interactive_demo


def _configure_runtime_device() -> None:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16 if device == "cuda" else torch.float32
    if interactive_demo.DEVICE != device or interactive_demo.DTYPE != dtype:
        interactive_demo._sam_cache.clear()
    interactive_demo.DEVICE = device
    interactive_demo.DTYPE = dtype


_run_segmentation = interactive_demo.run_segmentation
_run_gaussian_preview = interactive_demo.run_gaussian_preview
_run_glb_export = interactive_demo.run_glb_export


@spaces.GPU(duration=120)
def run_segmentation(
    image_prompts,
    model_choice,
    polygon_refinement,
    mask_threshold,
    request: gr.Request,
):
    _configure_runtime_device()
    return _run_segmentation(
        image_prompts,
        model_choice,
        polygon_refinement,
        mask_threshold,
        request,
    )


@spaces.GPU(duration=180)
def run_gaussian_preview(
    image_prompts,
    mask_path,
    seed,
    simplify,
    output_dir_text,
    request: gr.Request,
):
    _configure_runtime_device()
    return _run_gaussian_preview(
        image_prompts,
        mask_path,
        seed,
        simplify,
        output_dir_text,
        request,
    )


@spaces.GPU(duration=240)
def run_glb_export(state, simplify):
    _configure_runtime_device()
    yield from _run_glb_export(state, simplify)


interactive_demo.run_segmentation = run_segmentation
interactive_demo.run_gaussian_preview = run_gaussian_preview
interactive_demo.run_glb_export = run_glb_export
interactive_demo.MODEL_ID = os.environ.get("ISCENE_MODEL", interactive_demo.DEFAULT_MODEL)
interactive_demo.BASE_MODEL_ID = os.environ.get("ISCENE_BASE_MODEL") or None
interactive_demo.DEFAULT_OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
interactive_demo.UPLOAD_ROOT.mkdir(parents=True, exist_ok=True)

demo = interactive_demo.build_demo()
demo.queue()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)