unreal-sdk/src/types.rs

use anyhow::Context;
use itertools::Itertools;
use std::{hash::Hash, ptr::NonNull};

use crate::core_types::{self, FName};
use crate::tarray::TArray;

use self::traits::{AsUObject, AsUStruct, FromRaw};

#[repr(transparent)]
#[derive(Debug, Clone)]
pub struct UObject {
    inner: NonNull<core_types::UObject>,
}

unsafe impl Send for UObject {}
unsafe impl Sync for UObject {}

impl Eq for UObject {}

impl PartialEq for UObject {
    fn eq(&self, other: &Self) -> bool {
        unsafe {
            self.inner.as_ref().internal_index == other.inner.as_ref().internal_index
                && self.inner.as_ref().name == other.inner.as_ref().name
        }
    }
}

impl Hash for UObject {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.internal_index().hash(state);
        self.name().hash(state);
    }
}

#[repr(transparent)]
#[derive(Debug, Clone)]
pub struct UClass {
    pub(crate) inner: NonNull<core_types::UClass>,
}

impl Eq for UClass {}

impl PartialEq for UClass {
    fn eq(&self, other: &Self) -> bool {
        self.as_uobject().eq(&other.as_uobject())
    }
}

impl Hash for UClass {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.as_uobject().hash(state);
    }
}

#[repr(transparent)]
#[derive(Debug, Clone)]
pub struct UField {
    inner: NonNull<core_types::UField>,
}

impl Eq for UField {}

impl PartialEq for UField {
    fn eq(&self, other: &Self) -> bool {
        self.as_uobject().eq(&other.as_uobject())
    }
}

impl Hash for UField {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.as_uobject().hash(state);
    }
}

impl UObject {
    pub fn new(inner: NonNull<core_types::UObject>) -> Self {
        Self { inner }
    }

    pub fn maybe_new(inner: Option<NonNull<core_types::UObject>>) -> Option<Self> {
        inner.map(|inner| Self { inner })
    }

    pub fn to_inner(self) -> NonNull<core_types::UObject> {
        self.inner
    }

    pub const unsafe fn cast<T>(self) -> NonNull<T> {
        self.inner.cast::<T>()
    }

    /// returns a `Self` with the provided raw pointer, if the pointer is not null
    pub fn maybe_with_raw(raw: *mut core_types::UObject) -> Option<Self> {
        Self::maybe_new(NonNull::new(raw))
    }

    pub fn class(&self) -> Option<UClass> {
        unsafe { self.inner.as_ref().class.map(|inner| UClass::new(inner)) }
    }

    pub fn internal_index(&self) -> u32 {
        unsafe { self.inner.as_ref().internal_index }
    }

    pub fn name(&self) -> core_types::FName {
        unsafe { self.inner.as_ref().name }
    }

    pub fn get_name(&self) -> anyhow::Result<String> {
        let name = self.name();
        name.get_name()
            .map(|name_str| {
                if name.number > 0 {
                    format!("{}_{}", name_str, name.number)
                } else {
                    name_str
                }
            })
            .map(|name_str| {
                name_str
                    .rfind("/")
                    .map(|pos| name_str[(pos + 1)..].to_string())
                    .unwrap_or_else(|| name_str)
            })
    }

    pub fn get_name_or_default(&self) -> String {
        self.get_name()
            .unwrap_or_else(|_| "DefaultObjectName".to_string())
    }

    pub fn get_full_name(&self) -> anyhow::Result<String> {
        let tmp = self
            .iter_outer_objects()
            .map(|obj| obj.get_name())
            .fold_ok(String::new(), |acc, obj_name| {
                format!("{}.{}", obj_name, acc)
            })?;

        Ok(format!(
            "{} {}{}",
            self.class()
                .context("invalid class pointer")?
                .as_uobject()
                .get_name_or_default(),
            tmp,
            self.get_name_or_default()
        ))
    }

    pub fn get_full_name_or_default(&self) -> String {
        self.get_full_name()
            .unwrap_or_else(|_| "DefaultObjectFullName".to_string())
    }

    pub fn outer(&self) -> Option<UObject> {
        unsafe { self.inner.as_ref().outer.map(|inner| UObject::new(inner)) }
    }

    /// returns the package object of this object
    pub fn outermost(&self) -> Option<UObject> {
        self.iter_outer_objects().last()
    }

