"""Mesh I/O utilities for TRELLIS2 meshes."""

import logging
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union

import numpy as np
import torch

logger = logging.getLogger(__name__)

# Try to import trimesh for mesh I/O
try:
    import trimesh
    TRIMESH_AVAILABLE = True
except ImportError:
    TRIMESH_AVAILABLE = False
    logger.warning("trimesh not available. Install with: pip install trimesh")


def load_mesh(
    path: Union[str, Path],
    device: str = "cuda"
) -> Dict[str, Any]:
    """
    Load a mesh from file and convert to tensor format.

    Args:
        path: Path to mesh file (supports .obj, .ply, .glb, .gltf)
        device: Device for output tensors

    Returns:
        Dictionary containing mesh data:
            - vertices: [N, 3] tensor of vertex positions
            - faces: [F, 3] tensor of face indices
            - normals: [N, 3] tensor of vertex normals (if available)
            - uvs: [N, 2] tensor of texture coordinates (if available)
            - colors: [N, 3/4] tensor of vertex colors (if available)
            - textures: dict of texture images (if available)
    """
    if not TRIMESH_AVAILABLE:
        raise ImportError("trimesh is required for mesh loading. Install with: pip install trimesh")

    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Mesh file not found: {path}")

    logger.info(f"Loading mesh from {path}")

    # Load mesh
    mesh = trimesh.load(str(path), force='mesh')

    result = {}

    # Vertices
    vertices = torch.tensor(mesh.vertices, dtype=torch.float32, device=device)
    result["vertices"] = vertices

    # Faces
    faces = torch.tensor(mesh.faces, dtype=torch.int64, device=device)
    result["faces"] = faces

    # Normals
    if hasattr(mesh, 'vertex_normals') and mesh.vertex_normals is not None:
        normals = torch.tensor(mesh.vertex_normals, dtype=torch.float32, device=device)
        result["normals"] = normals

    # UVs
    if hasattr(mesh.visual, 'uv') and mesh.visual.uv is not None:
        uvs = torch.tensor(mesh.visual.uv, dtype=torch.float32, device=device)
        result["uvs"] = uvs

    # Vertex colors
    if hasattr(mesh.visual, 'vertex_colors') and mesh.visual.vertex_colors is not None:
        colors = torch.tensor(
            mesh.visual.vertex_colors[:, :3] / 255.0,
            dtype=torch.float32,
            device=device
        )
        result["colors"] = colors

    # Textures (if PBR material)
    if hasattr(mesh.visual, 'material'):
        material = mesh.visual.material
        textures = {}

        if hasattr(material, 'baseColorTexture') and material.baseColorTexture is not None:
            textures["albedo"] = np.array(material.baseColorTexture)

        if hasattr(material, 'metallicRoughnessTexture') and material.metallicRoughnessTexture is not None:
            textures["metallic_roughness"] = np.array(material.metallicRoughnessTexture)

        if hasattr(material, 'normalTexture') and material.normalTexture is not None:
            textures["normal"] = np.array(material.normalTexture)

        if textures:
            result["textures"] = textures

    # Store original mesh reference for advanced operations
    result["_trimesh"] = mesh

    logger.info(
        f"Loaded mesh: {len(vertices)} vertices, {len(faces)} faces"
    )

    return result


def save_mesh(
    mesh_data: Dict[str, Any],
    path: Union[str, Path],
    file_format: Optional[str] = None
) -> None:
    """
    Save mesh data to file.

    Args:
        mesh_data: Dictionary containing mesh tensors
        path: Output file path
        file_format: Optional format override (obj, ply, glb, gltf)
    """
    if not TRIMESH_AVAILABLE:
        raise ImportError("trimesh is required for mesh saving. Install with: pip install trimesh")

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    # If we have the original trimesh, use it
    if "_trimesh" in mesh_data:
        mesh = mesh_data["_trimesh"]
    else:
        # Reconstruct from tensors
        vertices = mesh_data["vertices"].cpu().numpy()
        faces = mesh_data["faces"].cpu().numpy()

        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)

        if "normals" in mesh_data:
            mesh.vertex_normals = mesh_data["normals"].cpu().numpy()

        if "colors" in mesh_data:
            colors = mesh_data["colors"].cpu().numpy()
            # Convert to 0-255 range and add alpha
            colors_255 = (colors * 255).astype(np.uint8)
            if colors_255.shape[1] == 3:
                alpha = np.full((len(colors_255), 1), 255, dtype=np.uint8)
                colors_255 = np.concatenate([colors_255, alpha], axis=1)
            mesh.visual.vertex_colors = colors_255

    # Save
    mesh.export(str(path), file_type=file_format)
    logger.info(f"Saved mesh to {path}")


