"""Hybrid texture generation pipeline: MLX renderer + PyTorch ML models.

Uses MeshRenderMLX for all rendering and texture baking (Metal GPU),
while keeping multiview diffusion and super-resolution in PyTorch.
The interface between them is PIL Images.
"""

import copy
import os
import warnings

import numpy as np
import trimesh
from PIL import Image
from typing import List

from DifferentiableRenderer.mesh_render_mlx import MeshRenderMLX
from utils.pipeline_utils_mlx import ViewProcessorMLX


def _convert_obj_to_glb(obj_path: str, glb_path: str) -> bool:
    """Convert OBJ + MTL + textures to GLB.

    Tries Blender (mesh_utils.convert_obj_to_glb) first for highest fidelity,
    falls back to trimesh (pure Python, no CUDA/Blender dependency) so the
    pipeline runs on Apple Silicon without `bpy`.
    """
    try:
        from DifferentiableRenderer.mesh_utils import convert_obj_to_glb
        return convert_obj_to_glb(obj_path, glb_path)
    except ImportError:
        pass

    import trimesh
    mesh = trimesh.load(obj_path, process=False)
    mesh.export(glb_path)
    return os.path.exists(glb_path)
try:
    from utils.simplify_mesh_utils import remesh_mesh
except ImportError:
    remesh_mesh = None  # pymeshlab unavailable; require use_remesh=False
try:
    from utils.uvwrap_utils import mesh_uv_wrap
except ImportError:
    mesh_uv_wrap = None  # xatlas unavailable; require pre-wrapped mesh

warnings.filterwarnings("ignore")


class Hunyuan3DPaintConfigMLX:
    def __init__(self, max_num_view, resolution):
        self.device = "cuda"  # for PyTorch ML models only

        self.multiview_cfg_path = "hy3dpaint/cfgs/hunyuan-paint-pbr.yaml"
        self.custom_pipeline = "hunyuanpaintpbr"
        self.multiview_pretrained_path = "tencent/Hunyuan3D-2.1"
        self.dino_ckpt_path = "facebook/dinov2-giant"
        self.realesrgan_ckpt_path = "ckpt/RealESRGAN_x4plus.pth"

        # Fully-MLX diffusion (Apple Silicon, no CUDA).
        # When True, loads HunyuanPaintModelMLX from mlx_weights_source and
        # skips the PyTorch multiview model / super-resolution.
        #
        # mlx_weights_source may be either a HuggingFace repo ID (auto-
        # downloaded via huggingface_hub) or an absolute local path. The
        # env var HUNYUAN3D_MLX_WEIGHTS_DIR overrides this value.
        self.use_mlx_diffusion = True
        self.mlx_weights_source = "dgrauet/hunyuan3d-2.1-mlx"
        self.mlx_num_inference_steps = 15
        self.mlx_guidance_scale = 3.0
        self.mlx_seed = 42

        # Fully-MLX super-resolution (RealESRGAN x4, ~17M params).
        # Upscales each 512^2 generated view to 2048^2 before baking so
        # the UV atlas keeps detail. Disable for faster iteration.
        self.use_mlx_super_res = True
        self.mlx_super_res_filename = "realesrgan_x4plus.safetensors"

        self.raster_mode = "mlx"
        self.bake_mode = "back_sample"
        self.render_size = 1024 * 2
        self.texture_size = 1024 * 4
        self.max_selected_view_num = max_num_view
        self.resolution = resolution
        self.bake_exp = 4
        self.merge_method = "fast"

        # View selection candidates
        self.candidate_camera_azims = [0, 90, 180, 270, 0, 180]
        self.candidate_camera_elevs = [0, 0, 0, 0, 90, -90]
        self.candidate_view_weights = [1, 0.1, 0.5, 0.1, 0.05, 0.05]

        for azim in range(0, 360, 30):
            self.candidate_camera_azims.append(azim)
            self.candidate_camera_elevs.append(20)
            self.candidate_view_weights.append(0.01)

            self.candidate_camera_azims.append(azim)
            self.candidate_camera_elevs.append(-20)
            self.candidate_view_weights.append(0.01)


