ui-automation 0.1.1

// Copyright (C) 2024 Tristan Gerritsen <tristan@thewoosh.org>
// All Rights Reserved.

use std::ptr::null;

use libc::pid_t;

use crate::{backend::macos::{core_foundation::{CFArray, CFArrayRef, CFString, CFStringRef, CFType, CFTypeRef, CFURL}, core_graphics::{CGPoint, CGSize}}, UIErrorKind, UIResult};

use super::{AXError, AXErrorI32, AXUIAttribute, AXValue, AXValueKind};

type AXUIElementRef = CFTypeRef;

#[derive(Debug)]
pub struct AXUIElement {
    inner: CFType,
}

impl AXUIElement {
    #[allow(unused)] // future reference, to get active window
    pub fn create_system_wide() -> UIResult<Self> {
        let ptr = unsafe { AXUIElementCreateSystemWide() };
        Self::try_from_ptr(ptr)
    }

    pub fn create_application(pid: u32) -> UIResult<Self> {
        let ptr = unsafe { AXUIElementCreateApplication(pid as _) };
        Self::try_from_ptr(ptr)
    }

    fn try_from_ptr(ptr: AXUIElementRef) -> UIResult<Self> {
        let Some(inner) = CFType::from_create(ptr) else {
            return Err(UIErrorKind::ProcessHasNoWindows.into());
        };

        let this = Self { inner };

        let role = this.attribute_string(AXUIAttribute::Role);
        if role.is_err() {
            // println!("Role Get Error: {role:?}, we are: {} with {:#?}", this.inner.description(), this.attribute_names());
            return Err(UIErrorKind::ProcessHasNoWindows.into());
        }

        Ok(this)
    }

    pub fn get_pid(&self) -> UIResult<pid_t> {
        let mut pid: pid_t = 0;
        let error = unsafe {
            AXUIElementGetPid(self.inner.ptr(), &mut pid)
        };

        if error != 0 {
            return Err(AXError::from(error).into());
        }

        Ok(pid)
    }

    pub fn window_count(&self) -> UIResult<usize> {
        let value = self.attribute(AXUIAttribute::Windows)?;
        let array = CFArray::from_ptr(value);
        Ok(array.len())
    }

    pub fn children<F: FnMut(Self)>(&self, mut f: F) -> UIResult<()> {
        let value = self.attribute(AXUIAttribute::Children)?;
        let array = CFArray::from_ptr(value);

        for child in array.iter::<CFType>() {
            f(AXUIElement {
                inner: child,
            })
        }

        Ok(())
    }

    #[allow(unused)] // great debugging/exploration feature
    pub fn attribute_names(&self) -> UIResult<Vec<String>> {
        let mut names: CFArrayRef = null();
        let error = unsafe {
            AXUIElementCopyAttributeNames(self.inner.ptr(), &mut names)
        };

        if error != 0 {
            return Err(AXError::from(error).into());
        }

        let names = CFArray::from_ptr(CFType::from_create(names).unwrap())
            .iter::<CFString>()
            .map(|x| x.to_string().unwrap())
            .collect();

        Ok(names)
    }

    fn attribute(&self, attribute: AXUIAttribute) -> UIResult<CFType> {
        let mut value: CFTypeRef = null();
        let error = unsafe {
            AXUIElementCopyAttributeValue(self.inner.ptr(), attribute.marshal(), &mut value)
        };

        if error != 0 {
            return Err(AXError::from(error).into());
        }

        CFType::from_get(value)
            .ok_or_else(|| {
                UIErrorKind::ConversionError { description: "failed to convert attribute value to CFType" }
                    .into()
            })
    }

    pub(crate) fn attribute_url(&self, attribute: AXUIAttribute) -> UIResult<String> {
        let attribute = self.attribute(attribute)?;

        if attribute.type_id() != "CFURL" {
            return Err(UIErrorKind::ConversionError { description: "attribute not convertible to URL" }.into());
        }

        let Some(value) = CFURL::from_get(attribute) else {
            return Err(UIErrorKind::ConversionError { description: "failed to convert AXUIElement attribute to CFURL" }.into());
        };

        let Ok(value) = value.to_string() else {
            return Err(UIErrorKind::ConversionError { description: "failed to convert AXUIElement attribute CFURL to CFString to Rust string" }.into());
        };

        Ok(value)
    }

    pub(crate) fn attribute_point(&self, attribute: AXUIAttribute) -> UIResult<CGPoint> {
        let value = self.attribute(attribute)?;
        let value = AXValue::from_get(value);
        match value.get() {
            AXValueKind::CGPoint(value) => Ok(value),
            actual => {
                eprintln!("Expected attribute {attribute:?} to be CGPoint, but was: {actual:?}");
                Err(UIErrorKind::ConversionError { description: "AXValue was different than expected" }.into())
            }
        }
    }

    pub(crate) fn attribute_size(&self, attribute: AXUIAttribute) -> UIResult<CGSize> {
        let value = self.attribute(attribute)?;
        let value = AXValue::from_get(value);
        match value.get() {
            AXValueKind::CGSize(value) => Ok(value),
            actual => {
                eprintln!("Expected attribute {attribute:?} to be CGSize, but was: {actual:?}");
                Err(UIErrorKind::ConversionError { description: "AXValue was different than expected" }.into())
            }
        }
    }

    pub(crate) fn attribute_string(&self, attribute: AXUIAttribute) -> UIResult<String> {
        let attribute = self.attribute(attribute)?;

        if attribute.type_id() != "CFString" {
            println!("Invalid type: {} desc={}", attribute.type_id(), attribute.description());
            return Err(UIErrorKind::ConversionError { description: "attribute not convertible to String" }.into());
        }

        let Some(value) = CFString::from_get(attribute) else {
            return Err(UIErrorKind::ConversionError { description: "failed to convert AXUIElement attribute to CFString" }.into());
        };

        let Ok(value) = value.to_string() else {
            return Err(UIErrorKind::ConversionError { description: "failed to convert AXUIElement attribute CFString to Rust string" }.into());
        };

        Ok(value)
    }
}

extern "C-unwind" {
    fn AXUIElementCreateApplication(pid: pid_t) -> AXUIElementRef;
    fn AXUIElementCreateSystemWide() -> AXUIElementRef;
    fn AXUIElementCopyAttributeNames(element: AXUIElementRef, names: *mut CFArrayRef) -> AXErrorI32;
    fn AXUIElementCopyAttributeValue(element: AXUIElementRef, attribute: CFStringRef, value: *mut CFTypeRef) -> AXErrorI32;
    fn AXUIElementGetPid(element: AXUIElementRef, pid: *mut pid_t) -> AXErrorI32;
}