alice-sdf 1.11.0

//! USD (Universal Scene Description) Export
//!
//! Exports meshes to USDA (ASCII) format compatible with:
//! - Pixar USD Toolset, NVIDIA Omniverse, Apple Reality Composer,
//!   Houdini, Maya (via USD plugin), Blender (via USD addon)
//!
//! # Supported Features
//! - Mesh geometry (positions, normals, UVs, indices)
//! - PBR materials (UsdPreviewSurface)
//! - Transform hierarchy
//!
//! Author: Moroya Sakamoto

use crate::io::IoError;
use crate::material::MaterialLibrary;
use crate::mesh::Mesh;
use std::io::Write;
use std::path::Path;

/// USD up-axis convention
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UsdUpAxis {
    /// Y-up (default for USD, matches Maya/Blender)
    Y,
    /// Z-up (matches 3ds Max, CAD)
    Z,
}

/// Configuration for USDA export
#[derive(Debug, Clone)]
pub struct UsdConfig {
    /// Up axis convention
    pub up_axis: UsdUpAxis,
    /// Meters per unit (1.0 = meters, 0.01 = centimeters)
    pub meters_per_unit: f64,
    /// Export PBR materials
    pub export_materials: bool,
    /// Export vertex normals
    pub export_normals: bool,
}

impl Default for UsdConfig {
    fn default() -> Self {
        Self {
            up_axis: UsdUpAxis::Y,
            meters_per_unit: 1.0,
            export_materials: true,
            export_normals: true,
        }
    }
}

