import os
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from spandrel import ModelLoader

from MV_Adapter.mvadapter.utils import image_to_tensor, make_image_grid, tensor_to_image
from MV_Adapter.mvadapter.utils.mesh_utils import (
    Camera,
    CameraProjection,
    NVDiffRastContextWrapper,
    SmartPainter,
    TexturedMesh,
    get_camera,
    get_orthogonal_camera,
    load_mesh,
    render,
    replace_mesh_texture_and_save,
)
from MV_Adapter.mvadapter.utils.mesh_utils.mesh_process import process_raw

# Import dilate function to remove texture seams
def torch_image_dilate(img):
    """
    Remove thin seams on generated texture
        img (torch): [H, W, C]
    """
    import cv2
    img = np.asarray(img.cpu().numpy(), dtype=np.float32)
    img = img * 255
    img = img.clip(0, 255)
    mask = np.sum(img.astype(np.float32), axis=-1, keepdims=True)
    mask = (mask <= 3.0).astype(np.float32)
    kernel = np.ones((3, 3), 'uint8')
    dilate_img = cv2.dilate(img, kernel, iterations=1)
    img = img * (1 - mask) + dilate_img * mask
    img = (img.clip(0, 255) / 255).astype(np.float32)
    return torch.from_numpy(img)


def clear():
    torch.cuda.empty_cache()


@contextmanager
def mesh_use_texture(mesh: TexturedMesh, texture: torch.FloatTensor):
    texture_ = mesh.texture
    mesh.texture = texture
    try:
        yield
    finally:
        mesh.texture = texture_


@dataclass
class ModProcessConfig:
    view_upscale: bool = False
    view_upscale_factor: int = 2
    inpaint_mode: str = "uv"  # in ["none", "uv", "view"]
    view_inpaint_max_view_score_thresh: float = 0.02
    view_inpaint_min_rounds: int = 4
    view_inpaint_max_rounds: int = 8
    view_inpaint_uv_padding_end: bool = True


@dataclass
class TexturePipelineOutput:
    shaded_model_save_path: Optional[str] = None
    pbr_model_save_path: Optional[str] = None


