kcl-lib 0.2.186

use anyhow::Result;
use kcmc::ImportFile;
use kcmc::ModelingCmd;
use kcmc::coord::KITTYCAD;
use kcmc::coord::System;
use kcmc::each_cmd as mcmd;
use kcmc::format::InputFormat3d;
use kcmc::shared::FileImportFormat;
use kcmc::units::UnitLength;
use kittycad_modeling_cmds as kcmc;
use serde::Deserialize;
use serde::Serialize;
use uuid::Uuid;

use crate::SourceRange;
use crate::errors::KclError;
use crate::errors::KclErrorDetails;
use crate::execution::ExecState;
use crate::execution::ExecutorContext;
use crate::execution::ImportedGeometry;
use crate::execution::ModelingCmdMeta;
use crate::execution::annotations;
use crate::execution::typed_path::TypedPath;
use crate::execution::types::length_from_str;
use crate::parsing::ast::types::Annotation;
use crate::parsing::ast::types::Node;
use crate::unit_conversion::ToKcmc;

// Zoo co-ordinate system.
//
// * Forward: -Y
// * Up: +Z
// * Handedness: Right
pub const ZOO_COORD_SYSTEM: System = *KITTYCAD;

pub async fn import_foreign(
    file_path: &TypedPath,
    format: Option<InputFormat3d>,
    exec_state: &mut ExecState,
    ctxt: &ExecutorContext,
    source_range: SourceRange,
) -> Result<PreImportedGeometry, KclError> {
    // Make sure the file exists.
    if !ctxt.fs.exists(file_path, source_range).await? {
        return Err(KclError::new_semantic(KclErrorDetails::new(
            format!("File `{}` does not exist.", file_path.display()),
            vec![source_range],
        )));
    }

    let ext_format = get_import_format_from_extension(&file_path.to_string_lossy())
        .map_err(|e| KclError::new_semantic(KclErrorDetails::new(e.to_string(), vec![source_range])))?;

    // Get the format type from the extension of the file.
    let format = if let Some(format) = format {
        // Validate the given format with the extension format.
        validate_extension_format(&file_path.to_string_lossy(), ext_format, format.clone())
            .map_err(|e| KclError::new_semantic(KclErrorDetails::new(e.to_string(), vec![source_range])))?;
        format
    } else {
        ext_format
    };

    // Get the file contents for each file path.
    let file_contents = ctxt
        .fs
        .read(file_path, source_range)
        .await
        .map_err(|e| KclError::new_semantic(KclErrorDetails::new(e.to_string(), vec![source_range])))?;

    // We want the file_path to be without the parent.
    let file_name = file_path.file_name().ok_or_else(|| {
        KclError::new_semantic(KclErrorDetails::new(
            format!("Could not get the file name from the path `{}`", file_path.display()),
            vec![source_range],
        ))
    })?;
    let mut import_files = vec![
        kcmc::ImportFile::builder()
            .path(file_name.to_string())
            .data(file_contents.clone())
            .build(),
    ];

    // In the case of a gltf importing a bin file we need to handle that! and figure out where the
    // file is relative to our current file.
    if let InputFormat3d::Gltf(..) = format {
        // Check if the file is a binary gltf file, in that case we don't need to import the bin
        // file.
        if !file_contents.starts_with(b"glTF") {
            let json = gltf_json::Root::from_slice(&file_contents)
                .map_err(|e| KclError::new_semantic(KclErrorDetails::new(e.to_string(), vec![source_range])))?;

            // Read the gltf file and check if there is a bin file.
            for buffer in json.buffers.iter() {
                if let Some(uri) = &buffer.uri
                    && !uri.starts_with("data:")
                {
                    // We want this path relative to the file_path given.
                    let bin_path = file_path.parent().map(|p| p.join(uri)).ok_or_else(|| {
                        KclError::new_semantic(KclErrorDetails::new(
                            format!("Could not get the parent path of the file `{}`", file_path.display()),
                            vec![source_range],
                        ))
                    })?;

                    let bin_contents =
                        ctxt.fs.read(&bin_path, source_range).await.map_err(|e| {
                            KclError::new_semantic(KclErrorDetails::new(e.to_string(), vec![source_range]))
                        })?;

                    import_files.push(ImportFile::builder().path(uri.to_string()).data(bin_contents).build());
                }
            }
        }
    }
    Ok(PreImportedGeometry {
        id: exec_state.next_uuid(),
        source_range,
        command: mcmd::ImportFiles::builder()
            .files(import_files.clone())
            .format(format)
            .build(),
    })
}

