strafe-client-jed/src/runner/instance.rs

use mlua::IntoLua;
use rbx_types::Ref;
use rbx_dom_weak::{InstanceBuilder,WeakDom};

use super::vector3::Vector3;

// LMAO look at this function!
fn dom<T>(lua:&mlua::Lua,mut f:impl FnMut(&mut WeakDom)->mlua::Result<T>)->mlua::Result<T>{
	let mut dom=lua.app_data_mut::<&'static mut WeakDom>().ok_or(mlua::Error::runtime("DataModel missing"))?;
	f(&mut *dom)
}

fn coerce_float32(value:&mlua::Value)->Option<f32>{
	match value{
		&mlua::Value::Integer(i)=>Some(i as f32),
		&mlua::Value::Number(f)=>Some(f as f32),
		_=>None,
	}
}

fn get_full_name(dom:&rbx_dom_weak::WeakDom,instance:&rbx_dom_weak::Instance)->String{
	let mut full_name=instance.name.clone();
	let mut pref=instance.parent();
	while let Some(parent)=dom.get_by_ref(pref){
		full_name.insert(0,'.');
		full_name.insert_str(0,parent.name.as_str());
		pref=parent.parent();
	}
	full_name
}

pub fn find_first_child_of_class<'a>(dom:&'a rbx_dom_weak::WeakDom,instance:&rbx_dom_weak::Instance,class:&str)->Option<&'a rbx_dom_weak::Instance>{
	instance.children().iter().filter_map(|&r|dom.get_by_ref(r)).find(|inst|inst.class==class)
}

//workaround until I have an enum of classes
struct Dereferent(Ref);
impl mlua::UserData for Dereferent{}
type_from_lua_userdata!(Dereferent);
impl Referent for Dereferent{
	fn referent(&self)->Ref{
		self.0
	}
}

trait Referent{
	fn referent(&self)->Ref;
	fn get<'a>(&self,dom:&'a WeakDom)->mlua::Result<&'a rbx_dom_weak::Instance>{
		dom.get_by_ref(self.referent()).ok_or(mlua::Error::runtime("Instance missing"))
	}
	fn get_mut<'a>(&self,dom:&'a mut WeakDom)->mlua::Result<&'a mut rbx_dom_weak::Instance>{
		dom.get_by_ref_mut(self.referent()).ok_or(mlua::Error::runtime("Instance missing"))
	}
}

macro_rules! class{
	($class:ident)=>{
		pub struct $class{
			referent:Ref,
		}
		impl $class{
			pub const fn new(referent:Ref)->Self{
				Self{referent}
			}
		}
		impl Referent for $class{
			fn referent(&self)->Ref{
				self.referent
			}
		}
		type_from_lua_userdata!($class);
	};
}
macro_rules! class_composition{
	($class:ident,($($superclass:ident),*))=>{
		impl mlua::UserData for $class{
			fn add_fields<'lua,F:mlua::UserDataFields<'lua,Self>>(fields:&mut F){
				fields.add_field_method_get("Referent",|_,this|{
					Ok(Dereferent(this.referent()))
				});
				$(
					$superclass::composition_add_fields(fields);
				)*
			}
			fn add_methods<'lua,M:mlua::UserDataMethods<'lua,Self>>(methods:&mut M){
				$(
					$superclass::composition_add_methods(methods);
				)*
			}
		}
	};
}

//TODO: update rbx_reflection and use dom.superclasses_iter
pub struct SuperClassIter<'a> {
    database: &'a rbx_reflection::ReflectionDatabase<'a>,
    descriptor: Option<&'a rbx_reflection::ClassDescriptor<'a>>,
}
impl<'a> SuperClassIter<'a> {
    fn next_descriptor(&self) -> Option<&'a rbx_reflection::ClassDescriptor<'a>> {
        let superclass = self.descriptor?.superclass.as_ref()?;
        self.database.classes.get(superclass)
    }
}
impl<'a> Iterator for SuperClassIter<'a> {
    type Item = &'a rbx_reflection::ClassDescriptor<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        let next_descriptor = self.next_descriptor();
        std::mem::replace(&mut self.descriptor, next_descriptor)
    }
}

class!(Instance);
class_composition!(Instance,(Instance));

