rust-usd 0.0.1

//! Rust bindings to OpenUSD (pxr).
//!
//! v0.0.5 surface:
//! - Stage: open with explicit `InitialLoadSet`, post-open `load`/`unload`,
//!   `prim_at_path`, `open_for_painting` (sublayer authoring), `save_edit_layer`.
//! - Prim tree walk (path, type_name, children, is_mesh, as_mesh).
//! - UsdGeomMesh attrs: points, face vertex counts/indices, normals,
//!   primvars:st (+indices), subdivision scheme, world-space transform.
//! - UsdShade: bound texture paths from a mesh's material.
//! - Variant sets: enumerate, read selection, set/clear selection.
//! - `forge://` URI resolution via [`AssetResolver`] + [`install_forge_resolver`].

use std::sync::{Arc, Mutex};

use cxx::UniquePtr;

#[cxx::bridge(namespace = "rust_usd")]
mod ffi {
    unsafe extern "C++" {
        include!("usd_bridge.h");

        type Stage;
        type Prim;
        type Mesh;
        type VariantSet;
        type Primvar;
        type Material;
        type Shader;

        fn open_stage(path: &str) -> Result<UniquePtr<Stage>>;
        fn open_stage_with_load(path: &str, load_all: bool) -> Result<UniquePtr<Stage>>;
        fn open_for_painting(asset_path: &str, edit_layer_path: &str)
            -> Result<UniquePtr<Stage>>;

        fn copy_prim(prim: &Prim) -> UniquePtr<Prim>;
        fn copy_mesh(mesh: &Mesh) -> UniquePtr<Mesh>;
        fn install_preferred_resolver();
        fn register_plugin_directory(dir: &str);

        fn pseudo_root(self: &Stage) -> UniquePtr<Prim>;
        fn prim_at_path(self: &Stage, sdf_path: &str) -> UniquePtr<Prim>;
        fn all_meshes(self: &Stage) -> UniquePtr<CxxVector<Mesh>>;
        fn load_path(self: &Stage, sdf_path: &str);
        fn unload_path(self: &Stage, sdf_path: &str);
        fn save_edit_layer(self: &Stage);
        fn edit_layer_path(self: &Stage) -> String;

        fn define_prim(self: &Stage, sdf_path: &str, type_name: &str) -> UniquePtr<Prim>;
        fn create_material(self: &Stage, sdf_path: &str) -> UniquePtr<Material>;
        fn material_at_path(self: &Stage, sdf_path: &str) -> UniquePtr<Material>;

        fn path(self: &Prim) -> String;
        fn type_name(self: &Prim) -> String;
        fn is_mesh(self: &Prim) -> bool;
        fn children(self: &Prim) -> UniquePtr<CxxVector<Prim>>;
        fn as_mesh(self: &Prim) -> UniquePtr<Mesh>;
        fn variant_set_names(self: &Prim) -> UniquePtr<CxxVector<CxxString>>;
        fn get_variant_set(self: &Prim, name: &str) -> UniquePtr<VariantSet>;

        fn name(self: &VariantSet) -> String;
        fn variant_names(self: &VariantSet) -> UniquePtr<CxxVector<CxxString>>;
        fn selection(self: &VariantSet) -> String;
        fn set_selection(self: &VariantSet, variant_name: &str) -> bool;
        fn clear_selection(self: &VariantSet);
        fn has_authored_selection(self: &VariantSet) -> bool;

        fn prim_path(self: &Mesh) -> String;
        fn subdivision_scheme(self: &Mesh) -> String;
        fn points(self: &Mesh) -> UniquePtr<CxxVector<f32>>;
        fn face_vertex_counts(self: &Mesh) -> UniquePtr<CxxVector<i32>>;
        fn face_vertex_indices(self: &Mesh) -> UniquePtr<CxxVector<i32>>;
        fn normals(self: &Mesh) -> UniquePtr<CxxVector<f32>>;
        fn normals_interpolation(self: &Mesh) -> String;
        fn st(self: &Mesh) -> UniquePtr<CxxVector<f32>>;
        fn st_indices(self: &Mesh) -> UniquePtr<CxxVector<i32>>;
        fn local_to_world(self: &Mesh) -> UniquePtr<CxxVector<f32>>;
        fn bound_texture_paths(self: &Mesh) -> UniquePtr<CxxVector<CxxString>>;

        fn primvar_names(self: &Mesh) -> UniquePtr<CxxVector<CxxString>>;
        fn primvar(self: &Mesh, name: &str) -> UniquePtr<Primvar>;
        fn create_primvar_float(self: &Mesh, name: &str, values: &[f32], interpolation: &str) -> bool;
        fn create_primvar_int(self: &Mesh, name: &str, values: &[i32], interpolation: &str) -> bool;
        fn create_primvar_vec2f(self: &Mesh, name: &str, values: &[f32], interpolation: &str) -> bool;
        fn create_primvar_vec3f(self: &Mesh, name: &str, values: &[f32], interpolation: &str) -> bool;
        fn create_primvar_color3f(self: &Mesh, name: &str, values: &[f32], interpolation: &str) -> bool;
        fn remove_primvar(self: &Mesh, name: &str) -> bool;

        fn name(self: &Primvar) -> String;
        fn interpolation(self: &Primvar) -> String;
        fn type_name(self: &Primvar) -> String;
        fn has_authored_value(self: &Primvar) -> bool;
        fn as_float_array(self: &Primvar) -> UniquePtr<CxxVector<f32>>;
        fn as_int_array(self: &Primvar) -> UniquePtr<CxxVector<i32>>;
        fn as_vec2f_array(self: &Primvar) -> UniquePtr<CxxVector<f32>>;
        fn as_vec3f_array(self: &Primvar) -> UniquePtr<CxxVector<f32>>;

        fn bind_material(self: &Mesh, material: &Material) -> bool;

        fn path(self: &Material) -> String;
        fn create_shader(self: &Material, name: &str, shader_id: &str) -> UniquePtr<Shader>;
        fn connect_surface(self: &Material, source: &Shader) -> bool;

        fn path(self: &Shader) -> String;
        fn shader_id(self: &Shader) -> String;
        fn set_input_float(self: &Shader, name: &str, value: f32) -> bool;
        fn set_input_color3f(self: &Shader, name: &str, r: f32, g: f32, b: f32) -> bool;
        fn set_input_asset(self: &Shader, name: &str, asset_path: &str) -> bool;
        fn set_input_token(self: &Shader, name: &str, value: &str) -> bool;
        fn connect_input(self: &Shader, input_name: &str, source: &Shader, output_name: &str) -> bool;
        fn declare_output(self: &Shader, output_name: &str, type_name: &str) -> bool;
    }

