from __future__ import annotations

import ctypes
import importlib
import os
import random
import shutil
import subprocess
import sys
import tempfile
import uuid
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import AutoencoderKL, DDPMScheduler, LCMScheduler, UNet2DConditionModel
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation

# Keep the Space image build simple: fetch a pinned MV-Adapter source snapshot
# at runtime instead of asking pip to resolve the entire upstream package and
# its dependency tree during the Hugging Face image build.
MV_ADAPTER_COMMIT = "4277e0018232bac82bb2c103caf0893cedb711be"
MV_ADAPTER_ROOT = Path("/tmp/basilisc-mv-adapter")

def _ensure_mvadapter_source():
    pkg = MV_ADAPTER_ROOT / "mvadapter"
    if not pkg.is_dir():
        import urllib.request
        import zipfile
        archive = Path("/tmp/basilisc-mv-adapter.zip")
        url = f"https://codeload.github.com/huanngzh/MV-Adapter/zip/{MV_ADAPTER_COMMIT}"
        print(f"Fetching pinned MV-Adapter source: {MV_ADAPTER_COMMIT}", flush=True)
        urllib.request.urlretrieve(url, archive)
        extract_root = Path("/tmp/basilisc-mv-adapter-extract")
        shutil.rmtree(extract_root, ignore_errors=True)
        extract_root.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(archive, "r") as zf:
            zf.extractall(extract_root)
        children = [x for x in extract_root.iterdir() if x.is_dir()]
        if len(children) != 1 or not (children[0] / "mvadapter").is_dir():
            raise RuntimeError("Downloaded MV-Adapter archive had an unexpected layout")
        shutil.rmtree(MV_ADAPTER_ROOT, ignore_errors=True)
        shutil.move(str(children[0]), str(MV_ADAPTER_ROOT))
    if str(MV_ADAPTER_ROOT) not in sys.path:
        sys.path.insert(0, str(MV_ADAPTER_ROOT))

_ensure_mvadapter_source()

from mvadapter.models.attention_processor import DecoupledMVRowColSelfAttnProcessor2_0
from mvadapter.pipelines.pipeline_mvadapter_i2mv_sdxl import MVAdapterI2MVSDXLPipeline
from mvadapter.schedulers.scheduling_shift_snr import ShiftSNRScheduler

DEVICE = "cuda"
DTYPE = torch.float16
NUM_VIEWS = 6
HEIGHT = 768
WIDTH = 768
NEGATIVE = "watermark, ugly, deformed, noisy, blurry, low contrast"
ROOT = Path(__file__).resolve().parent
CHECKPOINTS = ROOT / "checkpoints"
OUTPUTS = ROOT / "outputs"
CHECKPOINTS.mkdir(exist_ok=True)
OUTPUTS.mkdir(exist_ok=True)


def _prepare_pipeline():
    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix")
    pipe = MVAdapterI2MVSDXLPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0", vae=vae
    )
    pipe.scheduler = ShiftSNRScheduler.from_scheduler(
        pipe.scheduler, shift_mode="interpolated", shift_scale=8.0, scheduler_class=None
    )
    pipe.init_custom_adapter(
        num_views=NUM_VIEWS,
        self_attn_processor=DecoupledMVRowColSelfAttnProcessor2_0,
    )
    pipe.load_custom_adapter(
        "huanngzh/mv-adapter", weight_name="mvadapter_ig2mv_sdxl.safetensors"
    )
    pipe.to(device=DEVICE, dtype=DTYPE)
    pipe.cond_encoder.to(device=DEVICE, dtype=DTYPE)
    pipe.enable_vae_slicing()
    return pipe


