import gradio as gr
import spaces
from gradio_litmodel3d import LitModel3D

import os
import shutil
os.environ['SPCONV_ALGO'] = 'native'
from typing import *
import torch
import numpy as np
import imageio
import uuid
from easydict import EasyDict as edict
from PIL import Image
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils

# Constants
MAX_SEED = np.iinfo(np.int32).max
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
os.makedirs(TMP_DIR, exist_ok=True)

# Session Management Functions
def start_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    print(f'Creating user directory: {user_dir}')
    os.makedirs(user_dir, exist_ok=True)

def end_session(req: gr.Request):
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    print(f'Removing user directory: {user_dir}')
    shutil.rmtree(user_dir)

# Image Preprocessing Function
def preprocess_image(image: Image.Image) -> Image.Image:
    """
    Preprocess the input image.

    Args:
        image (Image.Image): The input image.

    Returns:
        Image.Image: The preprocessed image.
    """
    processed_image = pipeline.preprocess_image(image)
    return processed_image

# State Packing and Unpacking Functions
def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
    return {
        'gaussian': {
            **gs.init_params,
            '_xyz': gs._xyz.cpu().numpy(),
            '_features_dc': gs._features_dc.cpu().numpy(),
            '_scaling': gs._scaling.cpu().numpy(),
            '_rotation': gs._rotation.cpu().numpy(),
            '_opacity': gs._opacity.cpu().numpy(),
        },
        'mesh': {
            'vertices': mesh.vertices.cpu().numpy(),
            'faces': mesh.faces.cpu().numpy(),
        },
        'trial_id': trial_id,
    }

def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
    gs = Gaussian(
        aabb=state['gaussian']['aabb'],
        sh_degree=state['gaussian']['sh_degree'],
        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
        scaling_bias=state['gaussian']['scaling_bias'],
        opacity_bias=state['gaussian']['opacity_bias'],
        scaling_activation=state['gaussian']['scaling_activation'],
    )
    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
    
    mesh = edict(
        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
    )
    
    return gs, mesh, state['trial_id']

# Seed Management Function
def get_seed(randomize_seed: bool, seed: int) -> int:
    """
    Get the random seed.

    Args:
        randomize_seed (bool): Whether to randomize the seed.
        seed (int): The provided seed value.

    Returns:
        int: The final seed to use.
    """
    return np.random.randint(0, MAX_SEED) if randomize_seed else seed

# Core 3D Generation Function
@spaces.GPU
def image_to_3d(
    image: Image.Image,
    seed: int,
    ss_guidance_strength: float,
    ss_sampling_steps: int,
    slat_guidance_strength: float,
    slat_sampling_steps: int,
    req: gr.Request,
) -> Tuple[dict, str]:
    """
    Convert an image to a 3D model.

    Args:
        image (Image.Image): The input image.
        seed (int): The random seed.
        ss_guidance_strength (float): The guidance strength for sparse structure generation.
        ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
        slat_guidance_strength (float): The guidance strength for structured latent generation.
        slat_sampling_steps (int): The number of sampling steps for structured latent generation.
        req (gr.Request): Gradio request object.

    Returns:
        Tuple[dict, str]: The state dictionary and the path to the generated video.
    """
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    outputs = pipeline.run(
        image,
        seed=seed,
        formats=["gaussian", "mesh"],
        preprocess_image=False,
        sparse_structure_sampler_params={
            "steps": ss_sampling_steps,
            "cfg_strength": ss_guidance_strength,
        },
        slat_sampler_params={
            "steps": slat_sampling_steps,
            "cfg_strength": slat_guidance_strength,
        },
    )
    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
    trial_id = uuid.uuid4()
    video_path = os.path.join(user_dir, f"{trial_id}.mp4")
    imageio.mimsave(video_path, video, fps=15)
    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
    torch.cuda.empty_cache()
    return state, video_path