    pub fn is_package_object(&self) -> bool {
        unsafe { self.inner.as_ref().outer.is_none() }
    }

    pub fn iter_outer_objects(&self) -> OuterObjectIterator {
        OuterObjectIterator::new(self.clone())
    }

    pub fn is_a(&self, other: &UClass) -> bool {
        self.class()
            .map(|class| class.iter_super_classes().contains(other))
            .unwrap_or(false)
    }
}

pub struct OuterObjectIterator {
    object: UObject,
}

impl OuterObjectIterator {
    pub fn new(object: UObject) -> Self {
        Self { object }
    }
}

impl Iterator for OuterObjectIterator {
    type Item = UObject;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(outer) = self.object.outer() {
            self.object = outer.clone();

            Some(outer)
        } else {
            None
        }
    }
}

impl From<UObject> for UClass {
    fn from(obj: UObject) -> Self {
        Self::new(obj.inner.cast())
    }
}

impl AsUStruct for UClass {
    fn as_ustruct(&self) -> UStruct {
        UStruct::new(self.inner.clone().cast())
    }
}

impl AsRef<core_types::UClass> for UClass {
    fn as_ref(&self) -> &core_types::UClass {
        unsafe { self.inner.as_ref() }
    }
}

impl FromRaw<core_types::UClass> for UClass {
    fn from_non_null(inner: NonNull<core_types::UClass>) -> Self {
        Self::new(inner)
    }
}

impl UClass {
    pub fn new(inner: NonNull<core_types::UClass>) -> Self {
        Self { inner }
    }
    pub fn iter_super_classes(&self) -> SuperClassIter {
        SuperClassIter::new(self.clone())
    }
}

pub struct SuperClassIter {
    class: UClass,
}

impl SuperClassIter {
    pub fn new(class: UClass) -> Self {
        Self { class }
    }
}

impl Iterator for SuperClassIter {
    type Item = UClass;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(next) = self.class.super_struct().map(|spr| spr.as_uobject().into()) {
            let item = core::mem::replace(&mut self.class, next);
            Some(item)
        } else {
            None
        }
    }
}

impl From<UObject> for UField {
    fn from(obj: UObject) -> Self {
        Self::new(obj.inner.cast())
    }
}

impl UField {
    pub fn new(inner: NonNull<core_types::UField>) -> Self {
        Self { inner }
    }

    pub fn maybe_new(inner: Option<NonNull<core_types::UField>>) -> Option<Self> {
        inner.map(|inner| Self { inner })
    }

    pub fn maybe_with_raw(raw: *mut core_types::UField) -> Option<Self> {
        NonNull::new(raw).map(|inner| Self::new(inner))
    }

    pub fn as_uobject(&self) -> UObject {
        UObject::new(self.inner.clone().cast())
    }

    pub unsafe fn as_ref(&self) -> &core_types::UField {
        self.inner.as_ref()
    }

    pub fn next(&self) -> Option<Self> {
        Self::maybe_new(unsafe { self.as_ref().next })
    }
}

#[derive(Debug, Clone)]
pub struct UEnum {
    inner: NonNull<core_types::UEnum>,
}

impl From<UObject> for UEnum {
    fn from(obj: UObject) -> Self {
        Self::new(obj.inner.cast())
    }
}

impl UEnum {
    pub fn new(inner: NonNull<core_types::UEnum>) -> Self {
        Self { inner }
    }

    pub fn maybe_new(inner: Option<NonNull<core_types::UEnum>>) -> Option<Self> {
        inner.map(|inner| Self { inner })
    }

    pub fn maybe_with_raw(raw: *mut core_types::UEnum) -> Option<Self> {
        NonNull::new(raw).map(|inner| Self::new(inner))
    }

    pub fn as_uobject(&self) -> UObject {
        UObject::new(self.inner.clone().cast())
    }

    pub unsafe fn as_ref(&self) -> &core_types::UEnum {
        self.inner.as_ref()
    }

    pub fn get_names(&self) -> &TArray<FName> {
        unsafe { &self.as_ref().names }
    }
}

#[derive(Debug, Clone)]
pub struct UStruct {
    inner: NonNull<core_types::UStruct>,
}

