import os
import sys
import subprocess
import tempfile
import spaces

# Build pytorch3d and nvdiffrast against the installed torch the first time a GPU is allocated.
# Both packages publish no prebuilt wheels for torch >= 2.8 (Blackwell / CUDA 13 stack).
CUDA_HOME = "/cuda-image/usr/local/cuda-13.0"


def _build_env():
    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["TORCH_CUDA_ARCH_LIST"] = "12.0"
    return env


def _ensure_pkg(name, install_args):
    try:
        __import__(name)
        return True
    except ImportError:
        pass
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install",
         "--no-build-isolation",
         "--no-deps",
         *install_args],
        env=_build_env(),
    )


def _strip_pulsar_from_pytorch3d_src(src):
    """Drop pulsar from a cloned pytorch3d source tree.

    Pulsar fails to link on modern GCC because the explicit template
    instantiations of pulsar::Renderer::fill_bg<true>/forward<true>/...
    end up as hidden symbols and ld errors with:

        hidden symbol `..._fill_bg<true>...' isn't defined
        final link failed: bad value

    Unique3D never imports pulsar (only pytorch3d.renderer.{cameras,mesh,...}
    and pytorch3d.structures). Strip pulsar entirely before building.
    """
    import shutil

    # 1. Remove the C++/CUDA sources so setup.py's glob skips them.
    pulsar_csrc = os.path.join(src, "pytorch3d", "csrc", "pulsar")
    if os.path.isdir(pulsar_csrc):
        shutil.rmtree(pulsar_csrc)

    # 2. Remove the python-side pulsar package.
    pulsar_py = os.path.join(src, "pytorch3d", "renderer", "points", "pulsar")
    if os.path.isdir(pulsar_py):
        shutil.rmtree(pulsar_py)

    # 3. Rewrite ext.cpp to drop every pulsar-related include and the entire
    # pulsar pybind11 registration block. Use the well-defined "Pulsar." comment
    # marker that starts the registration block and run to the matching #endif.
    ext_cpp = os.path.join(src, "pytorch3d", "csrc", "ext.cpp")
    if os.path.isfile(ext_cpp):
        with open(ext_cpp, "r") as f:
            lines = f.readlines()

        # Phase A: drop the three pulsar include lines (they sit inside
        # `#if !defined(USE_ROCM)` blocks together with non-pulsar includes? No —
        # check the file: the top of ext.cpp has two #if !defined(USE_ROCM)
        # blocks that only contain pulsar includes. Drop the whole #if/#endif
        # pairs that contain only pulsar includes.
        out = []
        i = 0
        while i < len(lines):
            ln = lines[i]
            if ln.strip().startswith("#if !defined(USE_ROCM)"):
                # Look ahead for the matching #endif. Collect block body.
                j = i + 1
                block = []
                while j < len(lines) and not lines[j].strip().startswith("#endif"):
                    block.append(lines[j])
                    j += 1
                # Decide: if every non-blank/non-comment line in block mentions
                # "pulsar", drop the whole #if..#endif. Otherwise keep as-is.
                meaningful = [
                    b for b in block
                    if b.strip() and not b.strip().startswith("//")
                ]
                if meaningful and all("pulsar" in b.lower() for b in meaningful):
                    # Drop block (and matching #endif).
                    i = j + 1
                    continue
            out.append(ln)
            i += 1

        # Phase B: drop the pulsar pybind registration block.
        # Find the "// Pulsar." comment, walk back to its preceding blank line,
        # then forward to the matching #endif (the last #endif before `}`).
        text = "".join(out)
        marker = "  // Pulsar."
        idx = text.find(marker)
        if idx >= 0:
            # Find the start of the line containing the marker.
            line_start = text.rfind("\n", 0, idx) + 1
            # Find the trailing `#endif` of the pulsar block. Walk forward; the
            # block ends with a `#endif` on a line by itself, right before `}`.
            end_marker = "\n#endif\n}"
            end_idx = text.find(end_marker, line_start)
            if end_idx < 0:
                # Fall back: just drop to the next `}` on its own line.
                end_idx = text.find("\n}\n", line_start)
                if end_idx < 0:
                    raise RuntimeError("could not locate pulsar block end in ext.cpp")
                # Keep the `}`.
                text = text[:line_start] + text[end_idx + 1:]
            else:
                # Drop the `#endif` line but keep the `}`.
                text = text[:line_start] + text[end_idx + len("\n#endif"):]

        with open(ext_cpp, "w") as f:
            f.write(text)

    # 4. Patch renderer/points/__init__.py AND renderer/__init__.py: both have a
    # `if not torch.version.hip:` block that imports PulsarPointsRenderer; replace
    # the import with `pass` so the symbol is just absent.
    for init_path in [
        os.path.join(src, "pytorch3d", "renderer", "points", "__init__.py"),
        os.path.join(src, "pytorch3d", "renderer", "__init__.py"),
    ]:
        if not os.path.isfile(init_path):
            continue
        with open(init_path, "r") as f:
            init_lines = f.readlines()
        new_lines = []
        skip_block = False
        for ln in init_lines:
            if "torch.version.hip" in ln and ln.lstrip().startswith("if "):
                indent = " " * (len(ln) - len(ln.lstrip()))
                new_lines.append(ln)
                new_lines.append(indent + "    pass  # pulsar disabled (Blackwell linker)\n")
                skip_block = True
                continue
            if skip_block:
                # Drop any indented continuation lines of the if-block, including
                # multi-line `from .points import (` parenthesized imports.
                stripped = ln.strip()
                if stripped == "" or ln.startswith(" ") or ln.startswith("\t"):
                    if "pulsar" in ln.lower() or stripped == "" or stripped.startswith(")") or stripped.startswith("from ") or stripped.startswith("import "):
                        # part of the pulsar import block — drop
                        continue
                    skip_block = False
                else:
                    skip_block = False
            new_lines.append(ln)
        with open(init_path, "w") as f:
            f.writelines(new_lines)


