import os
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = '1'
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["ATTN_BACKEND"] = "flash_attn"
os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1' 

try:
    import spaces
    GPU = spaces.GPU
except ImportError:
    # Fallback for local development (no HuggingFace Spaces environment)
    spaces = None
    def GPU(duration=120):
        def decorator(fn):
            return fn
        return decorator

import gradio as gr

import shutil
import numpy as np
from PIL import Image
import torch
from datetime import datetime

from trellis2.modules.sparse import SparseTensor
from trellis2.pipelines import Trellis2ImageTo3DPipeline
from trellis2.utils import render_utils
import o_voxel


MAX_SEED = np.iinfo(np.int32).max
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
STEPS = 8


def start_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)


def end_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    shutil.rmtree(user_dir, ignore_errors=True)


def preprocess_image(input: Image.Image) -> Image.Image:
    """Preprocess the input image."""
    has_alpha = False
    if input.mode == 'RGBA':
        alpha = np.array(input)[:, :, 3]
        if not np.all(alpha == 255):
            has_alpha = True

    max_size = max(input.size)
    scale = min(1, 1024 / max_size)
    if scale < 1:
        input = input.resize(
            (int(input.width * scale), int(input.height * scale)),
            Image.Resampling.LANCZOS
        )

    if has_alpha:
        output = input
    else:
        # Simple background removal: assume near-white pixels are background
        output = input.convert('RGBA')
        arr = np.array(output)
        gray = arr[:, :, :3].mean(axis=2)
        mask = gray < 240
        arr[:, :, 3] = (mask * 255).astype(np.uint8)
        output = Image.fromarray(arr)

    output_np = np.array(output).astype(np.float32) / 255
    alpha = output_np[:, :, 3]
    bbox = np.argwhere(alpha > 0.8 * 255)
    if len(bbox) > 0:
        bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
        center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
        size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
        size = int(size * 1.15)
        bbox = (int(center[0] - size // 2), int(center[1] - size // 2),
                int(center[0] + size // 2), int(center[1] + size // 2))
        output = output.crop(bbox)
        output_np = np.array(output).astype(np.float32) / 255
        output = Image.fromarray(
            (output_np[:, :, :3] * output_np[:, :, 3:4] * 255).astype(np.uint8)
        )
    else:
        output = input.convert('RGB')

    return output


def pack_state(latents) -> dict:
    shape_slat, tex_slat, res = latents
    return {
        'shape_slat_feats': shape_slat.feats.cpu().numpy(),
        'tex_slat_feats': tex_slat.feats.cpu().numpy(),
        'coords': shape_slat.coords.cpu().numpy(),
        'res': res,
    }


def unpack_state(state: dict):
    shape_slat = SparseTensor(
        feats=torch.from_numpy(state['shape_slat_feats']).cuda(),
        coords=torch.from_numpy(state['coords']).cuda(),
    )
    tex_slat = shape_slat.replace(torch.from_numpy(state['tex_slat_feats']).cuda())
    return shape_slat, tex_slat, state['res']


def get_seed(randomize_seed: bool, seed: int) -> int:
    return np.random.randint(0, MAX_SEED) if randomize_seed else seed


@GPU(duration=120)
def image_to_3d(
    image: Image.Image,
    seed: int,
    resolution: str,
    ss_guidance_strength: float,
    ss_guidance_rescale: float,
    ss_sampling_steps: int,
    ss_rescale_t: float,
    shape_slat_guidance_strength: float,
    shape_slat_guidance_rescale: float,
    shape_slat_sampling_steps: int,
    shape_slat_rescale_t: float,
    tex_slat_guidance_strength: float,
    tex_slat_guidance_rescale: float,
    tex_slat_sampling_steps: int,
    tex_slat_rescale_t: float,
    req: gr.Request,
    progress=gr.Progress(track_tqdm=True),
):
    outputs, latents = pipeline.run(
        image,
        seed=seed,
        preprocess_image=False,
        sparse_structure_sampler_params={
            "steps": ss_sampling_steps,
            "guidance_strength": ss_guidance_strength,
            "guidance_rescale": ss_guidance_rescale,
            "rescale_t": ss_rescale_t,
        },
        shape_slat_sampler_params={
            "steps": shape_slat_sampling_steps,
            "guidance_strength": shape_slat_guidance_strength,
            "guidance_rescale": shape_slat_guidance_rescale,
            "rescale_t": shape_slat_rescale_t,
        },
        tex_slat_sampler_params={
            "steps": tex_slat_sampling_steps,
            "guidance_strength": tex_slat_guidance_strength,
            "guidance_rescale": tex_slat_guidance_rescale,
            "rescale_t": tex_slat_rescale_t,
        },
        pipeline_type={
            "512": "512",
            "1024": "1024_cascade",
            "1536": "1536_cascade",
        }[resolution],
        return_latent=True,
    )
    mesh = outputs[0]
    mesh.simplify(16777216)

    # Render preview images from multiple angles
    images = render_utils.render_snapshot(
        mesh, resolution=1024, r=2, fov=36, nviews=STEPS
    )
    render_key = 'normal' if 'normal' in images else list(images.keys())[0]
    preview_images = [Image.fromarray(img) for img in images[render_key]]

    state = pack_state(latents)
    torch.cuda.empty_cache()
    return state, preview_images


@GPU(duration=120)
def extract_glb(
    state: dict,
    decimation_target: int,
    texture_size: int,
    req: gr.Request,
    progress=gr.Progress(track_tqdm=True),
) -> str:
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    shape_slat, tex_slat, res = unpack_state(state)
    mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0]
    mesh.simplify(16777216)
    glb = o_voxel.postprocess.to_glb(
        vertices=mesh.vertices,
        faces=mesh.faces,
        attr_volume=mesh.attrs,
        coords=mesh.coords,
        attr_layout=pipeline.pbr_attr_layout,
        grid_size=res,
        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
        decimation_target=decimation_target,
        texture_size=texture_size,
        remesh=True,
        remesh_band=1,
        remesh_project=0,
        use_tqdm=True,
    )
    now = datetime.now()
    timestamp = now.strftime("%Y-%m-%dT%H%M%S") + f".{now.microsecond // 1000:03d}"
    os.makedirs(user_dir, exist_ok=True)
    glb_path = os.path.join(user_dir, f'sample_{timestamp}.glb')
    glb.export(glb_path, extension_webp=True)
    torch.cuda.empty_cache()
    return glb_path


with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    ## Image to 3D Asset with [TRELLIS.2-4B](https://huggingface.co/microsoft/TRELLIS.2-4B)
    
    Upload an image (preferably with an alpha-masked foreground object) and click **Generate** to create a high-fidelity 3D asset with PBR materials.
    
    Click **Extract GLB** to export and download the generated 3D model as a GLB file.
    
    *Model: [microsoft/TRELLIS.2-4B](https://huggingface.co/microsoft/TRELLIS.2-4B) | Paper: [arXiv:2512.14692](https://arxiv.org/abs/2512.14692) | License: MIT*
    """)

    with gr.Row():
        with gr.Column(scale=1, min_width=360):
            image_prompt = gr.Image(
                label="Image Prompt", format="png", image_mode="RGBA",
                type="pil", sources=["upload", "clipboard"], height=400
            )

            resolution = gr.Radio(
                ["512", "1024", "1536"], label="Resolution", value="1024",
                info="Higher = more detail but slower"
            )
            seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            decimation_target = gr.Slider(
                100000, 500000, label="Decimation Target (faces)",
                value=300000, step=10000
            )
            texture_size = gr.Slider(
                1024, 4096, label="Texture Size", value=2048, step=1024
            )

            generate_btn = gr.Button("Generate 3D Asset", variant="primary")

            with gr.Accordion(label="Advanced Settings", open=False):
                gr.Markdown("### Stage 1: Sparse Structure Generation")
                with gr.Row():
                    ss_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
                    ss_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.7, step=0.01)
                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
                    ss_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=5.0, step=0.1)

                gr.Markdown("### Stage 2: Shape Generation")
                with gr.Row():
                    shape_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
                    shape_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.5, step=0.01)
                    shape_slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
                    shape_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)

                gr.Markdown("### Stage 3: Material Generation")
                with gr.Row():
                    tex_slat_guidance_strength = gr.Slider(1.0, 10.0, label="Guidance Strength", value=1.0, step=0.1)
                    tex_slat_guidance_rescale = gr.Slider(0.0, 1.0, label="Guidance Rescale", value=0.0, step=0.01)
                    tex_slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
                    tex_slat_rescale_t = gr.Slider(1.0, 6.0, label="Rescale T", value=3.0, step=0.1)

        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Preview"):
                    preview_gallery = gr.Gallery(
                        label="3D Asset Preview (multiple viewing angles)",
                        columns=4, height=500,
                        show_label=True
                    )
                    extract_btn = gr.Button("Extract GLB", variant="secondary")

                with gr.TabItem("3D Viewer"):
                    glb_output = gr.Model3D(
                        label="Extracted GLB (interactive viewer)",
                        height=600, show_label=True,
                        display_mode="solid",
                        clear_color=(0.25, 0.25, 0.25, 1.0)
                    )
                    download_btn = gr.DownloadButton(label="Download GLB File")

            gr.Markdown("""
            *GLB extraction may take 30+ seconds. Face count is limited by the decimation target.*
            *The GLB is exported in OPAQUE mode by default. To enable transparency, import into your 3D software and connect the texture alpha channel to the material opacity input.*
            """)

    output_buf = gr.State()

    # Handlers
    demo.load(start_session)
    demo.unload(end_session)

    image_prompt.upload(
        preprocess_image,
        inputs=[image_prompt],
        outputs=[image_prompt],
    )

    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
    ).then(
        image_to_3d,
        inputs=[
            image_prompt, seed, resolution,
            ss_guidance_strength, ss_guidance_rescale, ss_sampling_steps, ss_rescale_t,
            shape_slat_guidance_strength, shape_slat_guidance_rescale, shape_slat_sampling_steps, shape_slat_rescale_t,
            tex_slat_guidance_strength, tex_slat_guidance_rescale, tex_slat_sampling_steps, tex_slat_rescale_t,
        ],
        outputs=[output_buf, preview_gallery],
    )

    extract_btn.click(
        extract_glb,
        inputs=[output_buf, decimation_target, texture_size],
        outputs=[glb_output, download_btn],
    )


if __name__ == "__main__":
    os.makedirs(TMP_DIR, exist_ok=True)

    pipeline = Trellis2ImageTo3DPipeline.from_pretrained('microsoft/TRELLIS.2-4B')
    pipeline.rembg_model = None
    pipeline.low_vram = True
    pipeline.cuda()

    demo.launch()