def mesh_to_tensors(
    mesh: Any,
    device: str = "cuda"
) -> Dict[str, torch.Tensor]:
    """
    Convert a trimesh/TRELLIS mesh object to tensor dictionary.

    Args:
        mesh: Mesh object (trimesh.Trimesh or TRELLIS mesh)
        device: Target device for tensors

    Returns:
        Dictionary of mesh tensors
    """
    result = {}

    # Handle different mesh types
    if hasattr(mesh, 'vertices'):
        # trimesh or similar
        if isinstance(mesh.vertices, torch.Tensor):
            result["vertices"] = mesh.vertices.to(device)
        else:
            result["vertices"] = torch.tensor(mesh.vertices, dtype=torch.float32, device=device)

        if hasattr(mesh, 'faces'):
            if isinstance(mesh.faces, torch.Tensor):
                result["faces"] = mesh.faces.to(device)
            else:
                result["faces"] = torch.tensor(mesh.faces, dtype=torch.int64, device=device)

    # Handle TRELLIS-specific mesh format
    elif hasattr(mesh, 'v_pos'):
        result["vertices"] = mesh.v_pos.to(device)
        if hasattr(mesh, 't_pos_idx'):
            result["faces"] = mesh.t_pos_idx.to(device)

    # Handle dictionary input
    elif isinstance(mesh, dict):
        for key, value in mesh.items():
            if isinstance(value, torch.Tensor):
                result[key] = value.to(device)
            elif isinstance(value, np.ndarray):
                result[key] = torch.tensor(value, device=device)

    return result


def get_mesh_info(mesh_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    Get information about a mesh.

    Args:
        mesh_data: Dictionary containing mesh tensors

    Returns:
        Dictionary with mesh properties
    """
    info = {}

    if "vertices" in mesh_data:
        vertices = mesh_data["vertices"]
        info["num_vertices"] = len(vertices)
        info["vertex_min"] = vertices.min(dim=0).values.cpu().tolist()
        info["vertex_max"] = vertices.max(dim=0).values.cpu().tolist()
        info["vertex_center"] = vertices.mean(dim=0).cpu().tolist()

        # Compute bounding box size
        bbox_min = vertices.min(dim=0).values
        bbox_max = vertices.max(dim=0).values
        bbox_size = bbox_max - bbox_min
        info["bbox_size"] = bbox_size.cpu().tolist()

    if "faces" in mesh_data:
        info["num_faces"] = len(mesh_data["faces"])

    if "normals" in mesh_data:
        info["has_normals"] = True

    if "uvs" in mesh_data:
        info["has_uvs"] = True

    if "colors" in mesh_data:
        info["has_colors"] = True

    if "textures" in mesh_data:
        info["texture_types"] = list(mesh_data["textures"].keys())

    return info


def normalize_mesh(
    mesh_data: Dict[str, torch.Tensor],
    target_scale: float = 1.0
) -> Dict[str, torch.Tensor]:
    """
    Normalize mesh to fit within a unit cube centered at origin.

    Args:
        mesh_data: Dictionary containing mesh tensors
        target_scale: Scale factor for the normalized mesh

    Returns:
        Normalized mesh data
    """
    vertices = mesh_data["vertices"]

    # Compute center and scale
    center = vertices.mean(dim=0, keepdim=True)
    vertices_centered = vertices - center

    max_extent = vertices_centered.abs().max()
    scale = target_scale / max_extent

    # Apply normalization
    result = mesh_data.copy()
    result["vertices"] = vertices_centered * scale

    return result


def combine_meshes(
    meshes: list,
    device: str = "cuda"
) -> Dict[str, torch.Tensor]:
    """
    Combine multiple meshes into a single mesh.

    Args:
        meshes: List of mesh dictionaries
        device: Target device

    Returns:
        Combined mesh dictionary
    """
    all_vertices = []
    all_faces = []
    vertex_offset = 0

    for mesh in meshes:
        vertices = mesh["vertices"]
        faces = mesh["faces"]

        all_vertices.append(vertices)
        all_faces.append(faces + vertex_offset)

        vertex_offset += len(vertices)

    result = {
        "vertices": torch.cat(all_vertices, dim=0).to(device),
        "faces": torch.cat(all_faces, dim=0).to(device),
    }

    # Combine other attributes if present
    if all(["normals" in m for m in meshes]):
        result["normals"] = torch.cat([m["normals"] for m in meshes], dim=0).to(device)

    if all(["colors" in m for m in meshes]):
        result["colors"] = torch.cat([m["colors"] for m in meshes], dim=0).to(device)

    return result