impl Instance{
	fn composition_add_fields<'lua,T:Referent,F:mlua::UserDataFields<'lua,T>>(fields:&mut F){
		fields.add_field_method_get("Parent",|lua,this|{
			dom(lua,|dom|{
				let instance=this.get(dom)?;
				Ok(Instance::new(instance.parent()))
			})
		});
		fields.add_field_method_set("Parent",|lua,this,val:mlua::AnyUserData|{
			let Dereferent(referent)=mlua::AnyUserDataExt::get(&val,"Referent")?;
			dom(lua,|dom|{
				dom.transfer_within(this.referent(),referent);
				Ok(())
			})
		});
		fields.add_field_method_get("Name",|lua,this|{
			dom(lua,|dom|{
				let instance=this.get(dom)?;
				Ok(instance.name.clone())
			})
		});
		fields.add_field_method_set("Name",|lua,this,val:mlua::String|{
			dom(lua,|dom|{
				let instance=this.get_mut(dom)?;
				//Why does this need to be cloned?
				instance.name=val.to_str()?.to_owned();
				Ok(())
			})
		});
		fields.add_field_method_get("ClassName",|lua,this|{
			dom(lua,|dom|{
				let instance=this.get(dom)?;
				Ok(instance.class.clone())
			})
		});
	}
	fn composition_add_methods<'lua,T:Referent,M:mlua::UserDataMethods<'lua,T>>(methods:&mut M){
		methods.add_method("GetChildren",|lua,this,_:()|
			dom(lua,|dom|{
				let instance=this.get(dom)?;
				let children:Vec<_>=instance
					.children()
					.iter()
					.copied()
					.map(Instance::new)
					.collect();
				Ok(children)
			})
		);
		let ffc=|lua,this:&T,(name,search_descendants):(mlua::String,mlua::Value)|{
			let name_str=name.to_str()?;
			let search_descendants=match search_descendants{
				mlua::Value::Nil=>false,
				mlua::Value::Boolean(b)=>b,
				_=>Err(mlua::Error::runtime("Invalid argument #3 bool expected"))?,
			};
			dom(lua,|dom|{
				let instance=this.get(dom)?;
				let child=match search_descendants{
					true=>dom.descendants_of(this.referent()).find(|inst|inst.name==name_str),
					false=>instance.children().iter().filter_map(|&r|
						dom.get_by_ref(r)
					).find(|inst|inst.name==name_str),
				};
				match child{
					Some(instance)=>Instance::new(instance.referent()).into_lua(lua),
					None=>mlua::Value::Nil.into_lua(lua),
				}
			})
		};
		methods.add_method("FindFirstChild",ffc);
		methods.add_method("WaitForChild",ffc);
		methods.add_method("FindFirstChildOfClass",|lua,this,(class,search_descendants):(mlua::String,mlua::Value)|{
			let class_str=class.to_str()?;
			let search_descendants=match search_descendants{
				mlua::Value::Nil=>false,
				mlua::Value::Boolean(b)=>b,
				_=>Err(mlua::Error::runtime("Invalid argument #3 bool expected"))?,
			};
			if search_descendants==true{
				return Err(mlua::Error::runtime("FFC of class searching descendants not supported get rekt"));
			}
			dom(lua,|dom|{
				Ok(find_first_child_of_class(dom,this.get(dom)?,class_str)
        			.map(|inst|(Instance::new(inst.referent())))
				)
			})
		});
		methods.add_method("GetDescendants",|lua,this,_:()|
			dom(lua,|dom|{
				let children:Vec<_>=dom
					.descendants_of(this.referent())
					.map(|instance|
						Instance::new(instance.referent())
					)
					.collect();
				Ok(children)
			})
		);
		methods.add_method("IsA",|lua,this,classname:mlua::String|
			dom(lua,|dom|{
				let instance=this.get(dom)?;
				Ok(crate::context::class_is_a(instance.class.as_str(),classname.to_str()?))
			})
		);
		methods.add_meta_function(mlua::MetaMethod::Index,|lua,(this,index):(mlua::AnyUserData,mlua::String)|{
			let index_str=index.to_str()?;
			let dereferent:Dereferent=mlua::AnyUserDataExt::get(&this,"Referent")?;
			dom(lua,|dom|{
				let instance=dereferent.get(dom)?;
				//find a child with a matching name
				let maybe_child=instance.children()
				.iter()
				.find(|&&r|
					dom.get_by_ref(r)
					.is_some_and(|instance|instance.name==index_str)
				);
				match maybe_child{
					Some(&referent)=>Instance::new(referent).into_lua(lua),
					None=>mlua::Value::Nil.into_lua(lua),
				}
			})
		});
		methods.add_meta_function(mlua::MetaMethod::NewIndex,|lua,(this,index,value):(mlua::AnyUserData,mlua::String,mlua::Value)|{
			let dereferent:Dereferent=mlua::AnyUserDataExt::get(&this,"Referent")?;
			dom(lua,|dom|{
				let instance=dereferent.get_mut(dom)?;
				//println!("__newindex t={} i={index:?} v={value:?}",instance.name);
				let index_str=index.to_str()?;
				let db=rbx_reflection_database::get();
				let class=db.classes.get(instance.class.as_str()).ok_or(mlua::Error::runtime("Class missing"))?;
				let mut iter=SuperClassIter{
					database:db,
					descriptor:Some(class),
				};
				let property=iter.find_map(|cls|cls.properties.get(index_str)).ok_or(mlua::Error::runtime(format!("Property '{index_str}' missing on class '{}'",class.name)))?;
				match &property.data_type{
					rbx_reflection::DataType::Value(rbx_types::VariantType::Vector3)=>{
						let typed_value:Vector3=value.as_userdata().ok_or(mlua::Error::runtime("Expected Userdata"))?.take()?;
						instance.properties.insert(index_str.to_owned(),rbx_types::Variant::Vector3(typed_value.into()));
					},
					rbx_reflection::DataType::Value(rbx_types::VariantType::Float32)=>{
						let typed_value:f32=coerce_float32(&value).ok_or(mlua::Error::runtime("Expected f32"))?;
						instance.properties.insert(index_str.to_owned(),rbx_types::Variant::Float32(typed_value));
					},
					rbx_reflection::DataType::Enum(enum_name)=>{
						let typed_value=match &value{
							&mlua::Value::Integer(int)=>Ok(rbx_types::Enum::from_u32(int as u32)),
							&mlua::Value::Number(num)=>Ok(rbx_types::Enum::from_u32(num as u32)),
							mlua::Value::String(s)=>{
								let e=db.enums.get(enum_name).ok_or(mlua::Error::runtime("Database DataType Enum name  does not exist"))?;
								Ok(rbx_types::Enum::from_u32(*e.items.get(s.to_str()?).ok_or(mlua::Error::runtime("Invalid enum item"))?))
							},
							mlua::Value::UserData(any_user_data)=>{
								let e:super::r#enum::Enum=any_user_data.take()?;
								Ok(e.into())
							},
							_=>Err(mlua::Error::runtime("Expected Enum")),
						}?;
						instance.properties.insert(index_str.to_owned(),rbx_types::Variant::Enum(typed_value));
					},
					rbx_reflection::DataType::Value(rbx_types::VariantType::Color3)=>{
						let typed_value:super::color3::Color3=value.as_userdata().ok_or(mlua::Error::runtime("Expected Color3"))?.take()?;
						instance.properties.insert(index_str.to_owned(),rbx_types::Variant::Color3(typed_value.into()));
					},
					rbx_reflection::DataType::Value(rbx_types::VariantType::Bool)=>{
						let typed_value=value.as_boolean().ok_or(mlua::Error::runtime("Expected boolean"))?;
						instance.properties.insert(index_str.to_owned(),rbx_types::Variant::Bool(typed_value));
					},
					rbx_reflection::DataType::Value(rbx_types::VariantType::String)=>{
						let typed_value=value.as_str().ok_or(mlua::Error::runtime("Expected boolean"))?;
						instance.properties.insert(index_str.to_owned(),rbx_types::Variant::String(typed_value.to_owned()));
					},
					other=>return Err(mlua::Error::runtime(format!("Unimplemented property type: {other:?}"))),
				}
				Ok(())
			})
		});
	}
}