# ZeroGPU emulation allows model placement on CUDA during startup. This is the
# recommended pattern because the real GPU is attached only while @spaces.GPU
# functions execute.
PIPE = _prepare_pipeline()
BIREFNET = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
).to(DEVICE)
BG_TRANSFORM = transforms.Compose(
    [
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)


def _ensure_checkpoints():
    real_esrgan = CHECKPOINTS / "RealESRGAN_x2plus.pth"
    if not real_esrgan.exists():
        hf_hub_download(
            "dtarnow/UPscaler",
            filename="RealESRGAN_x2plus.pth",
            local_dir=str(CHECKPOINTS),
        )
    lama = CHECKPOINTS / "big-lama.pt"
    if not lama.exists():
        subprocess.check_call(
            [
                "wget", "-q", "-O", str(lama),
                "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
            ]
        )
    return real_esrgan, lama


REAL_ESRGAN, BIG_LAMA = _ensure_checkpoints()
_NVD_READY = False
_TEXTURE_PIPE = None


def _cuda_home() -> str:
    candidates = [
        os.environ.get("CUDA_HOME"),
        "/cuda-image/usr/local/cuda-13.0",
        "/usr/local/cuda-13.0",
        "/usr/local/cuda",
    ]
    for item in candidates:
        if item and os.path.isdir(item):
            return item
    return "/cuda-image/usr/local/cuda-13.0"


def _ensure_nvdiffrast():
    """Compile/load nvdiffrast for current ZeroGPU Blackwell/PyTorch ABI."""
    global _NVD_READY
    if _NVD_READY:
        return
    try:
        import nvdiffrast.torch  # noqa: F401
        _NVD_READY = True
        return
    except Exception:
        pass

    cuda_home = _cuda_home()
    patch_dir = tempfile.mkdtemp(prefix="basilisc_torch_cuda_patch_")
    Path(patch_dir, "sitecustomize.py").write_text(
        "try:\n"
        " import torch.utils.cpp_extension as _c\n"
        " _c._check_cuda_version = lambda *a, **k: None\n"
        "except Exception:\n"
        " pass\n",
        encoding="utf-8",
    )
    env = os.environ.copy()
    env["CUDA_HOME"] = cuda_home
    env["CUDA_PATH"] = cuda_home
    env["PATH"] = os.path.join(cuda_home, "bin") + os.pathsep + env.get("PATH", "")
    env["PYTHONPATH"] = patch_dir + os.pathsep + env.get("PYTHONPATH", "")
    env["TORCH_CUDA_ARCH_LIST"] = "12.0"
    env["MAX_JOBS"] = "2"
    subprocess.check_call(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "--no-build-isolation",
            "--no-deps",
            "git+https://github.com/NVlabs/nvdiffrast/",
        ],
        env=env,
    )
    importlib.invalidate_caches()

    libdir = os.path.join(cuda_home, "lib64")
    for libname in ("libcudart.so.13", "libcudart.so"):
        lib = os.path.join(libdir, libname)
        if os.path.exists(lib):
            try:
                ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
                break
            except OSError:
                pass
    os.environ["LD_LIBRARY_PATH"] = libdir + os.pathsep + os.environ.get("LD_LIBRARY_PATH", "")
    import nvdiffrast.torch  # noqa: F401
    _NVD_READY = True


def _remove_bg(image: Image.Image) -> Image.Image:
    image = image.convert("RGB")
    original_size = image.size
    input_images = BG_TRANSFORM(image).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        preds = BIREFNET(input_images)[-1].sigmoid().cpu()
    mask = transforms.ToPILImage()(preds[0].squeeze()).resize(original_size)
    rgba = image.convert("RGBA")
    rgba.putalpha(mask)
    return rgba


def _preprocess_image(image: Image.Image, height: int, width: int) -> Image.Image:
    arr = np.array(image.convert("RGBA"))
    alpha = arr[..., 3] > 0
    ys, xs = np.where(alpha)
    if len(xs) == 0 or len(ys) == 0:
        return Image.new("RGB", (width, height), (127, 127, 127))
    y0, y1 = max(int(ys.min()) - 1, 0), min(int(ys.max()) + 2, arr.shape[0])
    x0, x1 = max(int(xs.min()) - 1, 0), min(int(xs.max()) + 2, arr.shape[1])
    crop = arr[y0:y1, x0:x1]
    ch, cw = crop.shape[:2]
    scale = min((height * 0.9) / max(ch, 1), (width * 0.9) / max(cw, 1))
    nh, nw = max(1, int(ch * scale)), max(1, int(cw * scale))
    crop = np.array(Image.fromarray(crop).resize((nw, nh), Image.Resampling.LANCZOS))
    canvas = np.zeros((height, width, 4), dtype=np.uint8)
    sy, sx = (height - nh) // 2, (width - nw) // 2
    canvas[sy:sy + nh, sx:sx + nw] = crop
    f = canvas.astype(np.float32) / 255.0
    rgb = f[..., :3] * f[..., 3:4] + (1.0 - f[..., 3:4]) * 0.5
    return Image.fromarray((rgb * 255).clip(0, 255).astype(np.uint8))