class Hunyuan3DPaintPipelineMLX:
    """Texture generation pipeline with MLX rendering backend.

    The rendering, rasterization, and texture baking run on Apple Silicon
    via MLX Metal kernels. The ML models (multiview diffusion, super-res)
    remain in PyTorch and communicate via PIL Images.
    """

    def __init__(self, config=None):
        self.config = config if config is not None else Hunyuan3DPaintConfigMLX(
            max_num_view=6, resolution=1024,
        )
        self.models = {}

        self.render = MeshRenderMLX(
            default_resolution=self.config.render_size,
            texture_size=self.config.texture_size,
        )
        self.view_processor = ViewProcessorMLX(self.config, self.render)
        self.load_models()

    def load_models(self):
        """Load diffusion model.

        MLX path: HunyuanPaintModelMLX on Apple Silicon.
        PyTorch path: diffusion + super-res on CUDA (legacy).
        """
        if self.config.use_mlx_diffusion:
            from hunyuanpaintpbr_mlx.load_model import HunyuanPaintModelMLX
            self.models["mlx_model"] = HunyuanPaintModelMLX.from_pretrained(
                self.config.mlx_weights_source,
            )
            print("MLX diffusion model loaded.")

            if getattr(self.config, "use_mlx_super_res", False):
                from utils.image_super_utils_mlx import imageSuperNetMLX
                sr_path = self._resolve_super_res_weights()
                self.models["mlx_super"] = imageSuperNetMLX(sr_path)
                print(f"MLX super-resolution loaded ({sr_path}).")
            return

        try:
            import torch
            torch.cuda.empty_cache()
            from utils.image_super_utils import imageSuperNet
            from utils.multiview_utils import multiviewDiffusionNet
            self.models["super_model"] = imageSuperNet(self.config)
            self.models["multiview_model"] = multiviewDiffusionNet(self.config)
            print("ML models loaded (PyTorch).")
        except (ImportError, RuntimeError) as e:
            print(f"ML models not loaded ({e}). "
                  "Pipeline can still run render/bake with pre-generated images.")

    def _resolve_super_res_weights(self) -> str:
        """Resolve RealESRGAN MLX weights path.

        Reuses the HF repo (or local dir) used for the diffusion weights so a
        single ``HUNYUAN3D_MLX_WEIGHTS_DIR`` override covers everything.
        """
        src = self.config.mlx_weights_source
        filename = self.config.mlx_super_res_filename
        env_override = os.environ.get("HUNYUAN3D_MLX_WEIGHTS_DIR")
        if env_override:
            return os.path.join(env_override, filename)
        if os.path.isdir(src):
            return os.path.join(src, filename)
        from huggingface_hub import hf_hub_download
        return hf_hub_download(repo_id=src, filename=filename)

    def _run_multiview_mlx(self, image_style, normal_maps, position_maps,
                            camera_azims):
        """Run MLX multiview diffusion.

        Args:
            image_style: list of PIL reference images (we use only the first).
            normal_maps: list of PIL normal maps (one per view).
            position_maps: list of PIL position maps.
            camera_azims: list of azimuth angles.

        Returns:
            dict with 'albedo' and 'mr', each a list of PIL images.
        """
        from hunyuanpaintpbr_mlx.inference import generate_multiview_pbr

        ref_np = np.array(image_style[0].convert("RGB"))

        def _pil_to_np_f32(img):
            """Convert a PIL image to (H, W, 3) float32 in [0, 1]."""
            arr = np.array(img.convert("RGB"), dtype=np.float32) / 255.0
            return arr

        n_maps = [_pil_to_np_f32(m) for m in normal_maps]
        p_maps = [_pil_to_np_f32(m) for m in position_maps]

        result = generate_multiview_pbr(
            model=self.models["mlx_model"],
            normal_maps=n_maps,
            position_maps=p_maps,
            reference_image=ref_np,
            camera_azims=camera_azims,
            num_inference_steps=self.config.mlx_num_inference_steps,
            guidance_scale=self.config.mlx_guidance_scale,
            view_size=self.config.resolution,
            seed=self.config.mlx_seed,
        )

        return {
            "albedo": [Image.fromarray(a) for a in result["albedo"]],
            "mr": [Image.fromarray(m) for m in result["mr"]],
        }

    def __call__(self, mesh_path=None, image_path=None,
                 output_mesh_path=None, use_remesh=True, save_glb=True):
        """Generate texture for 3D mesh using multiview diffusion.

        Args:
            mesh_path: Path to input mesh.
            image_path: Path to reference image or PIL Image.
            output_mesh_path: Output path for textured mesh.
            use_remesh: Whether to remesh before texturing.
            save_glb: Whether to also save as GLB.

        Returns:
            Path to the output mesh file.
        """
        # --- Prepare image prompt ---
        if isinstance(image_path, str):
            image_prompt = Image.open(image_path)
        elif isinstance(image_path, Image.Image):
            image_prompt = image_path
        else:
            image_prompt = image_path

        if not isinstance(image_prompt, list):
            image_prompt = [image_prompt]

        # --- Process mesh ---
        path = os.path.dirname(mesh_path)
        if use_remesh:
            processed_mesh_path = os.path.join(path, "white_mesh_remesh.obj")
            remesh_mesh(mesh_path, processed_mesh_path)
        else:
            processed_mesh_path = mesh_path

        if output_mesh_path is None:
            output_mesh_path = os.path.join(path, "textured_mesh.obj")

        mesh = trimesh.load(processed_mesh_path)
        # trimesh.load can return a Scene for .glb/.gltf; flatten to a single
        # Trimesh so we can inspect UVs (a Scene wrapper hides them).
        if isinstance(mesh, trimesh.Scene):
            mesh = mesh.dump(concatenate=True)

        # Only UV-wrap if the mesh has no UVs (xatlas creates fragmented
        # islands; a hand-authored UV layout produces much cleaner textures).
        uv = getattr(getattr(mesh, "visual", None), "uv", None)
        needs_wrap = uv is None or len(uv) == 0
        if needs_wrap:
            if mesh_uv_wrap is None:
                raise RuntimeError(
                    "Mesh has no UVs and `xatlas` is not installed. "
                    "Install xatlas or provide a pre-wrapped mesh."
                )
            mesh = mesh_uv_wrap(mesh)
        self.render.load_mesh(mesh=mesh)

        # --- View selection ---
        selected_elevs, selected_azims, selected_weights = \
            self.view_processor.bake_view_selection(
                self.config.candidate_camera_elevs,
                self.config.candidate_camera_azims,
                self.config.candidate_view_weights,
                self.config.max_selected_view_num,
            )

        # --- Render reference maps (MLX, Metal GPU) ---
        normal_maps = self.view_processor.render_normal_multiview(
            selected_elevs, selected_azims, use_abs_coor=True,
        )
        position_maps = self.view_processor.render_position_multiview(
            selected_elevs, selected_azims,
        )

        # --- Style prep ---
        image_style = []
        for image in image_prompt:
            image = image.resize((512, 512))
            if image.mode == "RGBA":
                white_bg = Image.new("RGB", image.size, (255, 255, 255))
                white_bg.paste(image, mask=image.getchannel("A"))
                image = white_bg
            image_style.append(image.convert("RGB"))

        # --- Multiview diffusion ---
        if self.config.use_mlx_diffusion:
            multiviews_pbr = self._run_multiview_mlx(
                image_style, normal_maps, position_maps, selected_azims,
            )
            enhance_images = {
                "albedo": list(multiviews_pbr["albedo"]),
                "mr": list(multiviews_pbr["mr"]),
            }

            # RealESRGAN x4 super-resolution per view (MLX)
            if "mlx_super" in self.models:
                sr = self.models["mlx_super"]
                for i in range(len(enhance_images["albedo"])):
                    enhance_images["albedo"][i] = sr(enhance_images["albedo"][i])
                    enhance_images["mr"][i] = sr(enhance_images["mr"][i])
        else:
            if "multiview_model" not in self.models:
                raise RuntimeError(
                    "Multiview diffusion model not loaded. "
                    "Install PyTorch with CUDA support."
                )
            multiviews_pbr = self.models["multiview_model"](
                image_style,
                normal_maps + position_maps,
                prompt="high quality",
                custom_view_size=self.config.resolution,
                resize_input=True,
            )

            # Super-resolution (CUDA-only)
            enhance_images = {
                "albedo": copy.deepcopy(multiviews_pbr["albedo"]),
                "mr": copy.deepcopy(multiviews_pbr["mr"]),
            }
            if "super_model" in self.models:
                for i in range(len(enhance_images["albedo"])):
                    enhance_images["albedo"][i] = self.models["super_model"](
                        enhance_images["albedo"][i]
                    )
                    enhance_images["mr"][i] = self.models["super_model"](
                        enhance_images["mr"][i]
                    )

        # --- Texture baking (MLX, Metal GPU) ---
        render_size = self.config.render_size
        for i in range(len(enhance_images["albedo"])):
            enhance_images["albedo"][i] = enhance_images["albedo"][i].resize(
                (render_size, render_size),
            )
            enhance_images["mr"][i] = enhance_images["mr"][i].resize(
                (render_size, render_size),
            )

        texture, mask = self.view_processor.bake_from_multiview(
            enhance_images["albedo"],
            selected_elevs, selected_azims, selected_weights,
        )
        mask_np = (mask[..., 0] * 255).astype(np.uint8)

        texture_mr, mask_mr = self.view_processor.bake_from_multiview(
            enhance_images["mr"],
            selected_elevs, selected_azims, selected_weights,
        )
        mask_mr_np = (mask_mr[..., 0] * 255).astype(np.uint8)

        # --- Inpainting (OpenCV, CPU) ---
        texture = self.view_processor.texture_inpaint(texture, mask_np)
        self.render.set_texture(texture)

        if "mr" in enhance_images:
            texture_mr = self.view_processor.texture_inpaint(
                texture_mr, mask_mr_np,
            )
            self.render.set_texture_mr(texture_mr)

        # --- Save ---
        # Do NOT downsample: cv2.resize blends across UV gutters and produces
        # visible color bleeding at island edges in the textured mesh.
        self.render.save_mesh(output_mesh_path, downsample=False)

        if save_glb:
            glb_path = output_mesh_path.replace(".obj", ".glb")
            # Export directly from the in-memory trimesh so the PBRMaterial
            # (metallicRoughnessTexture, doubleSided, sRGB baseColor) is
            # preserved. OBJ round-trip would drop all of it.
            in_mem = getattr(self.render, "_last_exported_mesh", None)
            if in_mem is not None:
                in_mem.export(glb_path)
            else:
                _convert_obj_to_glb(output_mesh_path, glb_path)

        return output_mesh_path