# Existing GLB Extraction Function
@spaces.GPU
def extract_glb(
    state: dict,
    mesh_simplify: float,
    texture_size: int,
    req: gr.Request,
) -> Tuple[dict, bytes]:
    """
    Extract a GLB file from the 3D model.

    Args:
        state (dict): The state of the generated 3D model.
        mesh_simplify (float): The mesh simplification factor.
        texture_size (int): The texture resolution.
        req (gr.Request): Gradio request object.

    Returns:
        Tuple[dict, bytes]: The model state and the GLB file bytes.
    """
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    gs, mesh, trial_id = unpack_state(state)
    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
    glb_path = os.path.join(user_dir, f"{trial_id}.glb")
    glb.export(glb_path)
    # Read the GLB file as bytes
    with open(glb_path, "rb") as f:
        glb_bytes = f.read()
    torch.cuda.empty_cache()
    return state, glb_bytes

# New High-Quality GLB Extraction Function
@spaces.GPU
def extract_glb_high_quality(
    state: dict,
    req: gr.Request,
) -> Tuple[dict, bytes]:
    """
    Extract a high-quality GLB file from the 3D model without polygon reduction.

    Args:
        state (dict): The state of the generated 3D model.
        req (gr.Request): Gradio request object.

    Returns:
        Tuple[dict, bytes]: The model state and the high-quality GLB file bytes.
    """
    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
    gs, mesh, trial_id = unpack_state(state)
    # Set simplify to 0.0 to disable polygon reduction
    # Set texture_size to 2048 for maximum texture quality
    glb = postprocessing_utils.to_glb(gs, mesh, simplify=0.0, texture_size=2048, verbose=False)
    glb_path = os.path.join(user_dir, f"{trial_id}_high_quality.glb")
    glb.export(glb_path)
    # Read the GLB file as bytes
    with open(glb_path, "rb") as f:
        glb_bytes = f.read()
    torch.cuda.empty_cache()
    return state, glb_bytes