pub(super) fn format_from_annotations(
    annotations: &[Node<Annotation>],
    path: &TypedPath,
    import_source_range: SourceRange,
) -> Result<Option<InputFormat3d>, KclError> {
    if annotations.is_empty() {
        return Ok(None);
    }

    let props = annotations.iter().flat_map(|a| a.properties.as_deref().unwrap_or(&[]));

    let mut result = None;
    for p in props.clone() {
        if p.key.name == annotations::IMPORT_FORMAT {
            result = Some(
                get_import_format_from_extension(annotations::expect_ident(&p.value)?).map_err(|_| {
                    KclError::new_semantic(KclErrorDetails::new(
                        format!(
                            "Unknown format for import, expected one of: {}",
                            crate::IMPORT_FILE_EXTENSIONS.join(", ")
                        ),
                        vec![p.as_source_range()],
                    ))
                })?,
            );
            break;
        }
    }

    let mut result = result
        .or_else(|| get_import_format_from_extension(&path.to_string_lossy()).ok())
        .ok_or(KclError::new_semantic(KclErrorDetails::new(
            "Unknown or missing extension, and no specified format for imported file".to_owned(),
            vec![import_source_range],
        )))?;

    for p in props {
        match p.key.name.as_str() {
            annotations::IMPORT_COORDS => {
                set_coords(&mut result, annotations::expect_ident(&p.value)?, p.as_source_range())?;
            }
            annotations::IMPORT_LENGTH_UNIT => {
                set_length_unit(&mut result, annotations::expect_ident(&p.value)?, p.as_source_range())?;
            }
            annotations::IMPORT_TARGET_REPRESENTATION => {
                set_target_representation(&mut result, annotations::expect_ident(&p.value)?, p.as_source_range())?;
            }
            annotations::IMPORT_FORMAT => {}
            _ => {
                return Err(KclError::new_semantic(KclErrorDetails::new(
                    format!(
                        "Unexpected annotation for import, expected one of: {}, {}, {}, {}",
                        annotations::IMPORT_FORMAT,
                        annotations::IMPORT_COORDS,
                        annotations::IMPORT_LENGTH_UNIT,
                        annotations::IMPORT_TARGET_REPRESENTATION
                    ),
                    vec![p.as_source_range()],
                )));
            }
        }
    }

    Ok(Some(result))
}

fn set_coords(fmt: &mut InputFormat3d, coords_str: &str, source_range: SourceRange) -> Result<(), KclError> {
    let mut coords = None;
    for (name, val) in annotations::IMPORT_COORDS_VALUES {
        if coords_str == name {
            coords = Some(*val);
        }
    }

    let Some(coords) = coords else {
        return Err(KclError::new_semantic(KclErrorDetails::new(
            format!(
                "Unknown coordinate system: {coords_str}, expected one of: {}",
                annotations::IMPORT_COORDS_VALUES
                    .iter()
                    .map(|(n, _)| *n)
                    .collect::<Vec<_>>()
                    .join(", ")
            ),
            vec![source_range],
        )));
    };

    match fmt {
        InputFormat3d::Obj(opts) => opts.coords = coords,
        InputFormat3d::Ply(opts) => opts.coords = coords,
        InputFormat3d::Step(opts) => opts.coords = coords,
        InputFormat3d::Stl(opts) => opts.coords = coords,
        _ => {
            return Err(KclError::new_semantic(KclErrorDetails::new(
                format!(
                    "`{}` option cannot be applied to the specified format",
                    annotations::IMPORT_COORDS
                ),
                vec![source_range],
            )));
        }
    }

    Ok(())
}

fn set_length_unit(fmt: &mut InputFormat3d, units_str: &str, source_range: SourceRange) -> Result<(), KclError> {
    let units = length_from_str(units_str, source_range)?;

    match fmt {
        InputFormat3d::Obj(opts) => opts.units = units.to_kcmc(),
        InputFormat3d::Ply(opts) => opts.units = units.to_kcmc(),
        InputFormat3d::Stl(opts) => opts.units = units.to_kcmc(),
        _ => {
            return Err(KclError::new_semantic(KclErrorDetails::new(
                format!(
                    "`{}` option cannot be applied to the specified format",
                    annotations::IMPORT_LENGTH_UNIT
                ),
                vec![source_range],
            )));
        }
    }

    Ok(())
}

