# -*- coding: utf-8 -*-
"""
USD Format Conversion Module
Author  : CraftPilot

Provides conversion functionality between USD files (.usd/.usda/.usdc) and PSON JSON format.
Supports complete USD scene parsing, including geometry, materials, textures, transforms, hierarchy, cameras, lights, etc.
"""

import os
import json
import logging
from typing import Dict, List, Optional, Tuple, Any, Union

logger = logging.getLogger(__name__)

# USD library import - lazy import for handling in test environment
try:
    from pxr import Usd, UsdGeom, UsdShade, Gf, Sdf, Tf
    USD_AVAILABLE = True
except ImportError:
    USD_AVAILABLE = False
    logger.warning("USD library not installed, USD conversion functionality will be unavailable. Please run: pip install usd-core")


# ============================================================================
# USD Support Check
# ============================================================================

def check_usd_support() -> bool:
    """
    Check if USD library is available
    
    Returns:
        bool: Whether USD library is available
    """
    return USD_AVAILABLE


def require_usd():
    """
    Check if USD library is available, raise exception if not
    
    Raises:
        ImportError: USD library not installed
    """
    if not USD_AVAILABLE:
        raise ImportError(
            "USD library not installed. Please run the following command to install: pip install usd-core"
        )


# ============================================================================
# USD File Format Detection
# ============================================================================

def is_usd_file(file_path: str) -> bool:
    """
    Detect if file is in USD format
    
    Args:
        file_path: File path
        
    Returns:
        bool: Whether it's a USD file (.usd, .usda, .usdc)
    """
    ext = os.path.splitext(file_path)[1].lower()
    return ext in {'.usd', '.usda', '.usdc'}


# ============================================================================
# USD Geometry Data Extraction
# ============================================================================

def extract_mesh_data(mesh_prim: 'UsdGeom.Mesh') -> Optional[Dict[str, Any]]:
    """
    Extract geometry data from USD Mesh Prim
    
    Args:
        mesh_prim: USD Mesh Prim object
        
    Returns:
        Dict: Dictionary containing vertices, faces, attributes, etc., in PSON format
        {
            "attribs": {
                "point": {"P": {"type": "float", "data": [[x,y,z], ...]}},
                "vertex": {},
                "prim": {},
                "detail": {}
            },
            "faces": [[v1,v2,v3], ...]
        }
    """
    require_usd()
    
    if not mesh_prim.IsA(UsdGeom.Mesh):
        logger.warning(f"Prim {mesh_prim.GetPath()} is not a Mesh type")
        return None
    
    mesh = UsdGeom.Mesh(mesh_prim)
    
    # Extract vertex positions
    points_attr = mesh.GetPointsAttr()
    if not points_attr:
        logger.warning(f"Mesh {mesh_prim.GetPath()} has no vertex data")
        return None
    
    points = points_attr.Get()
    if not points:
        return None
    
    # Convert vertex data to list format
    positions = [[float(p[0]), float(p[1]), float(p[2])] for p in points]
    
    # Extract face data
    face_vertex_counts = mesh.GetFaceVertexCountsAttr().Get()
    face_vertex_indices = mesh.GetFaceVertexIndicesAttr().Get()
    
    if not face_vertex_counts or not face_vertex_indices:
        logger.warning(f"Mesh {mesh_prim.GetPath()} has no face data")
        return None
    
    # Convert face data (supports triangles and polygons, polygons need triangulation)
    faces = []
    idx = 0
    for count in face_vertex_counts:
        face_indices = [int(face_vertex_indices[idx + i]) for i in range(count)]
        
        # Triangulate polygons (Fan triangulation)
        if count == 3:
            faces.append(face_indices)
        elif count > 3:
            # Decompose polygon into multiple triangles
            for i in range(1, count - 1):
                faces.append([face_indices[0], face_indices[i], face_indices[i + 1]])
        
        idx += count
    
    # Initialize result dictionary
    result = {
        "attribs": {
            "point": {
                "P": {
                    "type": "float",
                    "data": positions
                }
            },
            "vertex": {},
            "prim": {},
            "detail": {}
        },
        "faces": faces
    }
    
    # Extract normals (if any)
    normals_attr = mesh.GetNormalsAttr()
    if normals_attr:
        normals = normals_attr.Get()
        if normals:
            normals_list = [[float(n[0]), float(n[1]), float(n[2])] for n in normals]
            
            # Check normal interpolation type
            normals_interpolation = mesh.GetNormalsInterpolation()
            if normals_interpolation == 'vertex':
                result["attribs"]["vertex"]["N"] = {
                    "type": "float",
                    "data": normals_list
                }
            elif normals_interpolation == 'faceVarying':
                result["attribs"]["vertex"]["N"] = {
                    "type": "float",
                    "data": normals_list
                }
            else:  # uniform, constant, etc.
                result["attribs"]["point"]["N"] = {
                    "type": "float",
                    "data": normals_list
                }
    
    # Extract UV coordinates (if any)
    primvars = UsdGeom.PrimvarsAPI(mesh_prim)
    st_primvar = primvars.GetPrimvar("st")
    if not st_primvar:
        st_primvar = primvars.GetPrimvar("uv")
    
    if st_primvar:
        uvs = st_primvar.Get()
        if uvs:
            uvs_list = [[float(uv[0]), float(uv[1])] for uv in uvs]
            
            # Check UV interpolation type
            uv_interpolation = st_primvar.GetInterpolation()
            if uv_interpolation == 'vertex' or uv_interpolation == 'faceVarying':
                result["attribs"]["vertex"]["uv"] = {
                    "type": "float",
                    "data": uvs_list
                }
            else:
                result["attribs"]["point"]["uv"] = {
                    "type": "float",
                    "data": uvs_list
                }
    
    # Extract vertex colors (if any)
    color_primvar = primvars.GetPrimvar("displayColor")
    if color_primvar:
        colors = color_primvar.Get()
        if colors:
            colors_list = [[float(c[0]), float(c[1]), float(c[2])] for c in colors]
            result["attribs"]["point"]["Cd"] = {
                "type": "float",
                "data": colors_list
            }
    
    return result