class!(DataModel);
class_composition!(DataModel,(Instance,DataModel));
impl DataModel{
	fn composition_add_fields<'lua,T:Referent,F:mlua::UserDataFields<'lua,T>>(fields:&mut F){
		fields.add_field_method_get("workspace",|lua,this|{
			dom(lua,|dom|{
				Ok(find_first_child_of_class(dom,this.get(dom)?,"Workspace")
					.map(|inst|Workspace::new(inst.referent()))
				)
			})
		});
		fields.add_field_method_get("PlaceId",|lua,this|{
			Ok(mlua::Value::Integer(0))
		});
	}
	fn composition_add_methods<'lua,T,M:mlua::UserDataMethods<'lua,T>>(methods:&mut M){
		methods.add_method("GetService",|lua,this,service:String|
			dom(lua,|dom|{
				match service.as_str(){
					"Lighting"=>{
						let referent=dom.root()
							.children()
							.iter()
							.find(|&&c|
								dom.get_by_ref(c).is_some_and(|c|c.class=="Lighting")
							).map(|r|*r)
							.unwrap_or_else(||
								dom.insert(dom.root_ref(),InstanceBuilder::new("Lighting"))
							);
						Lighting::new(referent).into_lua(lua)
					},
					other=>Err::<mlua::Value,_>(mlua::Error::runtime(format!("Service '{other}' not supported"))),
				}
			})
		);
	}
}

class!(Workspace);
class_composition!(Workspace,(Instance));

class!(Lighting);
class_composition!(Lighting,(Instance));

class!(Script);
class_composition!(Script,(Instance));
impl Script{
	pub fn get_name_source(&self,lua:&mlua::Lua)->Result<(String,String),mlua::Error>{
		dom(lua,|dom|{
			let instance=self.get(dom)?;
			let source=match instance.properties.get("Source"){
				Some(rbx_dom_weak::types::Variant::String(s))=>s.clone(),
				_=>Err(mlua::Error::external("Missing script.Source"))?,
			};
			Ok((get_full_name(dom,instance),source))
		})
	}
}

class!(Terrain);
class_composition!(Terrain,(Instance,Terrain));
impl Terrain{
	fn composition_add_fields<'lua,T:Referent,F:mlua::UserDataFields<'lua,T>>(fields:&mut F){
	}
	fn composition_add_methods<'lua,T,M:mlua::UserDataMethods<'lua,T>>(methods:&mut M){
		methods.add_method("FillBlock",|lua,this,_:(super::cframe::CFrame,Vector3,super::r#enum::Enum)|
			Ok(())//Ok(mlua::Value::Nil)
		)
	}
}