mtl-rs 0.2.1

use objc2::{Message, extern_protocol, msg_send, rc::Retained, runtime::ProtocolObject};
use objc2_foundation::{NSError, NSObjectProtocol, NSString};

use crate::{
    MTLAccelerationStructureCommandEncoder, MTLAccelerationStructurePassDescriptor, MTLBlitCommandEncoder,
    MTLBlitPassDescriptor, MTLCommandBufferErrorOption, MTLCommandBufferHandler, MTLCommandBufferStatus,
    MTLCommandQueue, MTLComputeCommandEncoder, MTLComputePassDescriptor, MTLDevice, MTLDispatchType, MTLDrawable,
    MTLEvent, MTLLogContainer, MTLParallelRenderCommandEncoder, MTLRenderCommandEncoder, MTLRenderPassDescriptor,
    MTLResidencySet, MTLResourceStateCommandEncoder, MTLResourceStatePassDescriptor, MetalError,
    util::ref_slice_as_ptr,
};

extern_protocol!(
    /// A serial list of commands for the device to execute.
    ///
    /// Availability: macOS 10.11+, iOS 8.0+
    ///
    /// Thread safety: command buffers are thread-safe. Commit, completion,
    /// status, and wait operations may be used from different threads; command
    /// encoders remain single-threaded.
    ///
    /// # Safety
    ///
    /// Implementors must be Objective-C objects that conform to the
    /// `MTLCommandBuffer` protocol and uphold its thread-safety contract.
    pub unsafe trait MTLCommandBuffer: NSObjectProtocol + Send + Sync {
        /// The device this resource was created against.
        #[unsafe(method(device))]
        #[unsafe(method_family = none)]
        fn device(&self) -> Retained<ProtocolObject<dyn MTLDevice>>;

        /// The command queue this command buffer was created from.
        #[unsafe(method(commandQueue))]
        #[unsafe(method_family = none)]
        fn command_queue(&self) -> Retained<ProtocolObject<dyn MTLCommandQueue>>;

        /// If YES, this command buffer holds strong references to objects needed to execute this command buffer.
        #[unsafe(method(retainedReferences))]
        #[unsafe(method_family = none)]
        fn retained_references(&self) -> bool;

        /// The options that configure command-buffer error reporting.
        ///
        /// Availability: macOS 11.0+, iOS 14.0+
        #[unsafe(method(errorOptions))]
        #[unsafe(method_family = none)]
        fn error_options(&self) -> MTLCommandBufferErrorOption;

        /// Append this command buffer to the end of its MTLCommandQueue.
        #[unsafe(method(enqueue))]
        #[unsafe(method_family = none)]
        fn enqueue(&self);

        /// Commit a command buffer so it can be executed as soon as possible.
        #[unsafe(method(commit))]
        #[unsafe(method_family = none)]
        fn commit(&self);

        /// Synchronously wait for this command buffer to be scheduled.
        #[unsafe(method(waitUntilScheduled))]
        #[unsafe(method_family = none)]
        fn wait_until_scheduled(&self);

        /// Synchronously wait for this command buffer to complete.
        #[unsafe(method(waitUntilCompleted))]
        #[unsafe(method_family = none)]
        fn wait_until_completed(&self);

        /// Status reports the current stage in the lifetime of MTLCommandBuffer,
        /// as it proceeds to enqueued, committed, scheduled, and completed.
        #[unsafe(method(status))]
        #[unsafe(method_family = none)]
        fn status(&self) -> MTLCommandBufferStatus;

        /// Logs generated by the command buffer during execution of the GPU commands.
        /// Valid after GPU execution is completed.
        ///
        /// Availability: macOS 11.0+, iOS 14.0+
        #[unsafe(method(logs))]
        #[unsafe(method_family = none)]
        fn logs(&self) -> Retained<ProtocolObject<dyn MTLLogContainer>>;

        /// Add a drawable present that will be invoked when this command buffer has been scheduled for execution.
        #[unsafe(method(presentDrawable:))]
        #[unsafe(method_family = none)]
        fn present_drawable(
            &self,
            drawable: &ProtocolObject<dyn MTLDrawable>,
        );

        /// Presents a drawable at a specific host time.
        #[unsafe(method(presentDrawable:atTime:))]
        #[unsafe(method_family = none)]
        fn present_drawable_at_time(
            &self,
            drawable: &ProtocolObject<dyn MTLDrawable>,
            presentation_time: f64,
        );

        /// Presents a drawable after the previous frame has remained visible for at least `duration` seconds.
        ///
        /// Availability: macOS 10.15.4+, iOS 10.3+
        #[unsafe(method(presentDrawable:afterMinimumDuration:))]
        #[unsafe(method_family = none)]
        fn present_drawable_after_minimum_duration(
            &self,
            drawable: &ProtocolObject<dyn MTLDrawable>,
            duration: f64,
        );

        /// Encodes a command that pauses execution of this command buffer until the specified event reaches a given value.
        ///
        /// This method may only be called if there is no current command encoder on the receiver.
        ///
        /// Availability: macOS 10.14+, iOS 12.0+
        #[unsafe(method(encodeWaitForEvent:value:))]
        #[unsafe(method_family = none)]
        fn encode_wait_for_event_value(
            &self,
            event: &ProtocolObject<dyn MTLEvent>,
            value: u64,
        );

        /// Encodes a command that signals an event with a given value.
        ///
        /// This method may only be called if there is no current command encoder on the receiver.
        ///
        /// Availability: macOS 10.14+, iOS 12.0+
        #[unsafe(method(encodeSignalEvent:value:))]
        #[unsafe(method_family = none)]
        fn encode_signal_event_value(
            &self,
            event: &ProtocolObject<dyn MTLEvent>,
            value: u64,
        );

        /// Returns a render command encoder to encode into this command buffer.
        ///
        /// Availability: macOS 10.11+, iOS 8.0+
        #[unsafe(method(renderCommandEncoderWithDescriptor:))]
        #[unsafe(method_family = none)]
        fn render_command_encoder_with_descriptor(
            &self,
            render_pass_descriptor: &MTLRenderPassDescriptor,
        ) -> Option<Retained<ProtocolObject<dyn MTLRenderCommandEncoder>>>;

        /// Returns a parallel render command encoder for the render pass.
        #[unsafe(method(parallelRenderCommandEncoderWithDescriptor:))]
        #[unsafe(method_family = none)]
        fn parallel_render_command_encoder_with_descriptor(
            &self,
            render_pass_descriptor: &MTLRenderPassDescriptor,
        ) -> Option<Retained<ProtocolObject<dyn MTLParallelRenderCommandEncoder>>>;

        /// Returns a compute command encoder to encode into this command buffer.
        #[unsafe(method(computeCommandEncoder))]
        #[unsafe(method_family = none)]
        fn compute_command_encoder(&self) -> Option<Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>>;

        #[unsafe(method(computeCommandEncoderWithDispatchType:))]
        #[unsafe(method_family = none)]
        fn compute_command_encoder_with_dispatch_type(
            &self,
            dispatch_type: MTLDispatchType,
        ) -> Option<Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>>;

        /// Returns a blit command encoder to encode into this command buffer.
        #[unsafe(method(blitCommandEncoder))]
        #[unsafe(method_family = none)]
        fn blit_command_encoder(&self) -> Option<Retained<ProtocolObject<dyn MTLBlitCommandEncoder>>>;

        /// Returns a compute command encoder configured by the given descriptor.
        #[unsafe(method(computeCommandEncoderWithDescriptor:))]
        #[unsafe(method_family = none)]
        fn compute_command_encoder_with_descriptor(
            &self,
            compute_pass_descriptor: &MTLComputePassDescriptor,
        ) -> Option<Retained<ProtocolObject<dyn MTLComputeCommandEncoder>>>;

        /// Returns a blit command encoder configured by the given descriptor.
        #[unsafe(method(blitCommandEncoderWithDescriptor:))]
        #[unsafe(method_family = none)]
        fn blit_command_encoder_with_descriptor(
            &self,
            blit_pass_descriptor: &MTLBlitPassDescriptor,
        ) -> Option<Retained<ProtocolObject<dyn MTLBlitCommandEncoder>>>;

        /// Returns a resource state command encoder to encode into this command buffer.
        #[unsafe(method(resourceStateCommandEncoder))]
        #[unsafe(method_family = none)]
        fn resource_state_command_encoder(
            &self
        ) -> Option<Retained<ProtocolObject<dyn MTLResourceStateCommandEncoder>>>;

        /// Returns a resource state command encoder configured by the given descriptor.
        #[unsafe(method(resourceStateCommandEncoderWithDescriptor:))]
        #[unsafe(method_family = none)]
        fn resource_state_command_encoder_with_descriptor(
            &self,
            resource_state_pass_descriptor: &MTLResourceStatePassDescriptor,
        ) -> Option<Retained<ProtocolObject<dyn MTLResourceStateCommandEncoder>>>;

        /// Returns an acceleration structure command encoder to encode into this command buffer.
        #[unsafe(method(accelerationStructureCommandEncoder))]
        #[unsafe(method_family = none)]
        fn acceleration_structure_command_encoder(
            &self
        ) -> Option<Retained<ProtocolObject<dyn MTLAccelerationStructureCommandEncoder>>>;

        /// Returns an acceleration structure command encoder configured by the given descriptor.
        #[unsafe(method(accelerationStructureCommandEncoderWithDescriptor:))]
        #[unsafe(method_family = none)]
        fn acceleration_structure_command_encoder_with_descriptor(
            &self,
            descriptor: &MTLAccelerationStructurePassDescriptor,
        ) -> Retained<ProtocolObject<dyn MTLAccelerationStructureCommandEncoder>>;

        /// Marks a residency set as part of this command buffer's execution.
        ///
        /// Availability: macOS 15.0+, iOS 18.0+
        #[unsafe(method(useResidencySet:))]
        #[unsafe(method_family = none)]
        fn use_residency_set(
            &self,
            residency_set: &ProtocolObject<dyn MTLResidencySet>,
        );
    }
);