fn set_target_representation(
    fmt: &mut InputFormat3d,
    representation: &str,
    source_range: SourceRange,
) -> Result<(), KclError> {
    const MESH: &str = "mesh";
    const BREP: &str = "brep";
    const ALL_OPTIONS: [&str; 2] = [MESH, BREP];
    let target_representation = match representation {
        MESH => kcmc::format::step::TargetRepresentation::Mesh,
        BREP => kcmc::format::step::TargetRepresentation::Brep,
        _ => {
            return Err(KclError::new_semantic(KclErrorDetails::new(
                format!(
                    "Unknown target representation: {representation}, expected one of: {}",
                    ALL_OPTIONS.join(", ")
                ),
                vec![source_range],
            )));
        }
    };

    let InputFormat3d::Step(opts) = fmt else {
        return Err(KclError::new_semantic(KclErrorDetails::new(
            format!(
                "`{}` option cannot be applied to the specified format",
                annotations::IMPORT_TARGET_REPRESENTATION
            ),
            vec![source_range],
        )));
    };

    opts.target_representation = target_representation;
    Ok(())
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
pub struct PreImportedGeometry {
    id: Uuid,
    command: mcmd::ImportFiles,
    pub source_range: SourceRange,
}

pub async fn send_to_engine(
    pre: PreImportedGeometry,
    exec_state: &mut ExecState,
    ctxt: &ExecutorContext,
) -> Result<ImportedGeometry, KclError> {
    let imported_geometry = ImportedGeometry::new(
        pre.id,
        pre.command.files.iter().map(|f| f.path.to_string()).collect(),
        vec![pre.source_range.into()],
    );

    exec_state
        .async_modeling_cmd(
            ModelingCmdMeta::with_id(exec_state, ctxt, pre.source_range, pre.id),
            &ModelingCmd::from(pre.command.clone()),
        )
        .await?;

    Ok(imported_geometry)
}

/// Get the source format from a file path, extension, or canonical format name.
fn get_import_format_from_extension(ext: &str) -> Result<InputFormat3d> {
    let format = crate::import_format::import_format_from_path(ext)
        .or_else(|| crate::import_format::import_format_from_name(ext))
        .ok_or_else(|| {
            anyhow::anyhow!(
                "unknown source format for file extension: {ext}. Try setting the `--src-format` flag explicitly or use a valid format."
            )
        })?;

    // Make the default units millimeters.
    let ul = UnitLength::Millimeters;

    // Zoo co-ordinate system.
    //
    // * Forward: -Y
    // * Up: +Z
    // * Handedness: Right
    match format {
        FileImportFormat::Acis => Ok(InputFormat3d::Acis(kcmc::format::acis::import::Options::default())),
        FileImportFormat::Catia => Ok(InputFormat3d::Catia(kcmc::format::catia::import::Options::default())),
        FileImportFormat::Creo => Ok(InputFormat3d::Creo(kcmc::format::creo::import::Options::default())),
        FileImportFormat::Step => Ok(InputFormat3d::Step(
            kcmc::format::step::import::Options::builder()
                .coords(ZOO_COORD_SYSTEM)
                .split_closed_faces(false)
                .build(),
        )),
        FileImportFormat::Stl => Ok(InputFormat3d::Stl(
            kcmc::format::stl::import::Options::builder()
                .coords(ZOO_COORD_SYSTEM)
                .units(ul)
                .build(),
        )),
        FileImportFormat::Obj => Ok(InputFormat3d::Obj(
            kcmc::format::obj::import::Options::builder()
                .coords(ZOO_COORD_SYSTEM)
                .units(ul)
                .build(),
        )),
        FileImportFormat::Gltf => Ok(InputFormat3d::Gltf(kcmc::format::gltf::import::Options::default())),
        FileImportFormat::Inventor => Ok(InputFormat3d::Inventor(
            kcmc::format::inventor::import::Options::default(),
        )),
        FileImportFormat::Nx => Ok(InputFormat3d::Nx(kcmc::format::nx::import::Options::default())),
        FileImportFormat::Ply => Ok(InputFormat3d::Ply(
            kcmc::format::ply::import::Options::builder()
                .coords(ZOO_COORD_SYSTEM)
                .units(ul)
                .build(),
        )),
        FileImportFormat::Fbx => Ok(InputFormat3d::Fbx(kcmc::format::fbx::import::Options::default())),
        FileImportFormat::Parasolid => Ok(InputFormat3d::Parasolid(
            kcmc::format::parasolid::import::Options::default(),
        )),
        FileImportFormat::Sldprt => Ok(InputFormat3d::Sldprt(
            kcmc::format::sldprt::import::Options::builder()
                .split_closed_faces(false)
                .build(),
        )),
        other => anyhow::bail!("Unknown format {other}"),
    }
}

fn validate_extension_format(path: &str, ext: InputFormat3d, given: InputFormat3d) -> Result<()> {
    if crate::import_format::import_path_supports_format(path, given.clone().into()) {
        return Ok(());
    }

    anyhow::bail!(
        "The given format does not match the file extension. Expected: `{}`, Given: `{}`",
        ext.name(),
        given.name()
    )
}

#[cfg(test)]
mod test {
    use super::*;

    macro_rules! test_import_format_from_extension {
        ($name:ident, $xtn:expr, $fmt:path) => {
            #[test]
            fn $name() {
                let x = get_import_format_from_extension($xtn).unwrap();
                assert!(matches!(x, $fmt(_)));
            }
        };
    }

    test_import_format_from_extension!(test_xtn_step, "step", InputFormat3d::Step);
    test_import_format_from_extension!(test_xtn_stp, "stp", InputFormat3d::Step);
    test_import_format_from_extension!(test_xtn_step_upper, "STEP", InputFormat3d::Step);
    test_import_format_from_extension!(test_xtn_step_spongebob, "STeP", InputFormat3d::Step);
    test_import_format_from_extension!(test_xtn_fbx, "fbx", InputFormat3d::Fbx);
    test_import_format_from_extension!(test_xtn_gltf, "gltf", InputFormat3d::Gltf);
    test_import_format_from_extension!(test_xtn_sat, "sat", InputFormat3d::Acis);
    test_import_format_from_extension!(test_xtn_sab, "sab", InputFormat3d::Acis);
    test_import_format_from_extension!(test_xtn_catpart, "catpart", InputFormat3d::Catia);
    test_import_format_from_extension!(test_xtn_ipt, "ipt", InputFormat3d::Inventor);
    test_import_format_from_extension!(test_xtn_prt, "prt", InputFormat3d::Nx);
    test_import_format_from_extension!(test_xtn_obj, "obj", InputFormat3d::Obj);
    test_import_format_from_extension!(test_xtn_x_t, "x_t", InputFormat3d::Parasolid);
    test_import_format_from_extension!(test_xtn_x_b, "x_b", InputFormat3d::Parasolid);
    test_import_format_from_extension!(test_xtn_ply, "ply", InputFormat3d::Ply);
    test_import_format_from_extension!(test_xtn_sldprt, "sldprt", InputFormat3d::Sldprt);
    test_import_format_from_extension!(test_xtn_stl, "stl", InputFormat3d::Stl);
    test_import_format_from_extension!(test_format_creo, "creo", InputFormat3d::Creo);

    #[test]
    fn versioned_creo_import_format_from_path() {
        for path in ["part.prt.1", "nested/part.PRT.23"] {
            let format = get_import_format_from_extension(path).unwrap();
            assert_eq!(
                format,
                InputFormat3d::Creo(kcmc::format::creo::import::Options::default())
            );
        }
    }

    #[test]
    fn unversioned_prt_accepts_explicit_creo_format() {
        validate_extension_format(
            "part.prt",
            InputFormat3d::Nx(kcmc::format::nx::import::Options::default()),
            InputFormat3d::Creo(kcmc::format::creo::import::Options::default()),
        )
        .unwrap();
    }

    #[test]
    fn annotations() {
        let (_, issues) = crate::Program::parse("@(targetRepresentation = mesh)\nimport '../foo.step' as foo")
            .expect("program should parse");
        assert_eq!(issues.len(), 0);

        // no annotations
        assert!(
            format_from_annotations(&[], &TypedPath::from("../foo.txt"), SourceRange::default(),)
                .unwrap()
                .is_none()
        );

        // no format, no options
        let text = "@()\nimport '../foo.gltf' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.gltf"), SourceRange::default())
            .unwrap()
            .unwrap();
        assert_eq!(fmt, InputFormat3d::Gltf(kcmc::format::gltf::import::Options::default()));

        // Creo format inferred from a versioned part path.
        let text = "@()\nimport '../foo.prt.3' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.prt.3"), SourceRange::default())
            .unwrap()
            .unwrap();
        assert_eq!(fmt, InputFormat3d::Creo(kcmc::format::creo::import::Options::default()));

        // Creo's canonical format name remains valid in annotations.
        let text = "@(format = creo)\nimport '../foo.prt.3' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.prt.3"), SourceRange::default())
            .unwrap()
            .unwrap();
        assert_eq!(fmt, InputFormat3d::Creo(kcmc::format::creo::import::Options::default()));

        // format, no options
        let text = "@(format = gltf)\nimport '../foo.txt' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.txt"), SourceRange::default())
            .unwrap()
            .unwrap();
        assert_eq!(fmt, InputFormat3d::Gltf(kcmc::format::gltf::import::Options::default()));

        // format, no extension (wouldn't parse but might some day)
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo"), SourceRange::default())
            .unwrap()
            .unwrap();
        assert_eq!(fmt, InputFormat3d::Gltf(kcmc::format::gltf::import::Options::default()));

        // format, options
        let text = "@(format = obj, coords = vulkan, lengthUnit = ft)\nimport '../foo.txt' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.txt"), SourceRange::default())
            .unwrap()
            .unwrap();
        assert_eq!(
            fmt,
            InputFormat3d::Obj(
                kcmc::format::obj::import::Options::builder()
                    .coords(*kcmc::coord::VULKAN)
                    .units(kcmc::units::UnitLength::Feet)
                    .build()
            )
        );

        // STEP defaults are unchanged when no target representation is given.
        let text = "@()\nimport '../foo.step' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.step"), SourceRange::default())
            .unwrap()
            .unwrap();
        assert_eq!(fmt, get_import_format_from_extension("step").unwrap());

        // STEP target representation.
        let text = "@(targetRepresentation = mesh)\nimport '../foo.step' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.step"), SourceRange::default())
            .unwrap()
            .unwrap();
        let InputFormat3d::Step(opts) = fmt else {
            panic!("expected STEP import options");
        };
        assert_eq!(
            opts.target_representation,
            kcmc::format::step::TargetRepresentation::Mesh
        );

        let text = "@(targetRepresentation = brep)\nimport '../foo.step' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.step"), SourceRange::default())
            .unwrap()
            .unwrap();
        let InputFormat3d::Step(opts) = fmt else {
            panic!("expected STEP import options");
        };
        assert_eq!(
            opts.target_representation,
            kcmc::format::step::TargetRepresentation::Brep,
        );

        // no format, options
        let text = "@(coords = vulkan, lengthUnit = ft)\nimport '../foo.obj' as foo";
        let parsed = crate::Program::parse_no_errs(text).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let fmt = format_from_annotations(attrs, &TypedPath::from("../foo.obj"), SourceRange::default())
            .unwrap()
            .unwrap();
        assert_eq!(
            fmt,
            InputFormat3d::Obj(
                kcmc::format::obj::import::Options::builder()
                    .coords(*kcmc::coord::VULKAN)
                    .units(kcmc::units::UnitLength::Feet)
                    .build()
            )
        );

        // err - format, options, but no options for specified format
        assert_annotation_error(
            "@(format = gltf, lengthUnit = ft)\nimport '../foo.txt' as foo",
            "../foo.txt",
            "`lengthUnit` option cannot be applied",
        );
        assert_annotation_error(
            "@(targetRepresentation = brep)\nimport '../foo.obj' as foo",
            "../foo.obj",
            "`targetRepresentation` option cannot be applied",
        );
        assert_annotation_error(
            "@(targetRepresentation = voxels)\nimport '../foo.step' as foo",
            "../foo.step",
            "Unknown target representation",
        );
        // err - no format, options, but no options for specified format
        assert_annotation_error(
            "@(lengthUnit = ft)\nimport '../foo.gltf' as foo",
            "../foo.gltf",
            "lengthUnit` option cannot be applied",
        );
        // err - bad option
        assert_annotation_error(
            "@(format = obj, coords = vulkan, lengthUni = ft)\nimport '../foo.txt' as foo",
            "../foo.txt",
            "Unexpected annotation",
        );
        // err - bad format
        assert_annotation_error(
            "@(format = foo)\nimport '../foo.txt' as foo",
            "../foo.txt",
            "Unknown format for import",
        );
        // err - bad coord value
        assert_annotation_error(
            "@(format = gltf, coords = north)\nimport '../foo.txt' as foo",
            "../foo.txt",
            "Unknown coordinate system",
        );
        // err - bad unit value
        assert_annotation_error(
            "@(format = gltf, lengthUnit = gallons)\nimport '../foo.txt' as foo",
            "../foo.txt",
            "Unexpected value for length units",
        );
    }

    #[track_caller]
    fn assert_annotation_error(src: &str, path: &str, expected: &str) {
        let parsed = crate::Program::parse_no_errs(src).unwrap().ast;
        let attrs = parsed.body[0].get_attrs();
        let err = format_from_annotations(attrs, &TypedPath::from(path), SourceRange::default()).unwrap_err();
        assert!(
            err.message().contains(expected),
            "Expected: `{expected}`, found `{}`",
            err.message()
        );
    }
}