import os
import os.path as osp
import spaces
import gc
import trimesh
from PIL import Image
import logging as log
from omegaconf import OmegaConf
import random
import numpy as np
import hashlib
from typing import Optional

import torch
from torchvision import transforms
from pycg import vis, image
from pycg import render as pycg_render

import sys
sys.path.append('.')

from lib.util.render import BLENDER_PATH
from third_party.PartField.partfield.model_trainer_pvcnn_only_demo import Model
from lib.opt import appearance, self_similarity
from lib.util import generation, common, pointcloud
import third_party.TRELLIS.trellis.models as models
from demos.custom_utils import render_all_views

# Set BLENDER_HOME for pycg if not set
if "BLENDER_HOME" not in os.environ:
    if osp.exists(BLENDER_PATH):
        os.environ["BLENDER_HOME"] = BLENDER_PATH
    else:
        # Fallback to just 'blender' if path invalid, though this likely fails too if not in PATH
        os.environ["BLENDER_HOME"] = "blender"

log.getLogger().setLevel(log.INFO)
log.basicConfig(level=log.INFO,
                format='%(asctime)s - %(levelname)s - %(message)s',
                datefmt='%Y-%m-%d %H:%M:%S')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
partfield_config = 'third_party/PartField/config.yaml'
partfield_cfg = OmegaConf.load(partfield_config)

def file_sha256(path: str, chunk_size: int = 1 << 20) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(chunk_size), b""):
            h.update(chunk)
    return h.hexdigest()

# @spaces.GPU()
def init_partfield(obj_path):
    torch.manual_seed(0)
    random.seed(0)
    np.random.seed(0)

    partfield_model = Model(partfield_cfg, obj_path)
    partfield_model = partfield_model.to(device)

    ckpt = torch.load(partfield_cfg.continue_ckpt, map_location=device, weights_only=False)

    state_dict = ckpt.get("state_dict", ckpt)
    state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()}
    missing, unexpected = partfield_model.load_state_dict(state_dict, strict=False)

    if missing:
        print("[load_partfield_model] Missing keys:", missing)
    if unexpected:
        print("[load_partfield_model] Unexpected keys:", unexpected)

    partfield_model.eval()
    return partfield_model

