modeling-app/src/wasm-lib/kcl/src/kcl_value.rs

use std::collections::HashMap;

use anyhow::Result;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

use crate::{
    ast::types::{FunctionExpression, KclNone, TagDeclarator, TagNode},
    errors::KclErrorDetails,
    exec::{ProgramMemory, Sketch},
    executor::{Face, ImportedGeometry, MemoryFunction, Metadata, Plane, SketchSet, Solid, SolidSet, TagIdentifier},
    std::{args::Arg, FnAsArg},
    ExecState, ExecutorContext, KclError, SourceRange,
};

pub type KclObjectFields = HashMap<String, KclValue>;

/// Any KCL value.
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, ts_rs::TS, JsonSchema)]
#[ts(export)]
#[serde(tag = "type")]
pub enum KclValue {
    Uuid {
        value: ::uuid::Uuid,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
    Bool {
        value: bool,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
    Number {
        value: f64,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
    Int {
        value: i64,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
    String {
        value: String,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
    Array {
        value: Vec<KclValue>,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
    Object {
        value: KclObjectFields,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
    TagIdentifier(Box<TagIdentifier>),
    TagDeclarator(crate::ast::types::BoxNode<TagDeclarator>),
    Plane(Box<Plane>),
    Face(Box<Face>),
    Sketch {
        value: Box<Sketch>,
    },
    Sketches {
        value: Vec<Box<Sketch>>,
    },
    Solid(Box<Solid>),
    Solids {
        value: Vec<Box<Solid>>,
    },
    ImportedGeometry(ImportedGeometry),
    #[ts(skip)]
    Function {
        #[serde(skip)]
        func: Option<MemoryFunction>,
        expression: crate::ast::types::BoxNode<FunctionExpression>,
        memory: Box<ProgramMemory>,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
    KclNone {
        value: KclNone,
        #[serde(rename = "__meta")]
        meta: Vec<Metadata>,
    },
}

impl KclValue {
    pub(crate) fn metadata(&self) -> Vec<Metadata> {
        match self {
            KclValue::Uuid { value: _, meta } => meta.clone(),
            KclValue::Bool { value: _, meta } => meta.clone(),
            KclValue::Number { value: _, meta } => meta.clone(),
            KclValue::Int { value: _, meta } => meta.clone(),
            KclValue::String { value: _, meta } => meta.clone(),
            KclValue::Array { value: _, meta } => meta.clone(),
            KclValue::Object { value: _, meta } => meta.clone(),
            KclValue::TagIdentifier(x) => x.meta.clone(),
            KclValue::TagDeclarator(x) => vec![x.metadata()],
            KclValue::Plane(x) => x.meta.clone(),
            KclValue::Face(x) => x.meta.clone(),
            KclValue::Sketch { value } => value.meta.clone(),
            KclValue::Sketches { value } => value.iter().flat_map(|sketch| &sketch.meta).copied().collect(),
            KclValue::Solid(x) => x.meta.clone(),
            KclValue::Solids { value } => value.iter().flat_map(|sketch| &sketch.meta).copied().collect(),
            KclValue::ImportedGeometry(x) => x.meta.clone(),
            KclValue::Function { meta, .. } => meta.clone(),
            KclValue::KclNone { meta, .. } => meta.clone(),
        }
    }

    pub(crate) fn get_solid_set(&self) -> Result<SolidSet> {
        match self {
            KclValue::Solid(e) => Ok(SolidSet::Solid(e.clone())),
            KclValue::Solids { value } => Ok(SolidSet::Solids(value.clone())),
            KclValue::Array { value, .. } => {
                let solids: Vec<_> = value
                    .iter()
                    .enumerate()
                    .map(|(i, v)| {
                        v.as_solid().map(|v| v.to_owned()).map(Box::new).ok_or_else(|| {
                            anyhow::anyhow!(
                                "expected this array to only contain solids, but element {i} was actually {}",
                                v.human_friendly_type()
                            )
                        })
                    })
                    .collect::<Result<_, _>>()?;
                Ok(SolidSet::Solids(solids))
            }
            _ => anyhow::bail!("Not a solid or solids: {:?}", self),
        }
    }

    /// Human readable type name used in error messages.  Should not be relied
    /// on for program logic.
    pub(crate) fn human_friendly_type(&self) -> &'static str {
        match self {
            KclValue::Uuid { .. } => "Unique ID (uuid)",
            KclValue::TagDeclarator(_) => "TagDeclarator",
            KclValue::TagIdentifier(_) => "TagIdentifier",
            KclValue::Solid(_) => "Solid",
            KclValue::Solids { .. } => "Solids",
            KclValue::Sketch { .. } => "Sketch",
            KclValue::Sketches { .. } => "Sketches",
            KclValue::ImportedGeometry(_) => "ImportedGeometry",
            KclValue::Function { .. } => "Function",
            KclValue::Plane(_) => "Plane",
            KclValue::Face(_) => "Face",
            KclValue::Bool { .. } => "boolean (true/false value)",
            KclValue::Number { .. } => "number",
            KclValue::Int { .. } => "integer",
            KclValue::String { .. } => "string (text)",
            KclValue::Array { .. } => "array (list)",
            KclValue::Object { .. } => "object",
            KclValue::KclNone { .. } => "None",
        }
    }

    pub(crate) fn is_function(&self) -> bool {
        matches!(self, KclValue::Function { .. })
    }
}

impl From<SketchSet> for KclValue {
    fn from(sg: SketchSet) -> Self {
        match sg {
            SketchSet::Sketch(value) => KclValue::Sketch { value },
            SketchSet::Sketches(value) => KclValue::Sketches { value },
        }
    }
}

impl From<Vec<Box<Sketch>>> for KclValue {
    fn from(sg: Vec<Box<Sketch>>) -> Self {
        KclValue::Sketches { value: sg }
    }
}

impl From<SolidSet> for KclValue {
    fn from(eg: SolidSet) -> Self {
        match eg {
            SolidSet::Solid(eg) => KclValue::Solid(eg),
            SolidSet::Solids(egs) => KclValue::Solids { value: egs },
        }
    }
}

impl From<Vec<Box<Solid>>> for KclValue {
    fn from(eg: Vec<Box<Solid>>) -> Self {
        if eg.len() == 1 {
            KclValue::Solid(eg[0].clone())
        } else {
            KclValue::Solids { value: eg }
        }
    }
}
impl From<KclValue> for Vec<SourceRange> {
    fn from(item: KclValue) -> Self {
        match item {
            KclValue::TagDeclarator(t) => vec![SourceRange([t.start, t.end, t.module_id.0 as usize])],
            KclValue::TagIdentifier(t) => to_vec_sr(&t.meta),
            KclValue::Solid(e) => to_vec_sr(&e.meta),
            KclValue::Solids { value } => value.iter().flat_map(|eg| to_vec_sr(&eg.meta)).collect(),
            KclValue::Sketch { value } => to_vec_sr(&value.meta),
            KclValue::Sketches { value } => value.iter().flat_map(|eg| to_vec_sr(&eg.meta)).collect(),
            KclValue::ImportedGeometry(i) => to_vec_sr(&i.meta),
            KclValue::Function { meta, .. } => to_vec_sr(&meta),
            KclValue::Plane(p) => to_vec_sr(&p.meta),
            KclValue::Face(f) => to_vec_sr(&f.meta),
            KclValue::Bool { meta, .. } => to_vec_sr(&meta),
            KclValue::Number { meta, .. } => to_vec_sr(&meta),
            KclValue::Int { meta, .. } => to_vec_sr(&meta),
            KclValue::String { meta, .. } => to_vec_sr(&meta),
            KclValue::Array { meta, .. } => to_vec_sr(&meta),
            KclValue::Object { meta, .. } => to_vec_sr(&meta),
            KclValue::Uuid { meta, .. } => to_vec_sr(&meta),
            KclValue::KclNone { meta, .. } => to_vec_sr(&meta),
        }
    }
}

fn to_vec_sr(meta: &[Metadata]) -> Vec<SourceRange> {
    meta.iter().map(|m| m.source_range).collect()
}

impl From<&KclValue> for Vec<SourceRange> {
    fn from(item: &KclValue) -> Self {
        match item {
            KclValue::TagDeclarator(t) => vec![SourceRange([t.start, t.end, t.module_id.0 as usize])],
            KclValue::TagIdentifier(t) => to_vec_sr(&t.meta),
            KclValue::Solid(e) => to_vec_sr(&e.meta),
            KclValue::Solids { value } => value.iter().flat_map(|eg| to_vec_sr(&eg.meta)).collect(),
            KclValue::Sketch { value } => to_vec_sr(&value.meta),
            KclValue::Sketches { value } => value.iter().flat_map(|eg| to_vec_sr(&eg.meta)).collect(),
            KclValue::ImportedGeometry(i) => to_vec_sr(&i.meta),
            KclValue::Function { meta, .. } => to_vec_sr(meta),
            KclValue::Plane(p) => to_vec_sr(&p.meta),
            KclValue::Face(f) => to_vec_sr(&f.meta),
            KclValue::Bool { meta, .. } => to_vec_sr(meta),
            KclValue::Number { meta, .. } => to_vec_sr(meta),
            KclValue::Int { meta, .. } => to_vec_sr(meta),
            KclValue::String { meta, .. } => to_vec_sr(meta),
            KclValue::Uuid { meta, .. } => to_vec_sr(meta),
            KclValue::Array { meta, .. } => to_vec_sr(meta),
            KclValue::Object { meta, .. } => to_vec_sr(meta),
            KclValue::KclNone { meta, .. } => to_vec_sr(meta),
        }
    }
}

impl KclValue {
    /// Put the number into a KCL value.
    pub fn from_number(f: f64, meta: Vec<Metadata>) -> Self {
        Self::Number { value: f, meta }
    }

    /// Put the point into a KCL value.
    pub fn from_point2d(p: [f64; 2], meta: Vec<Metadata>) -> Self {
        Self::Array {
            value: vec![
                Self::Number {
                    value: p[0],
                    meta: meta.clone(),
                },
                Self::Number {
                    value: p[1],
                    meta: meta.clone(),
                },
            ],
            meta,
        }
    }

    pub(crate) fn as_usize(&self) -> Option<usize> {
        match self {
            KclValue::Int { value, .. } if *value > 0 => Some(*value as usize),
            KclValue::Number { value, .. } => crate::try_f64_to_usize(*value),
            _ => None,
        }
    }

    pub fn as_int(&self) -> Option<i64> {
        match self {
            KclValue::Int { value, .. } => Some(*value),
            KclValue::Number { value, .. } => crate::try_f64_to_i64(*value),
            _ => None,
        }
    }

    pub fn as_object(&self) -> Option<&KclObjectFields> {
        if let KclValue::Object { value, meta: _ } = &self {
            Some(value)
        } else {
            None
        }
    }

    pub fn into_object(self) -> Option<KclObjectFields> {
        if let KclValue::Object { value, meta: _ } = self {
            Some(value)
        } else {
            None
        }
    }

    pub fn as_str(&self) -> Option<&str> {
        if let KclValue::String { value, meta: _ } = &self {
            Some(value)
        } else {
            None
        }
    }

    pub fn as_array(&self) -> Option<&[KclValue]> {
        if let KclValue::Array { value, meta: _ } = &self {
            Some(value)
        } else {
            None
        }
    }

    pub fn as_point2d(&self) -> Option<[f64; 2]> {
        let arr = self.as_array()?;
        if arr.len() != 2 {
            return None;
        }
        let x = arr[0].as_f64()?;
        let y = arr[1].as_f64()?;
        Some([x, y])
    }

    pub fn as_uuid(&self) -> Option<uuid::Uuid> {
        if let KclValue::Uuid { value, meta: _ } = &self {
            Some(*value)
        } else {
            None
        }
    }

    pub fn as_plane(&self) -> Option<&Plane> {
        if let KclValue::Plane(value) = &self {
            Some(value)
        } else {
            None
        }
    }

    pub fn as_solid(&self) -> Option<&Solid> {
        if let KclValue::Solid(value) = &self {
            Some(value)
        } else {
            None
        }
    }

    pub fn as_f64(&self) -> Option<f64> {
        if let KclValue::Number { value, meta: _ } = &self {
            Some(*value)
        } else if let KclValue::Int { value, meta: _ } = &self {
            Some(*value as f64)
        } else {
            None
        }
    }

    pub fn as_bool(&self) -> Option<bool> {
        if let KclValue::Bool { value, meta: _ } = &self {
            Some(*value)
        } else {
            None
        }
    }

    /// If this value fits in a u32, return it.
    pub fn get_u32(&self, source_ranges: Vec<SourceRange>) -> Result<u32, KclError> {
        let u = self.as_int().and_then(|n| u64::try_from(n).ok()).ok_or_else(|| {
            KclError::Semantic(KclErrorDetails {
                message: "Expected an integer >= 0".to_owned(),
                source_ranges: source_ranges.clone(),
            })
        })?;
        u32::try_from(u).map_err(|_| {
            KclError::Semantic(KclErrorDetails {
                message: "Number was too big".to_owned(),
                source_ranges,
            })
        })
    }

    /// If this value is of type function, return it.
    pub fn get_function(&self) -> Option<FnAsArg<'_>> {
        let KclValue::Function {
            func,
            expression,
            memory,
            meta: _,
        } = &self
        else {
            return None;
        };
        Some(FnAsArg {
            func: func.as_ref(),
            expr: expression.to_owned(),
            memory: memory.to_owned(),
        })
    }

    /// Get a tag identifier from a memory item.
    pub fn get_tag_identifier(&self) -> Result<TagIdentifier, KclError> {
        match self {
            KclValue::TagIdentifier(t) => Ok(*t.clone()),
            _ => Err(KclError::Semantic(KclErrorDetails {
                message: format!("Not a tag identifier: {:?}", self),
                source_ranges: self.clone().into(),
            })),
        }
    }

    /// Get a tag declarator from a memory item.
    pub fn get_tag_declarator(&self) -> Result<TagNode, KclError> {
        match self {
            KclValue::TagDeclarator(t) => Ok((**t).clone()),
            _ => Err(KclError::Semantic(KclErrorDetails {
                message: format!("Not a tag declarator: {:?}", self),
                source_ranges: self.clone().into(),
            })),
        }
    }

    /// Get an optional tag from a memory item.
    pub fn get_tag_declarator_opt(&self) -> Result<Option<TagNode>, KclError> {
        match self {
            KclValue::TagDeclarator(t) => Ok(Some((**t).clone())),
            _ => Err(KclError::Semantic(KclErrorDetails {
                message: format!("Not a tag declarator: {:?}", self),
                source_ranges: self.clone().into(),
            })),
        }
    }

    /// If this KCL value is a bool, retrieve it.
    pub fn get_bool(&self) -> Result<bool, KclError> {
        let Self::Bool { value: b, .. } = self else {
            return Err(KclError::Type(KclErrorDetails {
                source_ranges: self.into(),
                message: format!("Expected bool, found {}", self.human_friendly_type()),
            }));
        };
        Ok(*b)
    }

    /// If this memory item is a function, call it with the given arguments, return its val as Ok.
    /// If it's not a function, return Err.
    pub async fn call_fn(
        &self,
        args: Vec<Arg>,
        exec_state: &mut ExecState,
        ctx: ExecutorContext,
    ) -> Result<Option<KclValue>, KclError> {
        let KclValue::Function {
            func,
            expression,
            memory: closure_memory,
            meta,
        } = &self
        else {
            return Err(KclError::Semantic(KclErrorDetails {
                message: "not a in memory function".to_string(),
                source_ranges: vec![],
            }));
        };
        if let Some(func) = func {
            func(
                args,
                closure_memory.as_ref().clone(),
                expression.clone(),
                meta.clone(),
                exec_state,
                ctx,
            )
            .await
        } else {
            crate::executor::call_user_defined_function(
                args,
                closure_memory.as_ref(),
                expression.as_ref(),
                exec_state,
                &ctx,
            )
            .await
        }
    }
}