def extract_transform(prim: 'Usd.Prim') -> Optional[Dict[str, Any]]:
    """
    Extract transform information from Prim
    
    Args:
        prim: USD Prim object
        
    Returns:
        Dict: Transform information dictionary
        {
            "translate": [x, y, z],
            "rotate": [rx, ry, rz],  # Euler angles (degrees)
            "scale": [sx, sy, sz],
            "matrix": [[...], [...], [...], [...]]  # 4x4 transform matrix
        }
    """
    require_usd()
    
    xformable = UsdGeom.Xformable(prim)
    if not xformable:
        return None
    
    # Get local transform matrix
    local_transform = xformable.GetLocalTransformation()
    
    # Decompose transform matrix
    translation = local_transform.ExtractTranslation()
    rotation = local_transform.ExtractRotation()
    scale_shear = local_transform.ExtractRotationMatrix()
    
    # Extract scale (simplified, assuming no shear)
    scale_x = Gf.Vec3d(scale_shear[0][0], scale_shear[1][0], scale_shear[2][0]).GetLength()
    scale_y = Gf.Vec3d(scale_shear[0][1], scale_shear[1][1], scale_shear[2][1]).GetLength()
    scale_z = Gf.Vec3d(scale_shear[0][2], scale_shear[1][2], scale_shear[2][2]).GetLength()
    
    # Convert rotation to Euler angles (XYZ order)
    rotation_xyz = rotation.GetAxis() * rotation.GetAngle()
    
    result = {
        "translate": [float(translation[0]), float(translation[1]), float(translation[2])],
        "rotate": [float(rotation_xyz[0]), float(rotation_xyz[1]), float(rotation_xyz[2])],
        "scale": [float(scale_x), float(scale_y), float(scale_z)],
        "matrix": [
            [float(local_transform[i][j]) for j in range(4)] for i in range(4)
        ]
    }
    
    return result


def extract_camera_data(camera_prim: 'UsdGeom.Camera') -> Optional[Dict[str, Any]]:
    """
    Extract USD camera data
    
    Args:
        camera_prim: USD Camera Prim object
        
    Returns:
        Dict: Camera parameters dictionary
    """
    require_usd()
    
    if not camera_prim.IsA(UsdGeom.Camera):
        return None
    
    camera = UsdGeom.Camera(camera_prim)
    
    result = {
        "type": "camera",
        "focal_length": float(camera.GetFocalLengthAttr().Get()),
        "horizontal_aperture": float(camera.GetHorizontalApertureAttr().Get()),
        "vertical_aperture": float(camera.GetVerticalApertureAttr().Get()),
        "projection": str(camera.GetProjectionAttr().Get())
    }
    
    # Extract transform information
    transform = extract_transform(camera_prim)
    if transform:
        result["transform"] = transform
    
    return result