impl From<UObject> for UStruct {
    fn from(obj: UObject) -> Self {
        Self::new(obj.inner.cast())
    }
}

impl FromRaw<core_types::UStruct> for UStruct {
    fn from_non_null(inner: NonNull<core_types::UStruct>) -> Self {
        Self::new(inner)
    }
}

impl UStruct {
    pub fn new(inner: NonNull<core_types::UStruct>) -> Self {
        Self { inner }
    }

    pub unsafe fn as_raw(&self) -> &core_types::UStruct {
        self.inner.as_ref()
    }

    fn super_struct(&self) -> Option<Self> {
        Self::from_maybe_non_null(unsafe { self.as_raw().super_field })
    }

    fn children(&self) -> Option<UField> {
        UField::maybe_new(unsafe { self.as_raw().children })
    }

    fn iter_fields(&self) -> Option<UStructFieldIterator> {
        self.children()
            .map(|field| UStructFieldIterator::new(field))
    }
}

impl traits::AsUStruct for UStruct {
    fn as_ustruct(&self) -> UStruct {
        self.clone()
    }
}

pub struct UStructFieldIterator {
    field: UField,
}

impl UStructFieldIterator {
    pub fn new(field: UField) -> Self {
        Self { field }
    }
}

impl Iterator for UStructFieldIterator {
    type Item = UField;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(mut next) = self.field.next() {
            std::mem::swap(&mut self.field, &mut next);

            Some(next)
        } else {
            None
        }
    }
}

#[derive(Debug, Clone)]
pub struct UScriptStruct {
    inner: NonNull<core_types::UStruct>,
}

impl From<UObject> for UScriptStruct {
    fn from(obj: UObject) -> Self {
        Self::new(obj.inner.cast())
    }
}

impl UScriptStruct {
    pub fn new(inner: NonNull<core_types::UStruct>) -> Self {
        Self { inner }
    }

    pub fn maybe_new(inner: Option<NonNull<core_types::UStruct>>) -> Option<Self> {
        inner.map(|inner| Self { inner })
    }

    pub fn maybe_with_raw(raw: *mut core_types::UStruct) -> Option<Self> {
        NonNull::new(raw).map(|inner| Self::new(inner))
    }

    pub fn as_uobject(&self) -> UObject {
        UObject::new(self.inner.clone().cast())
    }

    pub unsafe fn as_ref(&self) -> &core_types::UStruct {
        self.inner.as_ref()
    }

    pub fn iter_fields(&self) -> Option<UStructFieldIterator> {
        UField::maybe_new(unsafe { self.as_ref().children })
            .map(|field| UStructFieldIterator::new(field))
    }
}

#[derive(Debug, Clone)]
pub struct UProperty {
    inner: NonNull<core_types::UProperty>,
}

impl From<UObject> for UProperty {
    /// WARNING: There is no type check on this conversion, so it is most certainly unsafe
    fn from(obj: UObject) -> Self {
        Self::new(obj.inner.cast())
    }
}

impl UProperty {
    pub fn new(inner: NonNull<core_types::UProperty>) -> Self {
        Self { inner }
    }

    pub fn maybe_new(inner: Option<NonNull<core_types::UProperty>>) -> Option<Self> {
        inner.map(|inner| Self { inner })
    }

    pub fn maybe_with_raw(raw: *mut core_types::UProperty) -> Option<Self> {
        NonNull::new(raw).map(|inner| Self::new(inner))
    }

    pub fn as_uobject(&self) -> UObject {
        UObject::new(self.inner.clone().cast())
    }

    pub unsafe fn as_ref(&self) -> &core_types::UProperty {
        self.inner.as_ref()
    }
}

#[derive(Debug, Clone)]
pub struct UFunction {
    inner: NonNull<core_types::UFunction>,
}

impl From<UObject> for UFunction {
    /// WARNING: There is no type check on this conversion, so it is most certainly unsafe
    fn from(obj: UObject) -> Self {
        Self::new(obj.inner.cast())
    }
}

impl UFunction {
    pub fn new(inner: NonNull<core_types::UFunction>) -> Self {
        Self { inner }
    }

    pub fn maybe_new(inner: Option<NonNull<core_types::UFunction>>) -> Option<Self> {
        inner.map(|inner| Self { inner })
    }