#[allow(unused)]
pub trait MTLCommandBufferExt: MTLCommandBuffer + Message {
    /// The execution error, if the command buffer failed.
    fn error(&self) -> Option<MetalError>;

    /// A string to help identify this object.
    fn label(&self) -> Option<String>;
    /// Sets a string to help identify this object.
    fn set_label(
        &self,
        label: Option<&str>,
    );
    /// Adds a block to be called when this command buffer has been scheduled for execution.
    fn add_scheduled_handler(
        &self,
        handler: &MTLCommandBufferHandler,
    );
    /// Adds a block to be called when this command buffer has completed execution.
    fn add_completed_handler(
        &self,
        handler: &MTLCommandBufferHandler,
    );
    /// Push a new named string onto a stack of string labels.
    fn push_debug_group(
        &self,
        string: &str,
    );
    /// Pop the latest named string off of the stack.
    fn pop_debug_group(&self);
    /// Marks residency sets as part of this command buffer's execution.
    ///
    /// Availability: macOS 15.0+, iOS 18.0+
    fn use_residency_sets(
        &self,
        residency_sets: &[&ProtocolObject<dyn MTLResidencySet>],
    );
    /// The host time, in seconds, when the CPU began scheduling this command buffer for execution.
    ///
    /// Returns `None` if the command buffer has not been scheduled yet.
    ///
    /// Availability: macOS 10.15+, iOS 10.3+
    fn kernel_start_time(&self) -> Option<f64>;
    /// The host time, in seconds, when the CPU finished scheduling this command buffer for execution.
    ///
    /// Returns `None` if the command buffer has not been scheduled yet.
    ///
    /// Availability: macOS 10.15+, iOS 10.3+
    fn kernel_end_time(&self) -> Option<f64>;
    /// The host time, in seconds, when the GPU started executing this command buffer.
    ///
    /// Availability: macOS 10.15+, iOS 10.3+
    fn gpu_start_time(&self) -> f64;
    /// The host time, in seconds, when the GPU finished executing this command buffer.
    ///
    /// Availability: macOS 10.15+, iOS 10.3+
    fn gpu_end_time(&self) -> f64;
}