    extern "Rust" {
        fn forge_uri_resolve(uri: &str) -> String;
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum InitialLoadSet {
    All,
    None,
}

pub trait AssetResolver: Send + Sync {
    fn resolve(&self, uri: &str) -> Option<String>;
}

static FORGE_RESOLVER: Mutex<Option<Arc<dyn AssetResolver>>> = Mutex::new(None);

pub fn install_forge_resolver<R: AssetResolver + 'static>(resolver: R) {
    *FORGE_RESOLVER.lock().expect("FORGE_RESOLVER mutex poisoned") = Some(Arc::new(resolver));

    if let Err(e) = ensure_plugin_registered() {
        eprintln!("rust-usd: failed to register resolver plugInfo: {}", e);
        return;
    }
    ffi::install_preferred_resolver();
}

fn forge_uri_resolve(uri: &str) -> String {
    let resolver = {
        let guard = FORGE_RESOLVER.lock().expect("FORGE_RESOLVER mutex poisoned");
        guard.clone()
    };
    match resolver.and_then(|r| r.resolve(uri)) {
        Some(path) => path,
        None => String::new(),
    }
}

fn ensure_plugin_registered() -> std::io::Result<()> {
    use std::sync::atomic::{AtomicBool, Ordering};
    static REGISTERED: AtomicBool = AtomicBool::new(false);
    if REGISTERED.swap(true, Ordering::AcqRel) {
        return Ok(());
    }

    let exe = std::env::current_exe()?;
    let parent = exe.parent().ok_or_else(|| {
        std::io::Error::new(std::io::ErrorKind::NotFound, "exe has no parent dir")
    })?;
    let plugin_dir = parent.join("rust_usd_resolver_plug");
    std::fs::create_dir_all(&plugin_dir)?;

    let exe_escaped = exe
        .to_string_lossy()
        .replace('\\', "\\\\")
        .replace('"', "\\\"");

    let plug_info = format!(
        r#"{{
  "Plugins": [
    {{
      "Info": {{
        "Types": {{
          "rust_usd::ForgeAwareResolver": {{
            "bases": ["ArDefaultResolver"]
          }}
        }}
      }},
      "Name": "rustUsdForgeResolver",
      "LibraryPath": "{}",
      "Type": "library"
    }}
  ]
}}
"#,
        exe_escaped
    );

    std::fs::write(plugin_dir.join("plugInfo.json"), plug_info)?;
    ffi::register_plugin_directory(plugin_dir.to_string_lossy().as_ref());
    Ok(())
}