@spaces.GPU(duration=1500)
def _build_pytorch3d():
    # Probe the actual import path Unique3D uses; if the previously installed
    # pytorch3d still has a broken `renderer/__init__.py` (e.g. left a dangling
    # `from .points import PulsarPointsRenderer`), uninstall and rebuild.
    try:
        from pytorch3d.renderer import TexturesVertex  # noqa: F401
        return
    except Exception:
        pass
    subprocess.call(
        [sys.executable, "-m", "pip", "uninstall", "-y", "pytorch3d"],
    )

    src = tempfile.mkdtemp(prefix="pytorch3d_src_")
    subprocess.check_call(
        ["git", "clone", "--depth=1", "--branch=stable",
         "https://github.com/facebookresearch/pytorch3d.git", src],
    )
    _strip_pulsar_from_pytorch3d_src(src)

    # Force a fresh wheel build (don't reuse the cached broken wheel from a
    # previous run, which still has the pulsar import in renderer/__init__.py).
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install",
         "--no-build-isolation", "--no-deps",
         "--force-reinstall", "--no-cache-dir", src],
        env=_build_env(),
    )


@spaces.GPU(duration=600)
def _build_nvdiffrast():
    _ensure_pkg("nvdiffrast", ["git+https://github.com/NVlabs/nvdiffrast.git"])


# Run setup before importing anything that needs these packages.
_build_pytorch3d()
_build_nvdiffrast()


if __name__ == "__main__":
    from huggingface_hub import login
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        login(token=hf_token)

    sys.path.append(os.curdir)
    import torch
    torch.set_float32_matmul_precision('medium')
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_grad_enabled(False)

import fire
import gradio as gr
from gradio_app.gradio_3dgen import create_ui as create_3d_ui
from gradio_app.all_models import model_zoo


_TITLE = '''Unique3D: High-Quality and Efficient 3D Mesh Generation from a Single Image'''
_DESCRIPTION = '''

<div>
    <a style="display:inline-block" href='https://github.com/AiuniAI/Unique3D'><img alt="GitHub Repo stars" src="https://img.shields.io/github/stars/AiuniAI/Unique3D?style=social">
</a>
    <img alt="GitHub License" src="https://img.shields.io/github/license/AiuniAI/Unique3D">
</div>

# [Paper](https://arxiv.org/abs/2405.20343) | [Project page](https://wukailu.github.io/Unique3D/) | [Huggingface Demo](https://huggingface.co/spaces/Wuvin/Unique3D) | [Gradio Demo](http://unique3d.demo.avar.cn/) | [Online Demo](https://www.aiuni.ai/)

* High-fidelity and diverse textured meshes generated by Unique3D from single-view images.

* The demo is still under construction, and more features are expected to be implemented soon.

* If the Huggingface Demo is overcrowded or fails to produce stable results, you can use the Online Demo [aiuni.ai](https://www.aiuni.ai/), which is free to try (get the registration invitation code Join Discord: https://discord.gg/aiuni). However, the Online Demo is slightly different from the Gradio Demo, in that the inference speed is slower, but the generation is much more stable.
'''

def launch():
    model_zoo.init_models()

    with gr.Blocks(
        title=_TITLE,
        # theme=gr.themes.Monochrome(),
    ) as demo:
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown('# ' + _TITLE)
        gr.Markdown(_DESCRIPTION)
        create_3d_ui("wkl")

    demo.queue().launch(share=True)

if __name__ == '__main__':
    fire.Fire(launch)