def _path(value) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        return value.get("path") or value.get("name") or ""
    return getattr(value, "name", "") or ""


def _run_multiview(mesh_path: str, image_path: str, prompt: str, seed: int, steps: int, ref_strength: float):
    from mvadapter.utils import tensor_to_image
    from mvadapter.utils.mesh_utils import (
        NVDiffRastContextWrapper,
        get_orthogonal_camera,
        load_mesh,
        render,
    )

    cameras = get_orthogonal_camera(
        elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
        distance=[1.8] * NUM_VIEWS,
        left=-0.55,
        right=0.55,
        bottom=-0.55,
        top=0.55,
        azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
        device=DEVICE,
    )
    ctx = NVDiffRastContextWrapper(device=DEVICE, context_type="cuda")
    mesh = load_mesh(mesh_path, rescale=True, device=DEVICE)
    render_out = render(
        ctx,
        mesh,
        cameras,
        height=HEIGHT,
        width=WIDTH,
        render_attr=False,
        normal_background=0.0,
    )
    control_images = (
        torch.cat(
            [
                (render_out.pos + 0.5).clamp(0, 1),
                (render_out.normal / 2 + 0.5).clamp(0, 1),
            ],
            dim=-1,
        )
        .permute(0, 3, 1, 2)
        .to(DEVICE)
    )
    reference = _preprocess_image(_remove_bg(Image.open(image_path)), HEIGHT, WIDTH)
    generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
    images = PIPE(
        prompt or "high quality",
        height=HEIGHT,
        width=WIDTH,
        num_inference_steps=int(steps),
        guidance_scale=3.0,
        num_images_per_prompt=NUM_VIEWS,
        control_image=control_images,
        control_conditioning_scale=1.0,
        reference_image=reference,
        reference_conditioning_scale=float(ref_strength),
        negative_prompt=NEGATIVE,
        cross_attention_kwargs={"scale": 1.0},
        generator=generator,
    ).images
    del control_images, render_out, mesh, cameras, ctx
    torch.cuda.empty_cache()
    return images


def _texture_pipeline():
    global _TEXTURE_PIPE
    if _TEXTURE_PIPE is not None:
        return _TEXTURE_PIPE

    # MV-Adapter defaults to an OpenGL raster context in the texture pipeline.
    # ZeroGPU is headless, so patch those constructors to use CUDA rasterization.
    import mvadapter.pipelines.pipeline_texture as pt
    import mvadapter.utils.mesh_utils.projection as projection_mod
    import mvadapter.utils.mesh_utils.render as render_mod
    import mvadapter.utils.mesh_utils.smart_paint as smart_mod

    OrigCtx = render_mod.NVDiffRastContextWrapper
    OrigProjection = projection_mod.CameraProjection

    class CudaCtx(OrigCtx):
        def __init__(self, device: str, context_type: str = "cuda"):
            super().__init__(device=device, context_type="cuda")

    class CudaProjection(OrigProjection):
        def __init__(self, pb_backend: str, bg_remover, device: str, context_type: str = "cuda"):
            super().__init__(pb_backend=pb_backend, bg_remover=bg_remover, device=device, context_type="cuda")

    render_mod.NVDiffRastContextWrapper = CudaCtx
    projection_mod.NVDiffRastContextWrapper = CudaCtx
    projection_mod.CameraProjection = CudaProjection
    smart_mod.NVDiffRastContextWrapper = CudaCtx
    smart_mod.CameraProjection = CudaProjection
    pt.NVDiffRastContextWrapper = CudaCtx
    pt.CameraProjection = CudaProjection

    _TEXTURE_PIPE = pt.TexturePipeline(
        upscaler_ckpt_path=str(REAL_ESRGAN),
        inpaint_ckpt_path=str(BIG_LAMA),
        device=DEVICE,
    )
    return _TEXTURE_PIPE