# Gradio Interface Definition
with gr.Blocks(delete_cache=(600, 600)) as demo:
    gr.Markdown("""
    ## Image to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
    * Upload an image and click "Generate" to create a 3D asset. If the image has an alpha channel, it will be used as the mask. Otherwise, the background will be removed automatically.
    * If you find the generated 3D asset satisfactory, click "Extract GLB" to extract the GLB file and download it.
    * **New:** Click "Download High Quality GLB" to download the GLB file without any polygon reduction and with maximum texture quality.
    """)

    with gr.Row():
        with gr.Column():
            # Image Input
            image_prompt = gr.Image(
                label="Image Prompt",
                format="png",
                image_mode="RGBA",
                type="pil",
                height=300
            )
            
            # Generation Settings Accordion
            with gr.Accordion(label="Generation Settings", open=False):
                seed = gr.Slider(
                    0,
                    MAX_SEED,
                    label="Seed",
                    value=0,
                    step=1
                )
                randomize_seed = gr.Checkbox(
                    label="Randomize Seed",
                    value=True
                )
                gr.Markdown("### Stage 1: Sparse Structure Generation")
                with gr.Row():
                    ss_guidance_strength = gr.Slider(
                        0.0,
                        10.0,
                        label="Guidance Strength",
                        value=7.5,
                        step=0.1
                    )
                    ss_sampling_steps = gr.Slider(
                        1,
                        500,
                        label="Sampling Steps",
                        value=12,
                        step=1
                    )
                gr.Markdown("### Stage 2: Structured Latent Generation")
                with gr.Row():
                    slat_guidance_strength = gr.Slider(
                        0.0,
                        10.0,
                        label="Guidance Strength",
                        value=3.0,
                        step=0.1
                    )
                    slat_sampling_steps = gr.Slider(
                        1,
                        500,
                        label="Sampling Steps",
                        value=12,
                        step=1
                    )

            # Generate Button
            generate_btn = gr.Button("Generate")
            
            # GLB Extraction Settings Accordion
            with gr.Accordion(label="GLB Extraction Settings", open=False):
                mesh_simplify = gr.Slider(
                    0.0,
                    0.98,
                    label="Simplify",
                    value=0.95,
                    step=0.01
                )
                texture_size = gr.Slider(
                    512,
                    2048,
                    label="Texture Size",
                    value=1024,
                    step=512
                )
            
            # Existing Extract GLB Button
            extract_glb_btn = gr.Button("Extract GLB", interactive=False)
            
            # New Extract High Quality GLB Button
            extract_glb_high_quality_btn = gr.Button("Download High Quality GLB", interactive=False)

        with gr.Column():
            # Video Output
            video_output = gr.Video(
                label="Generated 3D Asset",
                autoplay=True,
                loop=True,
                height=300
            )
            # 3D Model Display
            model_output = LitModel3D(
                label="Extracted GLB",
                exposure=20.0,
                height=300
            )
            # Existing Download GLB Button
            download_glb = gr.DownloadButton(
                label="Download GLB",
                interactive=False  # Initially disabled
            )
            # New Download High Quality GLB Button
            download_high_quality_glb = gr.DownloadButton(
                label="Download High Quality GLB",
                interactive=False  # Initially disabled
            )

    # State Variables
    output_buf = gr.State()
    glb_bytes_state = gr.State()  # For standard GLB
    glb_high_quality_bytes_state = gr.State()  # For high-quality GLB

    # Example Images
    with gr.Row():
        examples = gr.Examples(
            examples=[
                f'assets/example_image/{image}'
                for image in os.listdir("assets/example_image")
            ],
            inputs=[image_prompt],
            fn=preprocess_image,
            outputs=[image_prompt],
            run_on_click=True,
            examples_per_page=64,
        )

    # Event Handlers
    demo.load(start_session)
    demo.unload(end_session)
    
    # Image Upload Handler
    image_prompt.upload(
        preprocess_image,
        inputs=[image_prompt],
        outputs=[image_prompt],
    )

    # Generate Button Click Handler
    generate_btn.click(
        get_seed,
        inputs=[randomize_seed, seed],
        outputs=[seed],
    ).then(
        image_to_3d,
        inputs=[
            image_prompt,
            seed,
            ss_guidance_strength,
            ss_sampling_steps,
            slat_guidance_strength,
            slat_sampling_steps
        ],
        outputs=[output_buf, video_output],
    ).then(
        # Enable both Extract GLB buttons after generation
        lambda: (gr.Button.update(interactive=True), gr.Button.update(interactive=True)),
        outputs=[extract_glb_btn, extract_glb_high_quality_btn],
    )

    # Existing Extract GLB Button Click Handler
    extract_glb_btn.click(
        extract_glb,
        inputs=[output_buf, mesh_simplify, texture_size],
        outputs=[model_output, glb_bytes_state],
    ).then(
        # Map the GLB bytes to the DownloadButton
        lambda glb_bytes: (glb_bytes, "model.glb"),
        inputs=[glb_bytes_state],
        outputs=[download_glb],
    ).then(
        # Enable the Download GLB button
        lambda: gr.DownloadButton.update(interactive=True),
        outputs=[download_glb],
    )

    # New Extract High Quality GLB Button Click Handler
    extract_glb_high_quality_btn.click(
        extract_glb_high_quality,
        inputs=[output_buf],
        outputs=[model_output, glb_high_quality_bytes_state],
    ).then(
        # Map the high-quality GLB bytes to the DownloadButton
        lambda glb_bytes: (glb_bytes, "model_high_quality.glb"),
        inputs=[glb_high_quality_bytes_state],
        outputs=[download_high_quality_glb],
    ).then(
        # Enable the Download High Quality GLB button
        lambda: gr.DownloadButton.update(interactive=True),
        outputs=[download_high_quality_glb],
    )

    # Handle Clearing of Video Output
    video_output.clear(
        lambda: (gr.Button.update(interactive=False), gr.Button.update(interactive=False)),
        outputs=[extract_glb_btn, extract_glb_high_quality_btn],
    )

    # Handle Clearing of Model Output
    model_output.clear(
        lambda: (gr.File.update(value=None), gr.File.update(value=None)),
        outputs=[download_glb, download_high_quality_glb],
    )

# Launch the Gradio app
if __name__ == "__main__":
    # Initialize the pipeline
    pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
    pipeline.cuda()
    try:
        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
    except:
        pass
    demo.launch()