import os
import random
import shutil
import uuid
from glob import glob
from pathlib import Path
import argparse

import torch
import gradio as gr
import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from PIL import Image
import trimesh
from transformers import AutoProcessor, AutoModelForImageClassification

# -------------------- Argumente --------------------
parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, default='tencent/Hunyuan3D-2mini')
parser.add_argument("--subfolder", type=str, default='hunyuan3d-dit-v2-mini-turbo')
parser.add_argument("--texgen_model_path", type=str, default='tencent/Hunyuan3D-2')
parser.add_argument('--port', type=int, default=7860)
parser.add_argument('--host', type=str, default='0.0.0.0')
parser.add_argument('--device', type=str, default=None)
parser.add_argument('--mc_algo', type=str, default='mc')
parser.add_argument('--cache_path', type=str, default='gradio_cache')
parser.add_argument('--enable_t23d', action='store_true')
parser.add_argument('--disable_tex', action='store_true')
parser.add_argument('--enable_flashvdm', action='store_true')
parser.add_argument('--compile', action='store_true')
parser.add_argument('--low_vram_mode', action='store_true')
args = parser.parse_args()

# -------------------- Device Setup --------------------
if args.device is None:
    if torch.cuda.is_available():
        args.device = "cuda"
    elif torch.backends.mps.is_available():  # macOS GPU
        args.device = "mps"
    else:
        args.device = "cpu"

print(f"Using device: {args.device}")

# -------------------- Pfade --------------------
SAVE_DIR = args.cache_path
os.makedirs(SAVE_DIR, exist_ok=True)
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

HTML_HEIGHT = 500
HTML_WIDTH = 500
MAX_SEED = int(1e7)

# -------------------- NSFW Modell --------------------
nsfw_processor = AutoProcessor.from_pretrained("Falconsai/nsfw_image_detection")
nsfw_model = AutoModelForImageClassification.from_pretrained(
    "Falconsai/nsfw_image_detection"
).to(args.device)

# -------------------- Hilfsfunktionen --------------------
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

def gen_save_folder(max_size=200):
    os.makedirs(SAVE_DIR, exist_ok=True)
    dirs = [f for f in Path(SAVE_DIR).iterdir() if f.is_dir()]
    if len(dirs) >= max_size:
        oldest_dir = min(dirs, key=lambda x: x.stat().st_ctime)
        shutil.rmtree(oldest_dir)
        print(f"Removed the oldest folder: {oldest_dir}")
    new_folder = os.path.join(SAVE_DIR, str(uuid.uuid4()))
    os.makedirs(new_folder, exist_ok=True)
    print(f"Created new folder: {new_folder}")
    return new_folder

def export_mesh(mesh, save_folder, textured=False, type='glb'):
    filename = f'textured_mesh.{type}' if textured else f'white_mesh.{type}'
    path = os.path.join(save_folder, filename)
    mesh.export(path, include_normals=textured)
    return path

def build_model_viewer_html(save_folder, height=660, width=790, textured=False):
    related_path = "textured_mesh.glb" if textured else "white_mesh.glb"
    template_name = './assets/modelviewer-textured-template.html' if textured else './assets/modelviewer-template.html'
    output_html_path = os.path.join(save_folder, f"{'textured' if textured else 'white'}_mesh.html")

    with open(os.path.join(CURRENT_DIR, template_name), 'r', encoding='utf-8') as f:
        template_html = f.read()

    offset = 50 if textured else 10
    template_html = template_html.replace('#height#', str(height - offset))
    template_html = template_html.replace('#width#', str(width))
    template_html = template_html.replace('#src#', f'./{related_path}/')

    with open(output_html_path, 'w', encoding='utf-8') as f:
        f.write(template_html)

    rel_path = os.path.relpath(output_html_path, SAVE_DIR)
    iframe_tag = f'<iframe src="/static/{rel_path}" height="{height}" width="100%" frameborder="0"></iframe>'

    return f"<div style='height: {height}; width: 100%;'>{iframe_tag}</div>"

# -------------------- Hy3Dgen Worker --------------------
from hy3dgen.shapegen import FaceReducer, FloaterRemover, DegenerateFaceRemover, MeshSimplifier, Hunyuan3DDiTFlowMatchingPipeline
from hy3dgen.shapegen.pipelines import export_to_trimesh
from hy3dgen.rembg import BackgroundRemover

rmbg_worker = BackgroundRemover()
i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
    args.model_path,
    subfolder=args.subfolder,
    use_safetensors=True,
    device=args.device
)

if args.enable_flashvdm:
    mc_algo = 'mc' if args.device in ['cpu', 'mps'] else args.mc_algo
    i23d_worker.enable_flashvdm(mc_algo=mc_algo)
if args.compile:
    i23d_worker.compile()

floater_remove_worker = FloaterRemover()
degenerate_face_remove_worker = DegenerateFaceRemover()
face_reduce_worker = FaceReducer()

# -------------------- NSFW Detection --------------------
def detect_nsfw(image: Image.Image, threshold: float = 0.5) -> bool:
    nsfw_score = 0  # Placeholder, optional: implement actual detection
    return nsfw_score > threshold

# -------------------- Mesh Generation --------------------
progress = gr.Progress()