def _estimate_gpu_duration(
    mesh_file,
    reference_file,
    prompt="high quality",
    seed=42,
    steps=25,
    reference_strength=1.0,
    uv_unwrap=False,
    uv_size=2048,
):
    """Reserve only the ZeroGPU time appropriate for the selected preset.

    Draft is intentionally kept below the current anonymous daily allowance;
    Standard and Final are intended for authenticated Hugging Face users.
    """
    try:
        step_count = int(steps)
    except Exception:
        step_count = 25
    try:
        texture_size = int(uv_size)
    except Exception:
        texture_size = 2048
    if step_count <= 15 and texture_size <= 1024:
        return 115
    if step_count <= 25 and texture_size <= 2048:
        return 210
    return 285


@spaces.GPU(duration=_estimate_gpu_duration)
@torch.no_grad()
def basilisc_texture(
    mesh_file,
    reference_file,
    prompt="high quality",
    seed=42,
    steps=25,
    reference_strength=1.0,
    uv_unwrap=False,
    uv_size=2048,
):
    """One-shot mesh + reference image -> textured GLB endpoint for Blender."""
    _ensure_nvdiffrast()
    mesh_path = _path(mesh_file)
    image_path = _path(reference_file)
    if not mesh_path or not os.path.isfile(mesh_path):
        raise gr.Error("Input mesh file was not received")
    if not image_path or not os.path.isfile(image_path):
        raise gr.Error("Reference image file was not received")

    steps = max(8, min(int(steps), 50))
    uv_size = max(1024, min(int(uv_size), 4096))
    # Texture projection operates best on 512-aligned sizes.
    uv_size = int(round(uv_size / 512.0) * 512)
    seed = int(seed) if int(seed) >= 0 else random.randint(0, 2_147_483_647)

    job_dir = OUTPUTS / uuid.uuid4().hex
    job_dir.mkdir(parents=True, exist_ok=True)
    images = _run_multiview(mesh_path, image_path, prompt, seed, steps, reference_strength)

    from mvadapter.utils import make_image_grid
    from mvadapter.pipelines.pipeline_texture import ModProcessConfig

    grid_path = job_dir / "multiview.png"
    make_image_grid(images, rows=1).save(grid_path)
    del images
    torch.cuda.empty_cache()

    tex = _texture_pipeline()
    result = tex(
        mesh_path=mesh_path,
        save_dir=str(job_dir),
        save_name="basilisc_texture",
        uv_unwarp=bool(uv_unwrap),
        preprocess_mesh=False,
        uv_size=uv_size,
        rgb_path=str(grid_path),
        rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"),
        camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
    )
    torch.cuda.empty_cache()
    output = result.shaded_model_save_path
    if not output or not os.path.isfile(output):
        raise gr.Error("Texture pipeline completed without a GLB output")
    return output


with gr.Blocks(title="BASILISC Texture Worker") as demo:
    gr.Markdown(
        "# BASILISC Texture Worker\n"
        "Reference image + existing mesh → MV-Adapter multi-view generation → baked GLB texture."
    )
    with gr.Row():
        mesh = gr.File(label="Mesh (GLB)", file_types=[".glb"], type="filepath")
        reference = gr.Image(label="Reference Image", type="filepath")
    with gr.Accordion("Settings", open=False):
        prompt = gr.Textbox(label="Prompt", value="high quality")
        seed = gr.Number(label="Seed", value=42, precision=0)
        steps = gr.Slider(8, 50, value=25, step=1, label="Diffusion Steps")
        strength = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Reference Strength")
        uv_unwrap = gr.Checkbox(label="Generate New UVs", value=False)
        uv_size = gr.Dropdown([1024, 2048, 4096], value=2048, label="Texture Size")
    run = gr.Button("Texture Model", variant="primary")
    output = gr.File(label="Textured GLB")
    run.click(
        fn=basilisc_texture,
        inputs=[mesh, reference, prompt, seed, steps, strength, uv_unwrap, uv_size],
        outputs=[output],
        api_name="basilisc_texture",
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(show_error=True)