    pub fn maybe_with_raw(raw: *mut core_types::UFunction) -> Option<Self> {
        NonNull::new(raw).map(|inner| Self::new(inner))
    }

    pub fn as_uobject(&self) -> UObject {
        UObject::new(self.inner.clone().cast())
    }

    pub unsafe fn as_ref(&self) -> &core_types::UFunction {
        self.inner.as_ref()
    }
}

#[derive(Debug, Clone)]
pub enum UAnyType {
    UObject(UObject),
    UClass(UClass),
    UField(UField),
    UScriptStruct(UScriptStruct),
    UProperty(UProperty),
    UEnum(UEnum),
    UStruct(UStruct),
    UFunction(UFunction),
}

impl UAnyType {
    pub fn as_uobject(&self) -> UObject {
        match self {
            UAnyType::UObject(rep) => rep.clone(),
            UAnyType::UClass(rep) => rep.as_uobject(),
            UAnyType::UField(rep) => rep.as_uobject(),
            UAnyType::UScriptStruct(rep) => rep.as_uobject(),
            UAnyType::UProperty(rep) => rep.as_uobject(),
            UAnyType::UEnum(rep) => rep.as_uobject(),
            UAnyType::UStruct(rep) => rep.as_uobject(),
            UAnyType::UFunction(rep) => rep.as_uobject(),
        }
    }
}

impl Hash for UAnyType {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.as_uobject().hash(state);
    }
}

pub mod traits {
    use crate::core_types::FName;
    use std::ptr::NonNull;

    use itertools::Itertools;

    use super::{OuterObjectIterator, UClass, UField, UObject, UStruct, UStructFieldIterator};

    pub unsafe fn null_ptr_cast<'a, T, U>(ptr: &'a NonNull<T>) -> &'a NonNull<U> {
        std::mem::transmute(ptr)
    }

    pub trait AsUObject: Sized {
        fn as_uobject(&self) -> UObject;

        fn class(&self) -> Option<UClass> {
            self.as_uobject().class()
        }

        fn internal_index(&self) -> u32 {
            self.as_uobject().internal_index()
        }

        fn name(&self) -> FName {
            self.as_uobject().name()
        }

        fn get_name(&self) -> anyhow::Result<String> {
            self.as_uobject().get_name()
        }

        fn get_name_or_default(&self) -> String {
            self.as_uobject().get_name_or_default()
        }

        fn get_full_name(&self) -> anyhow::Result<String> {
            self.as_uobject().get_full_name()
        }

        fn get_full_name_or_default(&self) -> String {
            self.as_uobject().get_full_name_or_default()
        }

        fn outer(&self) -> Option<UObject> {
            self.as_uobject().outer()
        }

        fn outermost(&self) -> Option<UObject> {
            self.as_uobject().outermost()
        }

        fn is_package_object(&self) -> bool {
            self.as_uobject().is_package_object()
        }

        fn iter_outer_objects(&self) -> OuterObjectIterator {
            OuterObjectIterator::new(self.as_uobject())
        }

        fn is_a(&self, other: &UClass) -> bool {
            self.class()
                .map(|class| class.iter_super_classes().contains(other))
                .unwrap_or(false)
        }
    }

    pub trait AsUStruct {
        fn as_ustruct(&self) -> UStruct;

        #[inline]
        fn children(&self) -> Option<UField> {
            self.as_ustruct().children()
        }

        #[inline]
        fn super_struct(&self) -> Option<UStruct> {
            self.as_ustruct().super_struct()
        }

        #[inline]
        fn iter_fields(&self) -> Option<UStructFieldIterator> {
            self.children()
                .map(|field| UStructFieldIterator::new(field))
        }
    }

    impl<T> AsUObject for T
    where
        T: AsUStruct,
    {
        fn as_uobject(&self) -> UObject {
            UObject::new(self.as_ustruct().inner.cast())
        }
    }

    pub trait FromRaw<R> {
        fn from_non_null(inner: NonNull<R>) -> Self;

        fn from_raw(raw: *mut R) -> Option<Self>
        where
            Self: Sized,
        {
            NonNull::new(raw).map(|inner| Self::from_non_null(inner))
        }

        fn from_maybe_non_null(inner: Option<NonNull<R>>) -> Option<Self>
        where
            Self: Sized,
        {
            inner.map(|inner| Self::from_non_null(inner))
        }
    }
}