class TexturePipeline:
    def __init__(self, upscaler_ckpt_path: str, inpaint_ckpt_path: str, device: str):
        self.device = device
        self.ctx = NVDiffRastContextWrapper(device=self.device)
        self.cam_proj = CameraProjection(
            pb_backend="torch-cuda", bg_remover=None, device=self.device
        )
        if upscaler_ckpt_path is not None:
            self.upscaler = ModelLoader().load_from_file(upscaler_ckpt_path)
            self.upscaler.to(self.device).eval().half()
        if inpaint_ckpt_path is not None:
            self.inpainter = ModelLoader().load_from_file(inpaint_ckpt_path)
            self.inpainter.to(self.device).eval()

        self.smart_painter = SmartPainter(self.device)

    def load_packed_images(self, packed_image_path: Optional[str]) -> List[Image.Image]:
        if packed_image_path is None:
            return None
        packed_image = Image.open(packed_image_path)
        images = np.array_split(np.array(packed_image), 6, axis=1)
        images = [Image.fromarray(im) for im in images]
        return images

    def maybe_upscale_image(
        self,
        tensor: Optional[torch.FloatTensor],
        upscale: bool,
        upscale_factor: int,
        batched: bool = False,
    ) -> Optional[torch.FloatTensor]:
        if upscale:
            with torch.no_grad():
                tensor = tensor.permute(0, 3, 1, 2)
                if batched:
                    tensor = self.upscaler(tensor.half()).float()
                else:
                    tensor = torch.concat(
                        [
                            self.upscaler(im.unsqueeze(0).half()).float()
                            for im in tensor
                        ],
                        dim=0,
                    )
                tensor = tensor.clamp(0, 1).permute(0, 2, 3, 1)
            clear()
        return tensor

    def view_inpaint(
        self,
        mod_name: str,
        mesh: TexturedMesh,
        uv_proj: torch.FloatTensor,
        uv_valid_mask: torch.BoolTensor,
        config: ModProcessConfig,
        debug_dir: Optional[str] = None,
    ) -> Tuple[torch.FloatTensor, torch.BoolTensor]:
        def inpaint_func(
            image: torch.FloatTensor, mask: torch.FloatTensor
        ) -> torch.FloatTensor:
            with torch.no_grad():
                return self.inpainter(image.to(torch.float32), mask.to(torch.float32))

        return self.smart_painter(
            mod_name,
            mesh,
            inpaint_func,
            uv_proj,
            ~uv_valid_mask,
            max_view_score_thresh=config.view_inpaint_max_view_score_thresh,
            min_rounds=config.view_inpaint_min_rounds,
            max_rounds=config.view_inpaint_max_rounds,
            uv_padding_end=config.view_inpaint_uv_padding_end,
            debug_dir=debug_dir,
            debug_visualize_details=False,
        )

    def __call__(
        self,
        mesh_path: str,
        save_dir: str,
        save_name: str = "default",
        # mesh load settings
        move_to_center: bool = False,
        front_x: bool = True,
        # uv unwarp
        uv_unwarp: bool = False,
        preprocess_mesh: bool = False,
        # projection
        uv_size: int = 4096,
        # modes
        rgb_path: Optional[str] = None,
        rgb_process_config: ModProcessConfig = ModProcessConfig(),
        base_color_path: Optional[str] = None,
        base_color_process_config: ModProcessConfig = ModProcessConfig(),
        orm_path: Optional[str] = None,
        orm_process_config: ModProcessConfig = ModProcessConfig(),
        normal_path: Optional[str] = None,
        normal_strength: float = 1.0,
        normal_process_config: ModProcessConfig = ModProcessConfig(),
        # inpaint
        uv_inpaint_use_network: bool = False,
        view_inpaint_include_occlusion_boundary: bool = False,
        poisson_reprojection: bool = False,
        # camera
        camera_projection_type: str = "ORTHO",
        camera_elevation_deg: List[float] = [0, 0, 0, 0, 89.99, -89.99],
        camera_azimuth_deg: List[float] = [0, 90, 180, 270, 180, 180],
        camera_distance: float = 1.0,
        camera_ortho_scale: float = 1.1,
        camera_fov_deg: float = 40,
        # debug
        debug_mode: bool = False,
        # seam removal
        apply_dilate: bool = True,
    ):
        clear()

        if debug_mode:
            debug_dir = os.path.join(save_dir, "debug")
            os.makedirs(debug_dir, exist_ok=True)

        if uv_unwarp:
            file_suffix = os.path.splitext(mesh_path)[-1]
            mesh_path_new = mesh_path.replace(file_suffix, f"_unwarp{file_suffix}")
            process_raw(mesh_path, mesh_path_new, preprocess=preprocess_mesh)
            mesh_path = mesh_path_new

        mesh: TexturedMesh = load_mesh(
            mesh_path,
            rescale=True,
            move_to_center=move_to_center,
            front_x_to_y=front_x,
            default_uv_size=uv_size,
            device=self.device,
        )

        # projection
        if camera_projection_type == "PERSP":
            raise NotImplementedError
        elif camera_projection_type == "ORTHO":
            cameras = get_orthogonal_camera(
                elevation_deg=camera_elevation_deg,
                distance=[camera_distance] * 6,
                left=-camera_ortho_scale / 2,
                right=camera_ortho_scale / 2,
                bottom=-camera_ortho_scale / 2,
                top=camera_ortho_scale / 2,
                azimuth_deg=[x - 90 for x in camera_azimuth_deg],  # -y as front
                device=self.device,
            )

        mod_kwargs = {
            "rgb": (rgb_path, rgb_process_config),
            "base_color": (base_color_path, base_color_process_config),
            "orm": (orm_path, orm_process_config),
            "normal": (normal_path, normal_process_config),
        }
        mod_uv_image, mod_uv_tensor = {}, {}
        for mod_name, (mod_path, mod_process_config) in mod_kwargs.items():
            if mod_path is None:
                mod_uv_image[mod_name] = None
                mod_uv_tensor[mod_name] = None
                continue
            mod_images = self.load_packed_images(mod_path)
            mod_tensor = image_to_tensor(mod_images, device=self.device)
            mod_tensor = self.maybe_upscale_image(
                mod_tensor,
                mod_process_config.view_upscale,
                mod_process_config.view_upscale_factor,
            )
            if mod_process_config.view_upscale and debug_mode:
                make_image_grid(tensor_to_image(mod_tensor, batched=True), rows=1).save(
                    os.path.join(debug_dir, f"{mod_name}_upscaled.jpg")
                )

            if mod_name == "normal":
                _, height, width, _ = mod_tensor.shape
                render_out = render(
                    self.ctx,
                    mesh,
                    cameras,
                    height=height,
                    width=width,
                    render_attr=False,
                    render_depth=False,
                    render_normal=True,
                    render_tangent=True,
                )

                # compute UV tangent space
                vN = render_out.normal
                vT = render_out.tangent
                vB = torch.cross(vN, vT, dim=-1)
                tang_space = F.normalize(torch.stack([vT, vB, vN], dim=-2), dim=-1)

                # compute geometry tangent space
                vGN = vN
                vGT = torch.as_tensor(
                    [
                        [1, 0, 0],
                        [0, 1, 0],
                        [-1, 0, 0],
                        [0, -1, 0],
                        [-1, 0, 0],
                        [-1, 0, 0],
                    ],
                    dtype=torch.float32,
                    device=self.device,
                )[:, None, None, :]
                vGB = torch.cross(vGN, vGT, dim=-1)
                vGT = torch.cross(vGB, vGN, dim=-1)
                geo_tang_space = F.normalize(
                    torch.stack([vGT, vGB, vGN], dim=-2), dim=-1
                )

                # restore world-space normal from geometry tangent space
                mod_tensor = mod_tensor * 2 - 1
                mod_tensor = F.normalize(
                    (
                        mod_tensor[:, :, :, None, :]
                        * geo_tang_space.permute(0, 1, 2, 4, 3)
                    ).sum(-1),
                    dim=-1,
                )

                # bake world-space normal to UV tangent space
                mod_tensor = F.normalize(
                    (mod_tensor[:, :, :, None, :] * tang_space).sum(-1), dim=-1
                )
                mod_tensor = (mod_tensor * 0.5 + 0.5).clamp(0, 1)

                uv_proj, uv_valid_mask = self.cam_proj(
                    mod_tensor,
                    mesh,
                    cameras,
                    from_scratch=mod_process_config.inpaint_mode != "none",
                    poisson_blending=False,
                    depth_grad_dilation=5,
                    uv_exp_blend_alpha=3,
                    uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1, 1, 1]),
                    aoi_cos_valid_threshold=0.2,
                    uv_size=uv_size,
                    return_uv_projection_mask=True,
                )
                uv_proj[~uv_valid_mask] = torch.as_tensor([0.5, 0.5, 1]).to(uv_proj)
            else:
                # TODO: tweak depth_grad_dilation
                cam_proj_out = self.cam_proj(
                    mod_tensor,
                    mesh,
                    cameras,
                    from_scratch=mod_process_config.inpaint_mode != "none",
                    poisson_blending=False,
                    depth_grad_dilation=5,
                    depth_grad_threshold=0.1,
                    uv_exp_blend_alpha=3,
                    uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1, 1, 1]),
                    aoi_cos_valid_threshold=0.2,
                    uv_size=uv_size,
                    uv_padding=not uv_inpaint_use_network,
                    return_dict=True,
                )
                uv_proj, uv_valid_mask, uv_depth_grad = (
                    cam_proj_out.uv_proj,
                    cam_proj_out.uv_proj_mask,
                    cam_proj_out.uv_depth_grad,
                )
                if uv_inpaint_use_network:
                    uv_inpaint_mask_input = 1 - uv_valid_mask[None, None].float()
                    uv_inpaint_image_input = uv_proj[None].permute(0, 3, 1, 2)
                    with torch.no_grad():
                        uv_inpaint_result = self.inpainter(
                            uv_inpaint_image_input, uv_inpaint_mask_input
                        )[0].permute(1, 2, 0)
                    clear()
                    if debug_mode:
                        make_image_grid(
                            [
                                tensor_to_image(uv_proj),
                                tensor_to_image(uv_valid_mask),
                                tensor_to_image(uv_inpaint_result),
                            ]
                        ).save(os.path.join(debug_dir, f"{mod_name}_uv_inpaint.jpg"))
                    uv_proj = uv_inpaint_result.contiguous()

                if mod_process_config.inpaint_mode == "view":
                    if view_inpaint_include_occlusion_boundary:
                        uv_max_depth_grad = uv_depth_grad.max(dim=0)[0]
                        uv_valid_mask = uv_valid_mask & (uv_max_depth_grad < 0.1)
                    uv_proj, uv_valid_mask = self.view_inpaint(
                        mod_name,
                        mesh,
                        uv_proj,
                        uv_valid_mask,
                        mod_process_config,
                        debug_dir=debug_dir if debug_mode else None,
                    )

                if poisson_reprojection:
                    # up and down
                    mesh.texture = uv_proj
                    uv_proj = self.cam_proj(
                        mod_tensor[4:5],
                        mesh,
                        cameras[4:5],
                        from_scratch=False,
                        poisson_blending=True,
                        pb_keep_original_border=True,
                        depth_grad_dilation=5,
                        uv_exp_blend_alpha=3,
                        uv_exp_blend_view_weight=torch.as_tensor([1, 1]),
                        aoi_cos_valid_threshold=0.2,
                        uv_size=uv_size,
                        uv_padding=True,
                        return_dict=False,
                    )
                    # front, sides and back
                    mesh.texture = uv_proj
                    uv_proj = self.cam_proj(
                        mod_tensor[0:4],
                        mesh,
                        cameras[0:4],
                        from_scratch=False,
                        poisson_blending=True,
                        pb_keep_original_border=True,
                        depth_grad_dilation=5,
                        uv_exp_blend_alpha=3,
                        uv_exp_blend_view_weight=torch.as_tensor([1, 1, 1, 1]),
                        aoi_cos_valid_threshold=0.2,
                        uv_size=uv_size,
                        uv_padding=True,
                        return_dict=False,
                    )

                if mod_name == "orm":
                    uv_proj[:, :, 0] = 1.0

            # Apply dilate to remove texture seams before final conversion
            if mod_name in ["rgb", "base_color"] and apply_dilate:  # Apply to color textures only
                uv_proj = torch_image_dilate(uv_proj)

            mod_uv_image[mod_name] = tensor_to_image(uv_proj)
            mod_uv_tensor[mod_name] = uv_proj
            clear()

        shaded_model_save_path = None
        if mod_uv_image["rgb"] is not None:
            shaded_model_save_path = os.path.join(save_dir, f"{save_name}_shaded.glb")
            replace_mesh_texture_and_save(
                mesh_path,
                shaded_model_save_path,
                texture=mod_uv_image["rgb"],
                backend="trimesh",
                task_id=save_name,
            )
        pbr_model_save_path = None
        if mod_uv_image["base_color"] is not None:
            pbr_model_save_path = os.path.join(save_dir, f"{save_name}_pbr.glb")
            replace_mesh_texture_and_save(
                mesh_path,
                pbr_model_save_path,
                texture=mod_uv_image["base_color"],
                metallic_roughness_texture=mod_uv_image["orm"],
                normal_texture=mod_uv_image["normal"],
                normal_strength=normal_strength,
                backend="trimesh",
                task_id=save_name,
            )

        clear()

        return TexturePipelineOutput(
            shaded_model_save_path=shaded_model_save_path,
            pbr_model_save_path=pbr_model_save_path,
        )