"""
Batch 3D model generation script for Hunyuan3D-2.1.
Optimized for RTX 3080 (20GB VRAM) by sequential model loading:
  Phase 1: Load shape model → generate all meshes → unload
  Phase 2: Load texture model → texture all meshes → unload
"""

import sys
sys.path.insert(0, './hy3dshape')
sys.path.insert(0, './hy3dpaint')

try:
    from torchvision_fix import apply_fix
    apply_fix()
except Exception as e:
    print(f"Warning: torchvision fix: {e}")

import os
import gc
import time
import glob
import torch
import trimesh
import numpy as np
from PIL import Image
from pathlib import Path

INPUT_DIR = "test/images"
OUTPUT_DIR = "test/models"
MODEL_PATH = "tencent/Hunyuan3D-2.1"
SUBFOLDER = "hunyuan3d-dit-v2-1"

SHAPE_STEPS = 50
GUIDANCE_SCALE = 7.5
SEED = 1234
OCTREE_RESOLUTION = 256
NUM_CHUNKS = 200000

TEXGEN_MAX_VIEWS = 6
TEXGEN_RESOLUTION = 512


def clear_gpu():
    """Aggressively free GPU memory."""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()


def get_image_files(input_dir):
    """Get all supported image files from input directory."""
    extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.webp']
    files = []
    for ext in extensions:
        files.extend(glob.glob(os.path.join(input_dir, ext)))
        files.extend(glob.glob(os.path.join(input_dir, ext.upper())))
    return sorted(set(files))