def extract_light_data(light_prim: 'Usd.Prim') -> Optional[Dict[str, Any]]:
    """
    Extract USD light data
    
    Args:
        light_prim: USD Light Prim object
        
    Returns:
        Dict: Light parameters dictionary
    """
    require_usd()
    
    # USD lights use UsdLux schema, but for simplicity, only extract basic info here
    result = {
        "type": "light",
        "light_type": light_prim.GetTypeName()
    }
    
    # Extract transform information
    transform = extract_transform(light_prim)
    if transform:
        result["transform"] = transform
    
    return result


# ============================================================================
# USD Scene Traversal
# ============================================================================

def traverse_stage(stage: 'Usd.Stage', 
                   include_meshes: bool = True,
                   include_cameras: bool = True,
                   include_lights: bool = True) -> Dict[str, Any]:
    """
    Traverse USD Stage, extract all geometry, cameras, lights, etc.
    
    Args:
        stage: USD Stage object
        include_meshes: Whether to include meshes
        include_cameras: Whether to include cameras
        include_lights: Whether to include lights
        
    Returns:
        Dict: Scene data dictionary
        {
            "meshes": {
                "mesh_name": {mesh_data},
                ...
            },
            "cameras": {
                "camera_name": {camera_data},
                ...
            },
            "lights": {
                "light_name": {light_data},
                ...
            },
            "hierarchy": {...}  # Hierarchy information
        }
    """
    require_usd()
    
    result = {
        "meshes": {},
        "cameras": {},
        "lights": {},
        "hierarchy": {}
    }
    
    # Traverse all Prims
    for prim in stage.Traverse():
        prim_path = str(prim.GetPath())
        prim_name = prim.GetName()
        
        # Extract meshes
        if include_meshes and prim.IsA(UsdGeom.Mesh):
            mesh_data = extract_mesh_data(prim)
            if mesh_data:
                # Use full path as unique identifier
                mesh_key = prim_path.replace('/', '_').strip('_')
                if not mesh_key:
                    mesh_key = prim_name
                
                result["meshes"][mesh_key] = mesh_data
                
                # Add transform information
                transform = extract_transform(prim)
                if transform:
                    result["meshes"][mesh_key]["transform"] = transform
        
        # Extract cameras
        elif include_cameras and prim.IsA(UsdGeom.Camera):
            camera_data = extract_camera_data(prim)
            if camera_data:
                camera_key = prim_path.replace('/', '_').strip('_')
                if not camera_key:
                    camera_key = prim_name
                result["cameras"][camera_key] = camera_data
        
        # Extract lights
        elif include_lights and prim.GetTypeName().startswith('Light'):
            light_data = extract_light_data(prim)
            if light_data:
                light_key = prim_path.replace('/', '_').strip('_')
                if not light_key:
                    light_key = prim_name
                result["lights"][light_key] = light_data
    
    return result


# ============================================================================
# USD to PSON JSON Conversion
# ============================================================================