/// Export mesh to USDA (ASCII USD) format
pub fn export_usda(
    mesh: &Mesh,
    path: impl AsRef<Path>,
    config: &UsdConfig,
    materials: Option<&MaterialLibrary>,
) -> Result<(), IoError> {
    let file = std::fs::File::create(path)?;
    let mut w = std::io::BufWriter::new(file);

    let up_axis = match config.up_axis {
        UsdUpAxis::Y => "Y",
        UsdUpAxis::Z => "Z",
    };

    // Header
    writeln!(w, "#usda 1.0")?;
    writeln!(w, "(")?;
    writeln!(w, "    upAxis = \"{}\"", up_axis)?;
    writeln!(w, "    metersPerUnit = {}", config.meters_per_unit)?;
    writeln!(w, "    doc = \"Generated by ALICE-SDF\"")?;
    writeln!(w, ")\n")?;

    writeln!(w, "def Xform \"Root\"")?;
    writeln!(w, "{{")?;

    // Mesh
    writeln!(w, "    def Mesh \"Mesh\"")?;
    writeln!(w, "    {{")?;

    // Face vertex counts (all triangles = 3)
    let tri_count = mesh.indices.len() / 3;
    write!(w, "        int[] faceVertexCounts = [")?;
    for i in 0..tri_count {
        if i > 0 {
            write!(w, ", ")?;
        }
        write!(w, "3")?;
    }
    writeln!(w, "]")?;

    // Face vertex indices
    write!(w, "        int[] faceVertexIndices = [")?;
    for (i, idx) in mesh.indices.iter().enumerate() {
        if i > 0 {
            write!(w, ", ")?;
        }
        write!(w, "{}", idx)?;
    }
    writeln!(w, "]")?;

    // Points
    write!(w, "        point3f[] points = [")?;
    for (i, v) in mesh.vertices.iter().enumerate() {
        if i > 0 {
            write!(w, ", ")?;
        }
        write!(w, "({}, {}, {})", v.position.x, v.position.y, v.position.z)?;
    }
    writeln!(w, "]")?;

    // Normals
    if config.export_normals {
        write!(w, "        normal3f[] normals = [")?;
        for (i, v) in mesh.vertices.iter().enumerate() {
            if i > 0 {
                write!(w, ", ")?;
            }
            write!(w, "({}, {}, {})", v.normal.x, v.normal.y, v.normal.z)?;
        }
        writeln!(w, "]")?;
        writeln!(
            w,
            "        uniform token normals:interpolation = \"vertex\""
        )?;
    }

    // UVs
    write!(w, "        texCoord2f[] primvars:st = [")?;
    for (i, v) in mesh.vertices.iter().enumerate() {
        if i > 0 {
            write!(w, ", ")?;
        }
        write!(w, "({}, {})", v.uv.x, v.uv.y)?;
    }
    writeln!(w, "]")?;
    writeln!(
        w,
        "        uniform token primvars:st:interpolation = \"vertex\""
    )?;

    writeln!(w, "        uniform token subdivisionScheme = \"none\"")?;

    if config.export_materials {
        writeln!(
            w,
            "        rel material:binding = </Root/Materials/Material_0>"
        )?;
    }

    writeln!(w, "    }}\n")?;

    // Materials
    if config.export_materials {
        writeln!(w, "    def Scope \"Materials\"")?;
        writeln!(w, "    {{")?;

        let default_mat = (0.8, 0.8, 0.8, 0.0, 0.5, 1.0);
        let (diffuse_r, diffuse_g, diffuse_b, metallic, roughness, opacity) = materials
            .and_then(|lib| lib.materials.first())
            .map_or(default_mat, |mat| {
                (
                    mat.base_color[0],
                    mat.base_color[1],
                    mat.base_color[2],
                    mat.metallic,
                    mat.roughness,
                    mat.opacity,
                )
            });

        writeln!(w, "        def Material \"Material_0\"")?;
        writeln!(w, "        {{")?;
        writeln!(w, "            token outputs:surface.connect = </Root/Materials/Material_0/PBR.outputs:surface>")?;
        writeln!(w)?;
        writeln!(w, "            def Shader \"PBR\"")?;
        writeln!(w, "            {{")?;
        writeln!(
            w,
            "                uniform token info:id = \"UsdPreviewSurface\""
        )?;
        writeln!(
            w,
            "                color3f inputs:diffuseColor = ({}, {}, {})",
            diffuse_r, diffuse_g, diffuse_b
        )?;
        writeln!(w, "                float inputs:metallic = {}", metallic)?;
        writeln!(w, "                float inputs:roughness = {}", roughness)?;
        writeln!(w, "                float inputs:opacity = {}", opacity)?;
        writeln!(w, "                token outputs:surface")?;
        writeln!(w, "            }}")?;
        writeln!(w, "        }}")?;
        writeln!(w, "    }}")?;
    }

    writeln!(w, "}}")?;
    w.flush()?;
    Ok(())
}

/// Imported USD data with mesh and optional material
#[derive(Debug, Clone)]
pub struct ImportedUsda {
    /// The imported mesh
    pub mesh: Mesh,
    /// Optional material extracted from USDA
    pub material: Option<ImportedUsdMaterial>,
}

/// Material extracted from USDA
#[derive(Debug, Clone)]
pub struct ImportedUsdMaterial {
    /// Diffuse color (RGB)
    pub diffuse_color: [f32; 3],
    /// Metallic value (0.0 = dielectric, 1.0 = metal)
    pub metallic: f32,
    /// Roughness value (0.0 = smooth, 1.0 = rough)
    pub roughness: f32,
    /// Opacity (0.0 = transparent, 1.0 = opaque)
    pub opacity: f32,
}