pub struct Stage {
    inner: UniquePtr<ffi::Stage>,
}

impl Stage {
    pub fn open(path: &str) -> Result<Self, cxx::Exception> {
        let _ = ensure_plugin_registered();
        Ok(Self {
            inner: ffi::open_stage(path)?,
        })
    }

    pub fn open_with_load_set(path: &str, load: InitialLoadSet) -> Result<Self, cxx::Exception> {
        let _ = ensure_plugin_registered();
        let load_all = matches!(load, InitialLoadSet::All);
        Ok(Self {
            inner: ffi::open_stage_with_load(path, load_all)?,
        })
    }

    /// Open a stage rooted at `edit_layer_path`, with `asset_path` mounted as
    /// a sublayer. Edits are directed to `edit_layer_path` so the asset file
    /// is never mutated. If `edit_layer_path` doesn't exist, a new layer is
    /// created (in-memory until [`Stage::save_edit_layer`] is called).
    pub fn open_for_painting(
        asset_path: &str,
        edit_layer_path: &str,
    ) -> Result<Self, cxx::Exception> {
        let _ = ensure_plugin_registered();
        Ok(Self {
            inner: ffi::open_for_painting(asset_path, edit_layer_path)?,
        })
    }

    pub fn load(&self, sdf_path: &str) {
        self.inner.load_path(sdf_path);
    }

    pub fn unload(&self, sdf_path: &str) {
        self.inner.unload_path(sdf_path);
    }

    pub fn pseudo_root(&self) -> Prim {
        Prim {
            inner: self.inner.pseudo_root(),
        }
    }

    pub fn prim_at_path(&self, sdf_path: &str) -> Option<Prim> {
        let p = self.inner.prim_at_path(sdf_path);
        if p.is_null() {
            None
        } else {
            Some(Prim { inner: p })
        }
    }

    pub fn meshes(&self) -> Vec<Mesh> {
        self.inner
            .all_meshes()
            .iter()
            .map(|m| Mesh {
                inner: ffi::copy_mesh(m),
            })
            .collect()
    }

    /// Persist the current edit target's layer to its on-disk identifier.
    /// No-op when the edit target is anonymous (in-memory only).
    pub fn save_edit_layer(&self) {
        self.inner.save_edit_layer();
    }

    pub fn edit_layer_path(&self) -> String {
        self.inner.edit_layer_path()
    }

    /// Defines a typed prim at `sdf_path`. Idempotent if the prim already
    /// exists with a matching type. Useful for explicitly authoring a
    /// `Scope` parent before [`create_material`].
    pub fn define_prim(&self, sdf_path: &str, type_name: &str) -> Option<Prim> {
        let p = self.inner.define_prim(sdf_path, type_name);
        if p.is_null() {
            None
        } else {
            Some(Prim { inner: p })
        }
    }

    pub fn create_material(&self, sdf_path: &str) -> Option<Material> {
        let m = self.inner.create_material(sdf_path);
        if m.is_null() {
            None
        } else {
            Some(Material { inner: m })
        }
    }

    pub fn material_at_path(&self, sdf_path: &str) -> Option<Material> {
        let m = self.inner.material_at_path(sdf_path);
        if m.is_null() {
            None
        } else {
            Some(Material { inner: m })
        }
    }
}

pub struct Prim {
    inner: UniquePtr<ffi::Prim>,
}

impl Prim {
    pub fn path(&self) -> String {
        self.inner.path()
    }

    pub fn type_name(&self) -> String {
        self.inner.type_name()
    }

    pub fn is_mesh(&self) -> bool {
        self.inner.is_mesh()
    }

    pub fn as_mesh(&self) -> Option<Mesh> {
        let m = self.inner.as_mesh();
        if m.is_null() {
            None
        } else {
            Some(Mesh { inner: m })
        }
    }

    pub fn children(&self) -> Vec<Prim> {
        self.inner
            .children()
            .iter()
            .map(|p| Prim {
                inner: ffi::copy_prim(p),
            })
            .collect()
    }

    pub fn variant_set_names(&self) -> Vec<String> {
        self.inner
            .variant_set_names()
            .iter()
            .map(|s| s.to_string())
            .collect()
    }

    pub fn variant_set(&self, name: &str) -> Option<VariantSet> {
        let v = self.inner.get_variant_set(name);
        if v.is_null() {
            None
        } else {
            Some(VariantSet { inner: v })
        }
    }
}