impl MTLCommandBufferExt for ProtocolObject<dyn MTLCommandBuffer> {
    fn error(&self) -> Option<MetalError> {
        let error: Option<Retained<NSError>> = unsafe { msg_send![self, error] };
        error.map(MetalError::from_nserror)
    }

    fn label(&self) -> Option<String> {
        let label: Option<Retained<NSString>> = unsafe { msg_send![self, label] };
        label.map(|s| s.to_string())
    }

    fn set_label(
        &self,
        label: Option<&str>,
    ) {
        unsafe {
            let _: () = msg_send![self, setLabel: label.map(NSString::from_str).as_deref()];
        }
    }

    fn add_scheduled_handler(
        &self,
        handler: &MTLCommandBufferHandler,
    ) {
        unsafe {
            let _: () = msg_send![self, addScheduledHandler: handler.as_block()];
        }
    }

    fn add_completed_handler(
        &self,
        handler: &MTLCommandBufferHandler,
    ) {
        unsafe {
            let _: () = msg_send![self, addCompletedHandler: handler.as_block()];
        }
    }

    fn push_debug_group(
        &self,
        string: &str,
    ) {
        unsafe {
            let _: () = msg_send![self, pushDebugGroup: &*NSString::from_str(string)];
        }
    }

    fn pop_debug_group(&self) {
        unsafe {
            let _: () = msg_send![self, popDebugGroup];
        }
    }

    fn use_residency_sets(
        &self,
        residency_sets: &[&ProtocolObject<dyn MTLResidencySet>],
    ) {
        let count = residency_sets.len();
        let residency_sets = ref_slice_as_ptr(residency_sets);
        unsafe {
            let _: () = msg_send![self, useResidencySets: residency_sets, count: count];
        }
    }

    fn kernel_start_time(&self) -> Option<f64> {
        let start_time: f64 = unsafe { msg_send![self, kernelStartTime] };
        if start_time > 0.0 {
            Some(start_time)
        } else {
            None
        }
    }

    fn kernel_end_time(&self) -> Option<f64> {
        let end_time: f64 = unsafe { msg_send![self, kernelEndTime] };
        if end_time > 0.0 {
            Some(end_time)
        } else {
            None
        }
    }

    fn gpu_start_time(&self) -> f64 {
        unsafe { msg_send![self, GPUStartTime] }
    }

    fn gpu_end_time(&self) -> f64 {
        unsafe { msg_send![self, GPUEndTime] }
    }
}

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

    fn assert_send_sync<T: ?Sized + Send + Sync>() {}

    #[test]
    fn command_buffer_is_send_and_sync() {
        assert_send_sync::<dyn MTLCommandBuffer>();
    }
}