def usd_to_pson_json(usd_file_path: str, 
                     output_json_path: Optional[str] = None,
                     include_cameras: bool = False,
                     include_lights: bool = False) -> Dict[str, Any]:
    """
    Convert USD file to PSON JSON format
    
    Args:
        usd_file_path: USD file path (.usd/.usda/.usdc)
        output_json_path: Output JSON file path (optional, if None, file is not saved)
        include_cameras: Whether to include camera data
        include_lights: Whether to include light data
        
    Returns:
        Dict: JSON data in PSON format
        
    Raises:
        FileNotFoundError: USD file does not exist
        ValueError: USD file format error or cannot be parsed
    """
    require_usd()
    
    if not os.path.exists(usd_file_path):
        raise FileNotFoundError(f"USD file does not exist: {usd_file_path}")
    
    if not is_usd_file(usd_file_path):
        raise ValueError(f"Not a valid USD file format: {usd_file_path}")
    
    try:
        # Open USD Stage
        stage = Usd.Stage.Open(usd_file_path)
        if not stage:
            raise ValueError(f"Cannot open USD file: {usd_file_path}")
        
        # Traverse Stage to extract data
        scene_data = traverse_stage(
            stage, 
            include_meshes=True,
            include_cameras=include_cameras,
            include_lights=include_lights
        )
        
        # Convert to PSON format
        pson_data = {
            "Output_Geo": {}
        }
        
        # If only one mesh, use Output_Geo directly as root
        if len(scene_data["meshes"]) == 1:
            mesh_name = list(scene_data["meshes"].keys())[0]
            mesh_data = scene_data["meshes"][mesh_name]
            pson_data["Output_Geo"] = {
                "attribs": mesh_data["attribs"],
                "faces": mesh_data["faces"]
            }
            if "transform" in mesh_data:
                pson_data["Output_Geo"]["transform"] = mesh_data["transform"]
        
        # If multiple meshes, store as array
        elif len(scene_data["meshes"]) > 1:
            pson_data["Output_Geo"] = {
                "meshes": scene_data["meshes"]
            }
        
        # Add cameras and lights (if any)
        if scene_data["cameras"]:
            pson_data["cameras"] = scene_data["cameras"]
        
        if scene_data["lights"]:
            pson_data["lights"] = scene_data["lights"]
        
        # Add metadata
        pson_data["metadata"] = {
            "source": os.path.basename(usd_file_path),
            "source_type": "usd",
            "mesh_count": len(scene_data["meshes"]),
            "camera_count": len(scene_data["cameras"]),
            "light_count": len(scene_data["lights"])
        }
        
        # Save JSON file (if output path is specified)
        if output_json_path:
            output_dir = os.path.dirname(output_json_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            with open(output_json_path, 'w', encoding='utf-8') as f:
                json.dump(pson_data, f, indent=2, ensure_ascii=False)
            
            logger.info(f"USD to JSON conversion complete: {usd_file_path} -> {output_json_path}")
        
        return pson_data
    
    except Exception as e:
        logger.error(f"USD conversion failed: {usd_file_path}, error: {str(e)}")
        raise ValueError(f"USD conversion failed: {str(e)}")


def usd_file_to_json_file(usd_file_path: str, 
                          output_json_path: Optional[str] = None) -> str:
    """
    Convert USD file to JSON file (simplified interface)
    
    Args:
        usd_file_path: USD file path
        output_json_path: Output JSON path (if None, auto-generated)
        
    Returns:
        str: Output JSON file path
    """
    if output_json_path is None:
        # Auto-generate output path (same directory as USD file, change extension to .json)
        base_name = os.path.splitext(usd_file_path)[0]
        output_json_path = f"{base_name}.json"
    
    usd_to_pson_json(usd_file_path, output_json_path)
    
    return output_json_path


# ============================================================================
# Batch Conversion
# ============================================================================

def batch_convert_usd_to_json(usd_dir: str, 
                               output_dir: Optional[str] = None,
                               recursive: bool = False) -> List[Tuple[str, str]]:
    """
    Batch convert USD files in directory to JSON
    
    Args:
        usd_dir: Directory containing USD files
        output_dir: Output directory (if None, output to original directory)
        recursive: Whether to recursively search subdirectories
        
    Returns:
        List[Tuple[str, str]]: List of successfully converted files [(usd_path, json_path), ...]
    """
    require_usd()
    
    if not os.path.exists(usd_dir):
        raise FileNotFoundError(f"Directory does not exist: {usd_dir}")
    
    converted_files = []
    
    # Find all USD files
    if recursive:
        usd_files = []
        for root, dirs, files in os.walk(usd_dir):
            for file in files:
                if is_usd_file(file):
                    usd_files.append(os.path.join(root, file))
    else:
        usd_files = [
            os.path.join(usd_dir, f) for f in os.listdir(usd_dir) 
            if os.path.isfile(os.path.join(usd_dir, f)) and is_usd_file(f)
        ]
    
    # Convert each file
    for usd_file in usd_files:
        try:
            if output_dir:
                # Maintain relative directory structure
                rel_path = os.path.relpath(usd_file, usd_dir)
                json_name = os.path.splitext(rel_path)[0] + '.json'
                output_json_path = os.path.join(output_dir, json_name)
            else:
                output_json_path = os.path.splitext(usd_file)[0] + '.json'
            
            usd_to_pson_json(usd_file, output_json_path)
            converted_files.append((usd_file, output_json_path))
            
        except Exception as e:
            logger.error(f"Conversion failed {usd_file}: {str(e)}")
            continue
    
    logger.info(f"Batch conversion complete: {len(converted_files)}/{len(usd_files)} files")
    
    return converted_files