def _gen_shape(
    image=None,
    steps=10,
    guidance_scale=7.5,
    seed=1234,
    octree_resolution=128,
    num_chunks=50000,
    target_face_num=2500,
    randomize_seed: bool = False,
):
    progress(0, desc="Starting")

    if image is None:
        return None, None, None, None, {"error": "Please provide an image.", "status": "failed"}

    rgb_image = image.convert('RGB')
    if detect_nsfw(rgb_image):
        return None, None, None, None, {"error": "NSFW content detected.", "status": "failed"}

    seed = int(randomize_seed_fn(seed, randomize_seed))
    save_folder = gen_save_folder()
    image = rmbg_worker(rgb_image)

    if args.device in ["cuda", "cpu"]:
        generator = torch.Generator(device=args.device).manual_seed(seed)
    else:
        generator = torch.Generator().manual_seed(seed)  # fallback für MPS

    outputs = i23d_worker(
        image=image,
        num_inference_steps=steps,
        guidance_scale=guidance_scale,
        generator=generator,
        octree_resolution=octree_resolution,
        num_chunks=num_chunks,
        output_type='mesh',
        callback=lambda step_idx, timestep, out: progress(((step_idx+1)/steps)*0.5, desc=f"Mesh generating {step_idx+1}/{steps}"),
        callback_steps=1
    )

    mesh = export_to_trimesh(outputs)[0]
    path = export_mesh(mesh, save_folder, textured=False)

    mesh = trimesh.load(path)
    progress(0.5, desc="Optimizing mesh")
    mesh = floater_remove_worker(mesh)
    mesh = degenerate_face_remove_worker(mesh)
    progress(0.6, desc="Reducing faces")
    mesh = face_reduce_worker(mesh, target_face_num)

    save_folder = gen_save_folder()
    source_obj_path = export_mesh(mesh, save_folder, textured=False, type="obj")
    model_viewer_html = build_model_viewer_html(save_folder, height=HTML_HEIGHT, width=HTML_WIDTH, textured=False)

    glb_path = export_mesh(mesh, save_folder, textured=False, type="glb")
    rel_glb_path = os.path.relpath(glb_path, SAVE_DIR)
    glb_path = "/static/" + rel_glb_path
    rel_obj_path = os.path.relpath(source_obj_path, SAVE_DIR)
    obj_path = "/static/" + rel_obj_path

    progress(1, desc="Complete")
    return model_viewer_html, gr.update(value=source_obj_path, interactive=True), glb_path, obj_path, {"status": "success"}

def gen_shape(*args, **kwargs):
    html, file_export, glb_path, obj_path, info = _gen_shape(*args, **kwargs)
    if info["status"] == "failed":
        raise gr.Error(info["error"])
    return html, file_export, glb_path, obj_path

# -------------------- Beispielbilder --------------------
def get_example_img_list():
    return sorted(glob('./assets/example_images/**/*.png', recursive=True))

example_imgs = get_example_img_list()

HTML_OUTPUT_PLACEHOLDER = f"""
<div style='height: {500}px; width: 100%; border-radius: 8px; border-color: #e5e7eb; border-style: solid; border-width: 1px; display: flex; justify-content: center; align-items: center;'>
    <div style='text-align: center; font-size: 16px; color: #6b7280;'>
    <p style="color: #8d8d8d;">No mesh here.</p>
    </div>
</div>
"""

# -------------------- Gradio UI --------------------
title = "## AI 3D Model Generator"
description = "Transforms 2D photos into AI-generated 3D models."

with gr.Blocks().queue() as demo:
    gr.Markdown(title)
    gr.Markdown(description)
    with gr.Row():
        with gr.Column(scale=3):
            image = gr.Image(sources=["upload"], label='Image', type='pil', image_mode='RGBA', height=290)
            gen_button = gr.Button(value='Generate Shape', variant='primary')
            with gr.Accordion("Advanced Options", open=False):
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=1234)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                num_steps = gr.Slider(maximum=100, minimum=1, value=5, step=1, label='Inference Steps')
                octree_resolution = gr.Slider(maximum=512, minimum=16, value=128, label='Octree Resolution')
                cfg_scale = gr.Slider(maximum=20.0, minimum=1.0, value=5.5, step=0.1, label='Guidance Scale')
                num_chunks = gr.Slider(maximum=50000, minimum=1000, value=2000, label='Number of Chunks')
                target_face_num = gr.Slider(maximum=1000000, minimum=100, value=2500, label='Target Face Number')
        with gr.Column(scale=6):
            html_export_mesh = gr.HTML(HTML_OUTPUT_PLACEHOLDER, label='Output')
            file_export = gr.DownloadButton(label="Download", variant='primary', interactive=False)
            objPath_output = gr.Text(label="Obj Path", interactive=False)
            glbPath_output = gr.Text(label="Glb Path", interactive=False)
        with gr.Column(scale=3):
            gr.Examples(examples=example_imgs, inputs=[image], examples_per_page=18)

    gen_button.click(
        fn=gen_shape,
        inputs=[image, num_steps, cfg_scale, seed, octree_resolution, num_chunks, target_face_num, randomize_seed],
        outputs=[html_export_mesh, file_export, glbPath_output, objPath_output]
    )

# -------------------- FastAPI + Gradio --------------------
if __name__ == "__main__":
    # Device Info
    print(f"Using device: {args.device}")

    # Optional: FastAPI static files (für Assets)
    app = FastAPI()
    static_dir = Path(SAVE_DIR).absolute()
    static_dir.mkdir(parents=True, exist_ok=True)
    app.mount("/static", StaticFiles(directory=static_dir, html=True), name="static")
    shutil.copytree('./assets/env_maps', os.path.join(static_dir, 'env_maps'), dirs_exist_ok=True)

    # Low VRAM cleanup
    if args.low_vram_mode and args.device == "cuda":
        torch.cuda.empty_cache()

    # Gradio Demo starten CPU-kompatibel, funktioniert auch in HF Spaces
    demo.launch(
        server_name="0.0.0.0",  # für Spaces oder lokal
        server_port=args.port,
        share=False            # erstellt einen öffentlichen Link wie HF Spaces
    )
    app = gr.mount_gradio_app(app, demo, path="/")
    # from spaces import zero
    # zero.startup()
    uvicorn.run(app, host=args.host, port=args.port)