/// Import mesh from USDA (ASCII USD) format
///
/// # Arguments
/// * `path` - Path to .usda file
///
/// # Returns
/// ImportedUsda containing mesh and optional material
pub fn import_usda(path: impl AsRef<Path>) -> Result<ImportedUsda, IoError> {
    use crate::mesh::Vertex;
    use glam::{Vec2, Vec3};

    let content = std::fs::read_to_string(path.as_ref())?;

    if !content.starts_with("#usda") {
        return Err(IoError::InvalidFormat(
            "Not a USDA file (missing #usda header)".into(),
        ));
    }

    // Parse helper: extract array content between [ ]
    fn extract_array(content: &str, key: &str) -> Option<String> {
        let search = format!("{} = [", key);
        if let Some(start) = content.find(&search) {
            let after_bracket = start + search.len();
            if let Some(end) = content[after_bracket..].find(']') {
                return Some(content[after_bracket..after_bracket + end].to_string());
            }
        }
        None
    }

    // Parse helper: extract single value
    fn extract_single_value(content: &str, key: &str) -> Option<String> {
        let search = format!("{} = ", key);
        if let Some(start) = content.find(&search) {
            let after_eq = start + search.len();
            let rest = &content[after_eq..];
            if let Some(end) = rest.find('\n') {
                return Some(rest[..end].trim().to_string());
            }
        }
        None
    }

    // Parse points: point3f[] points = [(x, y, z), ...]
    let points_str = extract_array(&content, "point3f[] points")
        .ok_or_else(|| IoError::InvalidFormat("Missing points array".into()))?;

    let mut positions = Vec::new();
    for tuple_str in points_str.split("),") {
        let cleaned = tuple_str
            .trim()
            .trim_start_matches('(')
            .trim_end_matches(')');
        if cleaned.is_empty() {
            continue;
        }
        let coords: Vec<f32> = cleaned
            .split(',')
            .filter_map(|s| s.trim().parse().ok())
            .collect();
        if coords.len() == 3 {
            positions.push(Vec3::new(coords[0], coords[1], coords[2]));
        }
    }

    if positions.is_empty() {
        return Err(IoError::InvalidFormat("No valid points found".into()));
    }

    // Parse indices: int[] faceVertexIndices = [0, 1, 2, ...]
    let indices_str = extract_array(&content, "int[] faceVertexIndices")
        .ok_or_else(|| IoError::InvalidFormat("Missing faceVertexIndices".into()))?;

    let raw_indices: Vec<u32> = indices_str
        .split(',')
        .filter_map(|s| s.trim().parse().ok())
        .collect();

    // Parse face counts: int[] faceVertexCounts = [3, 3, 3, ...]
    let face_counts_str = extract_array(&content, "int[] faceVertexCounts")
        .ok_or_else(|| IoError::InvalidFormat("Missing faceVertexCounts".into()))?;

    let face_counts: Vec<usize> = face_counts_str
        .split(',')
        .filter_map(|s| s.trim().parse().ok())
        .collect();

    // Triangulate faces (fan triangulation for quads/n-gons)
    let mut triangulated_indices = Vec::new();
    let mut idx_offset = 0;
    for &count in &face_counts {
        if count == 3 {
            // Already a triangle
            triangulated_indices.push(raw_indices[idx_offset]);
            triangulated_indices.push(raw_indices[idx_offset + 1]);
            triangulated_indices.push(raw_indices[idx_offset + 2]);
        } else if count > 3 {
            // Fan triangulation: (0, 1, 2), (0, 2, 3), (0, 3, 4), ...
            for i in 1..(count - 1) {
                triangulated_indices.push(raw_indices[idx_offset]);
                triangulated_indices.push(raw_indices[idx_offset + i]);
                triangulated_indices.push(raw_indices[idx_offset + i + 1]);
            }
        }
        idx_offset += count;
    }

    // Parse normals: normal3f[] normals = [(x, y, z), ...]
    let mut normals = vec![Vec3::Y; positions.len()]; // Default to Y-up
    if let Some(normals_str) = extract_array(&content, "normal3f[] normals") {
        let mut parsed_normals = Vec::new();
        for tuple_str in normals_str.split("),") {
            let cleaned = tuple_str
                .trim()
                .trim_start_matches('(')
                .trim_end_matches(')');
            if cleaned.is_empty() {
                continue;
            }
            let coords: Vec<f32> = cleaned
                .split(',')
                .filter_map(|s| s.trim().parse().ok())
                .collect();
            if coords.len() == 3 {
                let n = Vec3::new(coords[0], coords[1], coords[2]);
                parsed_normals.push(if n.length_squared() > 0.0001 {
                    n.normalize()
                } else {
                    Vec3::Y
                });
            }
        }
        if parsed_normals.len() == positions.len() {
            normals = parsed_normals;
        }
    }

    // Parse UVs: texCoord2f[] primvars:st = [(u, v), ...]
    let mut uvs = vec![Vec2::ZERO; positions.len()];
    if let Some(uvs_str) = extract_array(&content, "texCoord2f[] primvars:st") {
        let mut parsed_uvs = Vec::new();
        for tuple_str in uvs_str.split("),") {
            let cleaned = tuple_str
                .trim()
                .trim_start_matches('(')
                .trim_end_matches(')');
            if cleaned.is_empty() {
                continue;
            }
            let coords: Vec<f32> = cleaned
                .split(',')
                .filter_map(|s| s.trim().parse().ok())
                .collect();
            if coords.len() == 2 {
                parsed_uvs.push(Vec2::new(coords[0], coords[1]));
            }
        }
        if parsed_uvs.len() == positions.len() {
            uvs = parsed_uvs;
        }
    }

    // Build vertices
    let vertices: Vec<Vertex> = positions
        .iter()
        .zip(normals.iter())
        .zip(uvs.iter())
        .map(|((&pos, &norm), &uv)| {
            let mut v = Vertex::new(pos, norm);
            v.uv = uv;
            v
        })
        .collect();

    // Parse material (optional)
    let mut material = None;
    if let Some(diffuse_str) = extract_single_value(&content, "color3f inputs:diffuseColor") {
        let cleaned = diffuse_str
            .trim()
            .trim_start_matches('(')
            .trim_end_matches(')');
        let rgb: Vec<f32> = cleaned
            .split(',')
            .filter_map(|s| s.trim().parse().ok())
            .collect();

        if rgb.len() == 3 {
            let metallic = extract_single_value(&content, "float inputs:metallic")
                .and_then(|s| s.parse().ok())
                .unwrap_or(0.0);
            let roughness = extract_single_value(&content, "float inputs:roughness")
                .and_then(|s| s.parse().ok())
                .unwrap_or(0.5);
            let opacity = extract_single_value(&content, "float inputs:opacity")
                .and_then(|s| s.parse().ok())
                .unwrap_or(1.0);

            material = Some(ImportedUsdMaterial {
                diffuse_color: [rgb[0], rgb[1], rgb[2]],
                metallic,
                roughness,
                opacity,
            });
        }
    }

    let mesh = Mesh {
        vertices,
        indices: triangulated_indices,
    };

    Ok(ImportedUsda { mesh, material })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mesh::{sdf_to_mesh, MarchingCubesConfig};
    use crate::types::SdfNode;
    use glam::Vec3;

    #[test]
    fn test_usda_export() {
        let sphere = SdfNode::sphere(1.0);
        let mesh = sdf_to_mesh(
            &sphere,
            Vec3::splat(-2.0),
            Vec3::splat(2.0),
            &MarchingCubesConfig {
                resolution: 16,
                ..Default::default()
            },
        );

        let path = std::env::temp_dir().join("alice_test.usda");
        export_usda(&mesh, &path, &UsdConfig::default(), None).unwrap();

        let content = std::fs::read_to_string(&path).unwrap();
        assert!(content.contains("#usda 1.0"));
        assert!(content.contains("upAxis = \"Y\""));
        assert!(content.contains("def Mesh \"Mesh\""));
        assert!(content.contains("point3f[] points"));
        assert!(content.contains("UsdPreviewSurface"));

        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn test_usda_z_up() {
        let sphere = SdfNode::sphere(1.0);
        let mesh = sdf_to_mesh(
            &sphere,
            Vec3::splat(-2.0),
            Vec3::splat(2.0),
            &MarchingCubesConfig {
                resolution: 8,
                ..Default::default()
            },
        );

        let path = std::env::temp_dir().join("alice_test_zup.usda");
        let config = UsdConfig {
            up_axis: UsdUpAxis::Z,
            ..Default::default()
        };
        export_usda(&mesh, &path, &config, None).unwrap();

        let content = std::fs::read_to_string(&path).unwrap();
        assert!(content.contains("upAxis = \"Z\""));

        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn test_usda_import_roundtrip() {
        // Create a sphere mesh
        let sphere = SdfNode::sphere(1.0);
        let original_mesh = sdf_to_mesh(
            &sphere,
            Vec3::splat(-2.0),
            Vec3::splat(2.0),
            &MarchingCubesConfig {
                resolution: 16,
                ..Default::default()
            },
        );

        // Export to USDA
        let path = std::env::temp_dir().join("alice_test_import.usda");
        export_usda(&original_mesh, &path, &UsdConfig::default(), None).unwrap();

        // Import back
        let imported = import_usda(&path).unwrap();

        // Verify counts match
        assert_eq!(
            imported.mesh.vertices.len(),
            original_mesh.vertices.len(),
            "Vertex count mismatch"
        );
        assert_eq!(
            imported.mesh.indices.len(),
            original_mesh.indices.len(),
            "Index count mismatch"
        );

        // Verify positions are close (within epsilon)
        for (orig, imp) in original_mesh
            .vertices
            .iter()
            .zip(imported.mesh.vertices.iter())
        {
            let diff = (orig.position - imp.position).length();
            assert!(diff < 0.001, "Position difference too large: {}", diff);
        }

        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn test_usda_import_with_material() {
        use crate::material::{Material, MaterialLibrary};

        let sphere = SdfNode::sphere(1.0);
        let mesh = sdf_to_mesh(
            &sphere,
            Vec3::splat(-2.0),
            Vec3::splat(2.0),
            &MarchingCubesConfig {
                resolution: 8,
                ..Default::default()
            },
        );

        // Create material with specific values
        let mut mat_lib = MaterialLibrary::default();
        mat_lib.materials.push(Material {
            base_color: [0.9, 0.2, 0.1, 1.0],
            metallic: 0.7,
            roughness: 0.3,
            opacity: 0.8,
            ..Default::default()
        });

        // Export with material
        let path = std::env::temp_dir().join("alice_test_material.usda");
        export_usda(&mesh, &path, &UsdConfig::default(), Some(&mat_lib)).unwrap();

        // Import and verify material
        let imported = import_usda(&path).unwrap();
        assert!(imported.material.is_some(), "Material should be imported");

        let mat = imported.material.unwrap();
        assert!((mat.diffuse_color[0] - 0.9).abs() < 0.01, "Red mismatch");
        assert!((mat.diffuse_color[1] - 0.2).abs() < 0.01, "Green mismatch");
        assert!((mat.diffuse_color[2] - 0.1).abs() < 0.01, "Blue mismatch");
        assert!((mat.metallic - 0.7).abs() < 0.01, "Metallic mismatch");
        assert!((mat.roughness - 0.3).abs() < 0.01, "Roughness mismatch");
        assert!((mat.opacity - 0.8).abs() < 0.01, "Opacity mismatch");

        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn test_usda_import_invalid() {
        // Try to import a non-USDA file
        let path = std::env::temp_dir().join("alice_test_invalid.txt");
        std::fs::write(&path, "This is not a USDA file").unwrap();

        let result = import_usda(&path);
        assert!(result.is_err(), "Should fail on invalid file");

        if let Err(IoError::InvalidFormat(msg)) = result {
            assert!(
                msg.contains("missing #usda header"),
                "Should mention missing header"
            );
        } else {
            panic!("Expected InvalidFormat error");
        }

        std::fs::remove_file(&path).ok();
    }
}