pub struct VariantSet {
    inner: UniquePtr<ffi::VariantSet>,
}

impl VariantSet {
    pub fn name(&self) -> String {
        self.inner.name()
    }

    pub fn variants(&self) -> Vec<String> {
        self.inner
            .variant_names()
            .iter()
            .map(|s| s.to_string())
            .collect()
    }

    /// Returns the currently composed selection, or `None` if no opinion is
    /// authored anywhere in the layer stack.
    pub fn selection(&self) -> Option<String> {
        let s = self.inner.selection();
        if s.is_empty() {
            None
        } else {
            Some(s)
        }
    }

    /// Authors a selection in the current edit target. Returns `false` if the
    /// variant name doesn't exist in this set.
    pub fn set_selection(&self, variant: &str) -> bool {
        self.inner.set_selection(variant)
    }

    pub fn clear_selection(&self) {
        self.inner.clear_selection();
    }

    pub fn has_authored_selection(&self) -> bool {
        self.inner.has_authored_selection()
    }
}

pub struct Mesh {
    inner: UniquePtr<ffi::Mesh>,
}

impl Mesh {
    pub fn prim_path(&self) -> String {
        self.inner.prim_path()
    }

    pub fn subdivision_scheme(&self) -> String {
        self.inner.subdivision_scheme()
    }

    pub fn points(&self) -> Vec<f32> {
        cxx_to_vec(self.inner.points())
    }

    pub fn face_vertex_counts(&self) -> Vec<i32> {
        cxx_to_vec(self.inner.face_vertex_counts())
    }

    pub fn face_vertex_indices(&self) -> Vec<i32> {
        cxx_to_vec(self.inner.face_vertex_indices())
    }

    pub fn normals(&self) -> Vec<f32> {
        cxx_to_vec(self.inner.normals())
    }

    /// `"vertex"` / `"faceVarying"` / `"varying"` etc — load-bearing for
    /// downstream consumers that need to know whether the flat normals
    /// array is per-vertex or per-face-vertex.
    pub fn normals_interpolation(&self) -> String {
        self.inner.normals_interpolation()
    }

    pub fn st(&self) -> Vec<f32> {
        cxx_to_vec(self.inner.st())
    }

    pub fn st_indices(&self) -> Vec<i32> {
        cxx_to_vec(self.inner.st_indices())
    }

    pub fn local_to_world(&self) -> [[f32; 4]; 4] {
        let v = self.inner.local_to_world();
        let mut m = [[0.0f32; 4]; 4];
        for r in 0..4 {
            for c in 0..4 {
                m[r][c] = *v.get(r * 4 + c).expect("local_to_world returned <16 floats");
            }
        }
        m
    }

    pub fn bound_texture_paths(&self) -> Vec<String> {
        self.inner
            .bound_texture_paths()
            .iter()
            .map(|s| s.to_string())
            .collect()
    }

    /// All primvar names on this mesh (without the `"primvars:"` prefix).
    pub fn primvar_names(&self) -> Vec<String> {
        self.inner
            .primvar_names()
            .iter()
            .map(|s| s.to_string())
            .collect()
    }

    pub fn primvar(&self, name: &str) -> Option<Primvar> {
        let p = self.inner.primvar(name);
        if p.is_null() {
            None
        } else {
            Some(Primvar { inner: p })
        }
    }

    /// Author `primvars:<name>` of type `float[]`. Writes go to the stage's
    /// current edit target; pair with [`Stage::open_for_painting`].
    pub fn create_primvar_float(&self, name: &str, values: &[f32], interpolation: &str) -> bool {
        self.inner.create_primvar_float(name, values, interpolation)
    }

    pub fn create_primvar_int(&self, name: &str, values: &[i32], interpolation: &str) -> bool {
        self.inner.create_primvar_int(name, values, interpolation)
    }

    /// Author a `float2[]` primvar. `values.len()` must be a multiple of 2.
    pub fn create_primvar_vec2f(&self, name: &str, values: &[f32], interpolation: &str) -> bool {
        self.inner.create_primvar_vec2f(name, values, interpolation)
    }

    /// Author a `float3[]` primvar. `values.len()` must be a multiple of 3.
    pub fn create_primvar_vec3f(&self, name: &str, values: &[f32], interpolation: &str) -> bool {
        self.inner.create_primvar_vec3f(name, values, interpolation)
    }