def phase1_shape_generation(image_files, output_dir):
    """Phase 1: Load shape model, generate all meshes, unload."""
    print("\n" + "=" * 60)
    print("PHASE 1: Shape Generation")
    print("=" * 60)

    from hy3dshape.rembg import BackgroundRemover
    from hy3dshape.pipelines import Hunyuan3DDiTFlowMatchingPipeline
    from hy3dshape.pipelines import export_to_trimesh
    from hy3dshape import FaceReducer

    print("Loading shape model...")
    t0 = time.time()
    rmbg = BackgroundRemover()
    pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
        MODEL_PATH, subfolder=SUBFOLDER, use_safetensors=False, device='cuda'
    )
    face_reducer = FaceReducer()
    print(f"Shape model loaded in {time.time()-t0:.1f}s")
    print(f"GPU memory: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

    results = {}
    for i, img_path in enumerate(image_files):
        name = Path(img_path).stem
        item_dir = os.path.join(output_dir, name)
        os.makedirs(item_dir, exist_ok=True)

        mesh_path = os.path.join(item_dir, "white_mesh.obj")
        if os.path.exists(mesh_path):
            print(f"[{i+1}/{len(image_files)}] {name}: shape exists, skipping")
            results[name] = {"image": img_path, "mesh": mesh_path, "dir": item_dir}
            continue

        print(f"\n[{i+1}/{len(image_files)}] Generating shape for: {name}")
        t1 = time.time()
        try:
            image = Image.open(img_path).convert("RGBA")
            if image.mode == "RGB" or image.getchannel("A").getextrema()[0] > 250:
                image = rmbg(image.convert("RGB"))

            generator = torch.Generator().manual_seed(SEED)
            outputs = pipeline(
                image=image,
                num_inference_steps=SHAPE_STEPS,
                guidance_scale=GUIDANCE_SCALE,
                generator=generator,
                octree_resolution=OCTREE_RESOLUTION,
                num_chunks=NUM_CHUNKS,
                output_type='mesh',
            )
            mesh = export_to_trimesh(outputs)[0]

            # Face reduction for texture gen compatibility
            mesh = face_reducer(mesh)
            mesh.export(mesh_path, include_normals=False)

            # Save input image alongside mesh for texture gen
            input_copy = os.path.join(item_dir, "input.png")
            image.save(input_copy)

            results[name] = {"image": img_path, "mesh": mesh_path, "dir": item_dir}
            print(f"  Done in {time.time()-t1:.1f}s | faces: {mesh.faces.shape[0]}")

        except Exception as e:
            print(f"  ERROR: {e}")
            import traceback; traceback.print_exc()
            results[name] = {"image": img_path, "mesh": None, "dir": item_dir, "error": str(e)}

        clear_gpu()

    # Unload shape model
    print("\nUnloading shape model...")
    del pipeline, rmbg, face_reducer
    clear_gpu()
    print(f"GPU memory after unload: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

    return results


def phase2_texture_generation(results, output_dir):
    """Phase 2: Load texture model, texture all meshes, unload."""
    print("\n" + "=" * 60)
    print("PHASE 2: Texture Generation")
    print("=" * 60)

    meshes_to_texture = {k: v for k, v in results.items() if v.get("mesh")}
    if not meshes_to_texture:
        print("No meshes to texture!")
        return results

    from hy3dpaint.textureGenPipeline import Hunyuan3DPaintPipeline, Hunyuan3DPaintConfig
    from hy3dpaint.convert_utils import create_glb_with_pbr_materials

    print("Loading texture model...")
    t0 = time.time()
    conf = Hunyuan3DPaintConfig(TEXGEN_MAX_VIEWS, TEXGEN_RESOLUTION)
    conf.realesrgan_ckpt_path = "hy3dpaint/ckpt/RealESRGAN_x4plus.pth"
    conf.multiview_cfg_path = "hy3dpaint/cfgs/hunyuan-paint-pbr.yaml"
    conf.custom_pipeline = "hy3dpaint/hunyuanpaintpbr"
    tex_pipeline = Hunyuan3DPaintPipeline(conf)
    print(f"Texture model loaded in {time.time()-t0:.1f}s")
    print(f"GPU memory: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

    for i, (name, info) in enumerate(meshes_to_texture.items()):
        item_dir = info["dir"]
        mesh_path = info["mesh"]
        img_path = info["image"]

        textured_obj = os.path.join(item_dir, "textured_mesh.obj")
        textured_glb = os.path.join(item_dir, "textured_mesh.glb")

        if os.path.exists(textured_glb):
            print(f"[{i+1}/{len(meshes_to_texture)}] {name}: textured mesh exists, skipping")
            results[name]["textured_glb"] = textured_glb
            continue

        print(f"\n[{i+1}/{len(meshes_to_texture)}] Texturing: {name}")
        t1 = time.time()
        try:
            output_path = tex_pipeline(
                mesh_path=mesh_path,
                image_path=img_path,
                output_mesh_path=textured_obj,
                save_glb=False,
            )

            # Convert OBJ to GLB with PBR materials
            textures = {
                'albedo': output_path.replace('.obj', '.jpg'),
                'metallic': output_path.replace('.obj', '_metallic.jpg'),
                'roughness': output_path.replace('.obj', '_roughness.jpg'),
            }
            create_glb_with_pbr_materials(output_path, textures, textured_glb)

            results[name]["textured_obj"] = output_path
            results[name]["textured_glb"] = textured_glb
            print(f"  Done in {time.time()-t1:.1f}s")
            print(f"  Output: {textured_glb}")

        except Exception as e:
            print(f"  ERROR: {e}")
            import traceback; traceback.print_exc()
            results[name]["tex_error"] = str(e)

        clear_gpu()

    # Unload texture model
    print("\nUnloading texture model...")
    del tex_pipeline
    clear_gpu()
    print(f"GPU memory after unload: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

    return results


def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    image_files = get_image_files(INPUT_DIR)
    if not image_files:
        print(f"No images found in {INPUT_DIR}")
        return

    print(f"Found {len(image_files)} images:")
    for f in image_files:
        print(f"  - {os.path.basename(f)}")

    total_start = time.time()

    # Phase 1: Shape generation (shape model only in VRAM)
    results = phase1_shape_generation(image_files, OUTPUT_DIR)

    # Phase 2: Texture generation (texture model only in VRAM)
    results = phase2_texture_generation(results, OUTPUT_DIR)

    # Summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    total_time = time.time() - total_start
    success = sum(1 for v in results.values() if v.get("textured_glb"))
    shape_only = sum(1 for v in results.values() if v.get("mesh") and not v.get("textured_glb"))
    failed = sum(1 for v in results.values() if not v.get("mesh"))

    for name, info in results.items():
        status = "✓ textured" if info.get("textured_glb") else (
            "△ shape only" if info.get("mesh") else "✗ failed"
        )
        print(f"  {name}: {status}")

    print(f"\nTotal: {len(results)} | Success: {success} | Shape only: {shape_only} | Failed: {failed}")
    print(f"Total time: {total_time:.1f}s ({total_time/60:.1f}m)")
    print(f"Output directory: {os.path.abspath(OUTPUT_DIR)}")


if __name__ == "__main__":
    main()