import os
import sys
import shutil
import argparse
from omegaconf import OmegaConf
import trimesh
from PIL import Image
import torch
from torchvision import transforms
import logging as log

# Add project root to path
project_root = os.path.dirname(os.path.abspath(__file__))
if project_root not in sys.path:
    sys.path.insert(0, project_root)

from demos.pipeline_fn import GuideFlow3dPipeline
from lib.util import common, generation, pointcloud
from demos.custom_utils import render_all_views
import third_party.TRELLIS.trellis.models as models

# Setup logging
log.getLogger().setLevel(log.INFO)
log.basicConfig(level=log.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Examples (matching run_gradio_demo.py)
STRUCT_EXAMPLES = {
    "bench_chair": "example_data/structure_mesh/bench_chair.glb",
    "cabinet": "example_data/structure_mesh/cabinet.glb",
    "chair": "example_data/structure_mesh/chair.glb",
    "giraffe": "example_data/structure_mesh/giraffe.glb",
    "motorcycle": "example_data/structure_mesh/motorcycle.glb",
    "plane": "example_data/structure_mesh/plane.glb",
}

APP_MESH_EXAMPLES = {
    "B01DA8LC0A": "example_data/appearance_mesh/B01DA8LC0A.glb",
    "B01DJH73Y6": "example_data/appearance_mesh/B01DJH73Y6.glb",
    "B0728KSP33": "example_data/appearance_mesh/B0728KSP33.glb",
    "B07B4YXNR8": "example_data/appearance_mesh/B07B4YXNR8.glb",
    "B07QC84LP1": "example_data/appearance_mesh/B07QC84LP1.glb",
    "B07QFRSC8M": "example_data/appearance_mesh/B07QFRSC8M_zup.glb",
    "B082QC7YKR": "example_data/appearance_mesh/B082QC7YKR_zup.glb",
}

def generate_structure_cache(pipe, name, path, output_base):
    log.info(f"Processing Structure Example: {name}")
    output_dir = os.path.join(output_base, name)
    os.makedirs(output_dir, exist_ok=True)
    
    # Run preprocess to generate struct_renders, voxels, partfield
    # preprocess(self, structure_mesh, convert_yup_to_zup, output_dir)
    pipe.preprocess(
        structure_mesh=path,
        convert_yup_to_zup=False, # Assuming defaults
        output_dir=output_dir
    )
    log.info(f"Completed Structure Example: {name}")

def generate_appearance_cache(pipe, name, path, output_base):
    log.info(f"Processing Appearance Example: {name}")
    output_dir = os.path.join(output_base, name)
    os.makedirs(output_dir, exist_ok=True)

    # Logic adapted from pipe.run_appearance but focused on Appearance Mesh only
    
    # 1. Load and Save Z-up mesh
    cached_app_input_path = os.path.join(output_dir, "app_mesh_input.glb")
    shutil.copy2(path, cached_app_input_path)
    
    app_mesh_path = os.path.join(output_dir, "app_mesh.glb")
    app_mesh_zup_path = os.path.join(output_dir, "app_mesh_zup.glb")
    
    app_mesh = trimesh.load(path, force="mesh")
    app_mesh.export(app_mesh_path)
    # convert_appearance_yup_to_zup=False (Default)
    app_mesh.export(app_mesh_zup_path)
    
    # 2. Render Views
    app_render_dir = os.path.join(output_dir, 'app_renders')
    common.ensure_dir(app_render_dir)
    
    # Only render if not empty
    if not os.listdir(app_render_dir):
        render_all_views(
            app_mesh_zup_path,
            app_render_dir,
            num_views=pipe.cfg.num_views,
            num_workers=None
        )

    # 3. Voxelize
    app_voxel_dir = os.path.join(output_dir, "voxels")
    common.ensure_dir(app_voxel_dir)
    app_voxels_path = os.path.join(app_voxel_dir, "app_voxels.ply")
    app_mesh_ply_path = os.path.join(app_render_dir, "mesh.ply") # Generated by render_all_views
    
    if not os.path.exists(app_voxels_path) and os.path.exists(app_mesh_ply_path):
        pointcloud.voxelize_mesh(
            app_mesh_ply_path,
            save_path=app_voxels_path,
        )

    # 4. DinoV2 Features
    features_dir = os.path.join(output_dir, "features", pipe.cfg.feature_name)
    common.ensure_dir(features_dir)
    
    if not os.listdir(features_dir):
        dinov2_model = torch.hub.load(pipe.cfg.dinov2_repo, pipe.cfg.feature_name)
        dinov2_model.eval().cuda()
        transform = transforms.Compose([transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
        generation.extract_feature(output_dir, dinov2_model, transform)
        del dinov2_model
        torch.cuda.empty_cache()

    # 5. SLAT Latent
    latents_dir = os.path.join(output_dir, "latents", pipe.cfg.latent_name)
    common.ensure_dir(latents_dir)
    
    if not os.listdir(latents_dir):
        encoder = models.from_pretrained(pipe.cfg.enc_pretrained).eval().cuda()
        generation.get_latent(output_dir, pipe.cfg.feature_name, pipe.cfg.latent_name, encoder)
        del encoder
        torch.cuda.empty_cache()

    # 6. PartField (Appearance)
    app_partfield_dir = os.path.join(output_dir, "partfield")
    common.ensure_dir(app_partfield_dir)
    
    existing_pf = [f for f in os.listdir(app_partfield_dir) if f.startswith("part_feat_app_mesh_zup")]
    if not existing_pf:
        from demos.pipeline_fn import partfield_pipeline_predict
        partfield_pipeline_predict(
            app_mesh_zup_path,
            app_partfield_dir,
            "app_mesh_zup"
        )
    
    log.info(f"Completed Appearance Example: {name}")

def main():
    cfg_path = os.path.join(project_root, 'config', 'default.yaml')
    cfg = OmegaConf.load(cfg_path)
    pipe = GuideFlow3dPipeline().from_pretrained(cfg)
    
    output_base = os.path.join(project_root, "all_outputs")
    
    for name, path in STRUCT_EXAMPLES.items():
        abs_path = os.path.join(project_root, path)
        generate_structure_cache(pipe, name, abs_path, output_base)
        
    for name, path in APP_MESH_EXAMPLES.items():
        abs_path = os.path.join(project_root, path)
        generate_appearance_cache(pipe, name, abs_path, output_base)

if __name__ == "__main__":
    main()