    /// Author a `color3f[]` primvar. `values.len()` must be a multiple of 3.
    pub fn create_primvar_color3f(&self, name: &str, values: &[f32], interpolation: &str) -> bool {
        self.inner.create_primvar_color3f(name, values, interpolation)
    }

    pub fn remove_primvar(&self, name: &str) -> bool {
        self.inner.remove_primvar(name)
    }

    /// Bind this mesh to a material via UsdShadeMaterialBindingAPI.
    pub fn bind_material(&self, material: &Material) -> bool {
        self.inner.bind_material(&material.inner)
    }
}

pub struct Material {
    inner: UniquePtr<ffi::Material>,
}

impl Material {
    pub fn path(&self) -> String {
        self.inner.path()
    }

    /// Define a child shader under this material with `info:id = shader_id`
    /// (e.g. `"UsdPreviewSurface"`, `"UsdUVTexture"`, `"UsdPrimvarReader_float2"`).
    pub fn create_shader(&self, name: &str, shader_id: &str) -> Option<Shader> {
        let s = self.inner.create_shader(name, shader_id);
        if s.is_null() {
            None
        } else {
            Some(Shader { inner: s })
        }
    }

    /// Wire this material's `outputs:surface` to `source.outputs:surface`.
    pub fn connect_surface(&self, source: &Shader) -> bool {
        self.inner.connect_surface(&source.inner)
    }
}

pub struct Shader {
    inner: UniquePtr<ffi::Shader>,
}

impl Shader {
    pub fn path(&self) -> String {
        self.inner.path()
    }

    pub fn shader_id(&self) -> String {
        self.inner.shader_id()
    }

    pub fn set_input_float(&self, name: &str, value: f32) -> bool {
        self.inner.set_input_float(name, value)
    }

    pub fn set_input_color3f(&self, name: &str, r: f32, g: f32, b: f32) -> bool {
        self.inner.set_input_color3f(name, r, g, b)
    }

    /// Set an asset-typed input (e.g. `inputs:file = @./tex.<UDIM>.png@`).
    pub fn set_input_asset(&self, name: &str, asset_path: &str) -> bool {
        self.inner.set_input_asset(name, asset_path)
    }

    pub fn set_input_token(&self, name: &str, value: &str) -> bool {
        self.inner.set_input_token(name, value)
    }

    /// Connect `inputs:<input_name>` on this shader to
    /// `source.outputs:<output_name>`. The input is auto-created as `color3f`
    /// if it didn't already exist — for non-color3f inputs, declare with
    /// `set_input_*` first.
    pub fn connect_input(&self, input_name: &str, source: &Shader, output_name: &str) -> bool {
        self.inner
            .connect_input(input_name, &source.inner, output_name)
    }

    /// Declare an output. `type_name` accepts: `"float"`, `"float2"`,
    /// `"float3"`, `"color3f"`, `"asset"`, `"token"`, `"int"`, `"string"`.
    pub fn declare_output(&self, output_name: &str, type_name: &str) -> bool {
        self.inner.declare_output(output_name, type_name)
    }
}

pub struct Primvar {
    inner: UniquePtr<ffi::Primvar>,
}

impl Primvar {
    pub fn name(&self) -> String {
        self.inner.name()
    }

    pub fn interpolation(&self) -> String {
        self.inner.interpolation()
    }

    /// Type token, e.g. `"float[]"`, `"float3[]"`, `"color3f[]"`.
    pub fn type_name(&self) -> String {
        self.inner.type_name()
    }

    pub fn has_authored_value(&self) -> bool {
        self.inner.has_authored_value()
    }

    /// Returns `Vec<f32>` if this primvar's type is `float[]`, otherwise empty.
    pub fn as_float_array(&self) -> Vec<f32> {
        cxx_to_vec(self.inner.as_float_array())
    }

    pub fn as_int_array(&self) -> Vec<i32> {
        cxx_to_vec(self.inner.as_int_array())
    }

    /// Flattened `[u, v, u, v, ...]` if type is `float2[]`, else empty.
    pub fn as_vec2f_array(&self) -> Vec<f32> {
        cxx_to_vec(self.inner.as_vec2f_array())
    }

    /// Flattened `[x, y, z, ...]` if type is `float3[]` or `color3f[]`, else empty.
    pub fn as_vec3f_array(&self) -> Vec<f32> {
        cxx_to_vec(self.inner.as_vec3f_array())
    }
}

fn cxx_to_vec<T: cxx::vector::VectorElement + Copy>(v: UniquePtr<cxx::CxxVector<T>>) -> Vec<T> {
    v.iter().copied().collect()
}