@spaces.GPU
def partfield_pipeline_predict(obj_path, output_dir):
    
    log.info("Extracting PartField feature planes...")

    seed = int(partfield_cfg.seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    partfield_model = init_partfield(obj_path)
    dataloader = partfield_model.predict_dataloader()
    batch = next(iter(dataloader))

    with torch.no_grad():
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            batch = {
                k: (v.to(device) if torch.is_tensor(v) else v)
                for k, v in batch.items()
            }
            part_planes, uid = partfield_model.predict_step(batch, batch_idx=0)

    os.makedirs(output_dir, exist_ok=True)
    print("UID VALUE: ", uid)
    partfield_save_path = f'{output_dir}/part_feat_{uid}_batch_part_plane.npy'
    print("SAVING PART FIELD TO: ", partfield_save_path)
    np.save(partfield_save_path, part_planes)

    del partfield_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

    return partfield_save_path

class GuideFlow3dPipeline:
    def __init__(self):
        self.cfg = None

    def from_pretrained(self, config):
        self.cfg = config
        return self

    # @spaces.GPU(duration=360)
    def preprocess(
        self,
        structure_mesh: str,
        convert_yup_to_zup: bool,
        output_dir: str,
    ) -> None:
        log.info("Loading structure mesh...")

        if not structure_mesh.endswith('.glb'):
            log.error("Meshes must be in .glb format")
            return

        struct_hash_path = osp.join(output_dir, "struct_mesh.hash")
        current_struct_hash = file_sha256(structure_mesh)
        cached_struct_hash = None
        if osp.exists(struct_hash_path):
            with open(struct_hash_path, "r") as f:
                cached_struct_hash = f.read().strip()
        
        use_struct_cache = (cached_struct_hash == current_struct_hash)

        struct_mesh_path = structure_mesh
        struct_mesh_zup_path = osp.join(output_dir, "struct_mesh_zup.glb")

        if use_struct_cache and osp.exists(struct_mesh_zup_path):
            log.info("Using cached structure mesh (z-up).")
            struct_mesh = trimesh.load(struct_mesh_zup_path, force="mesh")
        else:
            struct_mesh = trimesh.load(structure_mesh, force='mesh')
            struct_mesh.export(struct_mesh_path)

            if convert_yup_to_zup:
                struct_mesh = pointcloud.convert_mesh_yup_to_zup(struct_mesh)
            struct_mesh.export(struct_mesh_zup_path)

            with open(struct_hash_path, "w") as f:
                f.write(current_struct_hash)

        if convert_yup_to_zup:
            struct_mesh = pointcloud.convert_mesh_yup_to_zup(struct_mesh)
        struct_mesh.export(osp.join(output_dir, 'struct_mesh_zup.glb'))

        log.info(f"Rendering structure mesh for {self.cfg.num_views // 10} views...")
        struct_render_dir = osp.join(output_dir, 'struct_renders')
        common.ensure_dir(struct_render_dir)

        struct_mesh_ply_path = osp.join(struct_render_dir, "mesh.ply")

        if use_struct_cache and osp.exists(struct_mesh_ply_path):
            log.info("Using cached structure renders.")
            out_renderviews = sorted(
                [
                    osp.join(struct_render_dir, f)
                    for f in os.listdir(struct_render_dir)
                    if f.lower().endswith((".png", ".jpg", ".jpeg"))
                ]
            )
        else:
            out_renderviews = render_all_views(
                struct_mesh_zup_path,
                struct_render_dir,
                num_views=self.cfg.num_views // 10,
                num_workers=None # Let custom_utils decide best worker count
            )

        if not out_renderviews:
            log.error("Structure rendering failed! Aborting pipeline.")
            return None

        voxel_dir = osp.join(output_dir, 'voxels')
        common.ensure_dir(voxel_dir)
        log.info("Voxelizing structure mesh...")
        struct_voxels_path = osp.join(voxel_dir, "struct_voxels.ply")

        if use_struct_cache and osp.exists(struct_voxels_path):
            log.info("Using cached structure voxels.")
        else:
            pointcloud.voxelize_mesh(
                struct_mesh_ply_path,
                save_path=struct_voxels_path,
            )

        log.info("Extracting Structure Mesh PartField feature planes...")
        partfield_dir = osp.join(output_dir, 'partfield')
        common.ensure_dir(partfield_dir)

        existing = [
            f for f in os.listdir(partfield_dir)
            if f.startswith("part_feat_struct_mesh_zup") and f.endswith("_batch_part_plane.npy")
        ]
        if use_struct_cache and existing:
            partfield_save_path = osp.join(partfield_dir, existing[0])
            log.info(f"Using cached Structure PartField at {partfield_save_path}")
        else:
            print("PREDICTING STRUCTURE PART FIELD...")
            partfield_save_path = partfield_pipeline_predict(
                struct_mesh_zup_path,
                partfield_dir,
            )

        if not out_renderviews:
            log.info("Structure rendering failed!")

        return {
            "struct_mesh": struct_mesh,
            "render_out": out_renderviews,
            "partfield_structure_predictions_save_path": partfield_save_path,
            "voxel_dir": voxel_dir
        }

    @spaces.GPU(duration=120)
    def run_appearance(
        self,
        structure_mesh: str,
        convert_target_yup_to_zup: bool,
        convert_appearance_yup_to_zup: bool,
        output_dir: str,
        appearance_mesh: str,
        appearance_image: str,
    ) -> Optional[str]:
        _ = self.preprocess(
            structure_mesh=structure_mesh,
            convert_yup_to_zup=convert_target_yup_to_zup,
            output_dir=output_dir,
        )

        app_hash_path = osp.join(output_dir, "app_mesh.hash")
        current_app_hash = file_sha256(appearance_mesh)
        cached_app_hash = None
        if osp.exists(app_hash_path):
            with open(app_hash_path, "r") as f:
                cached_app_hash = f.read().strip()
        use_app_cache = (cached_app_hash == current_app_hash)

        blender_cache_dir = osp.join(output_dir, "blender_cache")
        os.makedirs(blender_cache_dir, exist_ok=True)
        os.environ["XDG_CACHE_HOME"] = blender_cache_dir

        log.info("Running appearance-guided optimization...")
        
        # Load appearance mesh
        log.info("Loading appearance mesh...")
        
        if not appearance_mesh.endswith('.glb'):
            log.error("Meshes must be in .glb format")
            return None
        
        if not osp.exists(appearance_mesh):
            log.error(f"Appearance mesh not found: {appearance_mesh}")
            return None

        app_mesh_path = osp.join(output_dir, "app_mesh.glb")
        app_mesh_zup_path = osp.join(output_dir, "app_mesh_zup.glb")

        if use_app_cache and osp.exists(app_mesh_zup_path):
            log.info("Using cached appearance mesh (z-up).")
            app_mesh = trimesh.load(app_mesh_zup_path, force="mesh")
        else:
            app_mesh = trimesh.load(appearance_mesh, force="mesh")
            app_mesh.export(app_mesh_path)

            if convert_appearance_yup_to_zup:
                app_mesh = pointcloud.convert_mesh_yup_to_zup(app_mesh)
            app_mesh.export(app_mesh_zup_path)

            with open(app_hash_path, "w") as f:
                f.write(current_app_hash)

        # Load appearance image
        log.info("Loading appearance image...")
        if appearance_image:
            app_image = Image.open(appearance_image).convert('RGB')
            app_image.save(osp.join(output_dir, 'app_image.png'))
        else:
            mesh = vis.from_file(osp.join(output_dir, 'app_mesh.glb'), load_obj_textures=True)
            mesh.paint_uniform_color([0.5, 0.5, 0.5])
            scene = pycg_render.Scene(up_axis='+Y')
            scene.add_object(mesh)
            scene.quick_camera(w=512, h=512, pitch_angle=30, plane_angle=-45.0, fov=40)
            pycg_render.ThemeDiffuseShadow(None, sun_tilt_right=0.0, sun_tilt_back=0.0, sun_angle=60.0).apply_to(scene)
            rendering = scene.render_blender(quality=512)
            rendering = image.alpha_compositing(rendering, image.solid(rendering.shape[1], rendering.shape[0]))
            image.write(osp.join(output_dir, 'app_image.png'), rendering)

        # Render views for DinoV2 feature extraction
        log.info(f"Rendering appearance mesh for {self.cfg.num_views} views...")
        app_render_dir = osp.join(output_dir, 'app_renders')
        common.ensure_dir(app_render_dir)
        app_mesh_ply_path = osp.join(app_render_dir, "mesh.ply")

        if use_app_cache and osp.exists(app_mesh_ply_path):
            log.info("Using cached appearance renders.")
            out_renderviews = sorted(
                [
                    osp.join(app_render_dir, f)
                    for f in os.listdir(app_render_dir)
                    if f.lower().endswith((".png", ".jpg", ".jpeg"))
                ]
            )
        else:
            out_renderviews = render_all_views(
                app_mesh_zup_path,
                app_render_dir,
                num_views=self.cfg.num_views,
                num_workers=None # Let custom_utils decide best worker count
            )

        if not out_renderviews:
            log.info("Appearance rendering failed!")
            return None

        # Voxelise mesh
        log.info("Voxelizing appearance mesh...")
        app_voxel_dir = osp.join(output_dir, "voxels")
        common.ensure_dir(app_voxel_dir)
        app_voxels_path = osp.join(app_voxel_dir, "app_voxels.ply")

        if use_app_cache and osp.exists(app_voxels_path):
            log.info("Using cached appearance voxels.")
        else:
            pointcloud.voxelize_mesh(
                app_mesh_ply_path,
                save_path=app_voxels_path,
            )

        # Extract DinoV2 Features
        log.info("Extracting DinoV2 features...")
        features_dir = osp.join(output_dir, "features", self.cfg.feature_name)
        common.ensure_dir(features_dir)

        if use_app_cache and os.listdir(features_dir):
            log.info("Using cached DINOv2 features.")
        else:
            log.info("Extracting DinoV2 features...")
        
            dinov2_model = torch.hub.load(self.cfg.dinov2_repo, self.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)
            torch.cuda.empty_cache()

            del dinov2_model
            gc.collect() # Free up memory

        # Extract SLAT Latent
        log.info("Extracting SLAT latent...")
        latents_dir = osp.join(output_dir, "latents", self.cfg.latent_name)
        common.ensure_dir(latents_dir)

        if use_app_cache and os.listdir(latents_dir):
            log.info("Using cached SLAT latent.")
        else:
            log.info("Extracting SLAT latent...")
            encoder = models.from_pretrained(self.cfg.enc_pretrained).eval().cuda()

            generation.get_latent(output_dir, self.cfg.feature_name, self.cfg.latent_name, encoder)

            del encoder
            gc.collect() # Free up memory
        
        # Extract PartField features for appearance mesh
        log.info("Extracting Appearance Mesh PartField feature planes...")
        app_partfield_dir = osp.join(output_dir, "partfield")
        common.ensure_dir(app_partfield_dir)

        existing_app_pf = [
            f for f in os.listdir(app_partfield_dir)
            if f.startswith("part_feat_app_mesh_zup") and f.endswith("_batch_part_plane.npy")
        ]
        if use_app_cache and existing_app_pf:
            appearance_partfield_save_path = osp.join(
                app_partfield_dir, existing_app_pf[0]
            )
            log.info(
                f"Using cached Appearance PartField at {appearance_partfield_save_path}"
            )
        else:
            appearance_partfield_save_path = partfield_pipeline_predict(
                app_mesh_zup_path,
                app_partfield_dir,
            )
    
        # Appearance Optimization
        appearance.optimize_appearance(self.cfg, output_dir)
        
        # Return the output mesh path
        output_mesh_path = osp.join(output_dir, 'out_app.glb')
        output_video_path = osp.join(output_dir, 'out_gaussian_app.mp4')
        if not osp.exists(output_mesh_path) or not osp.exists(output_video_path):
            log.error(f"Output mesh or video not found at {output_mesh_path} or {output_video_path}")
            return None, None
        return output_mesh_path, output_video_path

    @spaces.GPU(duration=120)
    def run_self_similarity(
        self,
        structure_mesh: str,
        convert_target_yup_to_zup: bool,
        output_dir: str,
        appearance_text: str,
    ) -> Optional[str]:
        _ = self.preprocess(
            structure_mesh=structure_mesh,
            convert_yup_to_zup=convert_target_yup_to_zup,
            output_dir=output_dir,
        )
        log.info("Running similarity-guided optimization...")

        # Self-Similarity Optimization
        self_similarity.optimize_self_similarity(self.cfg, appearance_text, output_dir)
        
        # Return the output mesh path
        output_mesh_path = osp.join(output_dir, 'out_sim.glb')
        output_video_path = osp.join(output_dir, 'out_gaussian_sim.mp4')
        if not osp.exists(output_mesh_path) or not osp.exists(output_video_path):
            log.error(f"Output mesh or video not found at {output_mesh_path} or {output_video_path}")
            return None, None
        return output_mesh_path, output_video_path

def main():
    args = {
        "structure_mesh": os.path.join(os.getcwd(), "structure_mesh.glb"),
        "output_dir": os.path.join(os.getcwd(), "all_outputs", "pipeline_outputs"),
        "convert_target_yup_to_zup": True,
        "convert_appearance_yup_to_zup": True,
        "appearance_mesh": os.path.join(os.getcwd(), "appearance_mesh.glb"),
        "appearance_image": os.path.join(os.getcwd(), "appearance_image.jpg"),
        "appearance_text": "",
    }

    cfg = OmegaConf.load('config/default.yaml')
    
    common.ensure_dir(args["output_dir"])

    pipe = GuideFlow3dPipeline.from_pretrained(cfg)

    if args["guidance_mode"] == "appearance":
        out = pipe.run_appearance(
            **args
        )
    else:
        out = pipe.run_self_similarity(
            **args
        )