//! CommandEncoder dependency and submission rules from MTPLX's pinned runtime
//! (backend/metal/device.cpp). Native code only performs Metal API operations.
//! This is still test-bound; complete model/stream routing remains open.
use super::*;
use std::collections::{HashMap, HashSet};
use std::ffi::CStr;
use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex, OnceLock};

#[derive(Clone, Copy)]
pub(super) enum Access {
    Input(usize),
    Output(usize),
    Raw,
}

impl Binding<'_> {
    /// Byte offset relative to the bound backing, like set_input/output_array.
    /// Geometry validation belongs to the array; native binding checks the base.
    pub(super) fn at_byte_offset(mut self, offset: u64) -> Self {
        assert!(self.allocation.is_some());
        self.raw.offset = offset;
        self
    }

    pub(super) fn input(mut self, elements: usize) -> Self {
        assert!(self.allocation.is_some());
        self.access = Some(Access::Input(elements));
        self
    }

    pub(super) fn output(mut self, elements: usize) -> Self {
        assert!(self.allocation.is_some());
        self.access = Some(Access::Output(elements));
        self
    }

    pub(super) fn raw_buffer(mut self) -> Self {
        assert!(self.allocation.is_some());
        self.access = Some(Access::Raw);
        self
    }
}

thread_local! {
    static ACTIVE: RefCell<Option<Rc<RefCell<Encoder>>>> = const { RefCell::new(None) };
}

/// One encoding-thread stream. The scope owns its queue across operations and
/// routes the existing operators/allocations, without adding dispatch boundaries.
pub(in super::super) struct Scope(Rc<RefCell<Encoder>>);

impl Scope {
    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(in super::super) fn new() -> Result<Self, String> {
        if ACTIVE.with(|active| active.borrow().is_some()) {
            return Err("MTPLX encoder scope is already active".into());
        }
        let encoder = Rc::new(RefCell::new(Encoder::new(allocator::Allocator::new()?)?));
        Self::activate(encoder)
    }

    pub(super) fn activate(encoder: Rc<RefCell<Encoder>>) -> Result<Self, String> {
        if ACTIVE.with(|active| active.borrow().is_some()) {
            return Err("MTPLX encoder scope is already active".into());
        }
        ACTIVE.with(|active| *active.borrow_mut() = Some(encoder.clone()));
        Ok(Self(encoder))
    }

    pub(in super::super) fn synchronize(&self) -> Result<(), String> {
        self.0.borrow_mut().synchronize()
    }

    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(in super::super) fn commits(&self) -> usize {
        self.0.borrow().commits
    }

    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(in super::super) fn primitive_commits(&self) -> usize {
        self.0.borrow().primitive_commits
    }
}

impl Drop for Scope {
    fn drop(&mut self) {
        let owner = ACTIVE.with(|active| active.borrow_mut().take()).unwrap();
        assert!(Rc::ptr_eq(&owner, &self.0));
        // The final Rc drains Encoder after routing is removed. Never leave an
        // expired thread-local owner, including during unwinding.
    }
}

pub(super) fn with_active<T>(operation: impl FnOnce(&mut Encoder) -> T) -> Option<T> {
    let owner = ACTIVE.with(|active| active.borrow().clone());
    owner.map(|owner| operation(&mut owner.borrow_mut()))
}

pub(in super::super) fn active_allocator() -> Option<Arc<allocator::Allocator>> {
    with_active(|encoder| encoder.allocator.clone())
}

/// Both CPU and GPU evaluators retain inputs and siblings by Data identity,
/// excluding the primary output (including an input donated to that output).
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
pub(super) fn primitive_buffers<'a>(
    inputs: &[&'a Buffer],
    outputs: &'a [Buffer],
    primary: usize,
) -> Result<Vec<&'a Buffer>, String> {
    let selected = outputs
        .get(primary)
        .ok_or("invalid evaluator primary output")?;
    Ok(data_buffers(
        inputs.iter().copied().chain(outputs.iter()),
        Some(selected),
    ))
}

pub(super) fn data_buffers<'a>(
    buffers: impl Iterator<Item = &'a Buffer>,
    primary: Option<&Buffer>,
) -> Vec<&'a Buffer> {
    let mut seen = HashSet::new();
    if let Some(primary) = primary {
        seen.insert(primary.data_identity());
    }
    buffers
        .filter(|buffer| seen.insert(buffer.data_identity()))
        .collect()
}

pub(super) fn add_temporaries(buffers: &[Buffer]) {
    with_active(|encoder| {
        for buffer in buffers {
            encoder.add_temporary(buffer);
        }
    });
}

/// Concatenate's disjoint slice writers share an output allocation. Keep their
/// dependency join outside the dispatch borrow, including on early return.
pub(super) struct ConcurrentScope(Option<Rc<RefCell<Encoder>>>);

pub(super) fn concurrent_scope() -> ConcurrentScope {
    let owner = ACTIVE.with(|active| active.borrow().clone());
    if let Some(owner) = &owner {
        owner.borrow_mut().dependencies.concurrent = true;
    }
    ConcurrentScope(owner)
}

impl Drop for ConcurrentScope {
    fn drop(&mut self) {
        if let Some(owner) = &self.0 {
            owner.borrow_mut().dependencies.end_concurrent();
        }
    }
}

#[derive(Default)]
struct Dependencies {
    prev_inputs: HashSet<usize>,
    prev_outputs: HashSet<usize>,
    next_inputs: HashSet<usize>,
    next_outputs: HashSet<usize>,
    concurrent_outputs: HashSet<usize>,
    all_inputs: HashSet<usize>,
    all_outputs: HashSet<usize>,
    needs_barrier: bool,
    concurrent: bool,
    ops: i32,
    // The original adds array.data_size(), measured in ELEMENTS, not bytes.
    data_size: usize,
}

impl Dependencies {
    fn bind(&mut self, resource: usize, access: Access) {
        match access {
            Access::Raw => {
                self.all_inputs.insert(resource);
                self.all_outputs.insert(resource);
            }
            Access::Input(elements) | Access::Output(elements) => {
                if self.all_inputs.insert(resource) {
                    self.data_size = self.data_size.wrapping_add(elements);
                }
                self.next_inputs.insert(resource);
                self.needs_barrier |= self.prev_outputs.contains(&resource);
                if matches!(access, Access::Output(_)) {
                    self.register_output(resource);
                }
            }
        }
    }

    fn register_output(&mut self, resource: usize) {
        self.all_outputs.insert(resource);
        if self.concurrent {
            self.concurrent_outputs.insert(resource);
        } else {
            self.next_outputs.insert(resource);
            self.needs_barrier |= self.prev_inputs.contains(&resource);
        }
    }

    fn dispatched(&mut self) {
        if self.needs_barrier {
            self.needs_barrier = false;
            std::mem::swap(&mut self.prev_inputs, &mut self.next_inputs);
            std::mem::swap(&mut self.prev_outputs, &mut self.next_outputs);
        } else {
            self.prev_inputs.extend(self.next_inputs.iter());
            self.prev_outputs.extend(self.next_outputs.iter());
        }
        self.next_inputs.clear();
        self.next_outputs.clear();
        self.ops += 1;
    }

    fn end_concurrent(&mut self) {
        self.concurrent = false;
        self.prev_outputs.extend(self.concurrent_outputs.drain());
    }

    fn end_encoding(&mut self) {
        self.needs_barrier = false;
        self.concurrent = false;
        self.prev_inputs.clear();
        self.next_inputs.clear();
        self.prev_outputs.clear();
        self.next_outputs.clear();
        self.concurrent_outputs.clear();
        self.all_inputs.clear();
        // all_outputs is moved to the completion handler, like the reference.
        assert!(self.all_outputs.is_empty());
    }

    fn needs_commit(&self, [ops, mb]: [i32; 2]) -> bool {
        self.ops > ops || (self.data_size >> 20) > mb as usize
    }

    fn committed(&mut self) {
        self.ops = 0;
        self.data_size = 0;
    }
}

struct Native(NonNull<c_void>);

impl Native {
    fn new(pointer: *mut c_void, what: &str) -> Result<Self, String> {
        NonNull::new(pointer)
            .map(Self)
            .ok_or_else(|| format!("failed to create MTPLX {what}"))
    }

    fn raw(&self) -> *mut c_void {
        self.0.as_ptr()
    }
}

impl Drop for Native {
    fn drop(&mut self) {
        unsafe { ds4_gpu_mtplx_object_free(self.raw()) };
    }
}

struct Fence(Native);
// SAFETY: only immutable handles/ownership cross threads. Fence encoding is
// confined to Encoder's thread; completion callbacks only drop retained handles.
unsafe impl Send for Fence {}
unsafe impl Sync for Fence {}

type Producers = Arc<Mutex<HashMap<usize, Arc<Fence>>>>;
pub(super) type SharedError = Arc<Mutex<Option<String>>>;

pub(super) struct Event {
    native: Native,
    pub(super) error: Mutex<Option<SharedError>>,
}

// SAFETY: SharedEvent CPU operations are thread-safe native APIs; error pointer
// replacement is mutex protected. No mutable encoder crosses a thread boundary.
unsafe impl Send for Event {}
unsafe impl Sync for Event {}

impl Event {
    pub(super) fn new() -> Result<Arc<Self>, String> {
        Ok(Arc::new(Self {
            native: Native::new(unsafe { ds4_gpu_mtplx_event_create() }, "shared event")?,
            error: Mutex::new(None),
        }))
    }

    pub(super) fn check_error(&self) -> Result<(), String> {
        let error = self.error.lock().unwrap().clone();
        if let Some(error) = error {
            return error.lock().unwrap().take().map_or(Ok(()), Err);
        }
        Ok(())
    }

    pub(super) fn signal(&self, value: u64) {
        unsafe { ds4_gpu_mtplx_event_signal(self.native.raw(), value) };
    }

    pub(super) fn wait(&self, value: u64) -> Result<(), String> {
        self.check_error()?;
        self.wait_raw(value);
        self.check_error()
    }

    pub(super) fn wait_raw(&self, value: u64) {
        unsafe { ds4_gpu_mtplx_event_wait(self.native.raw(), value) };
    }

    pub(super) fn is_signaled(&self, value: u64) -> bool {
        unsafe { ds4_gpu_mtplx_event_value(self.native.raw()) >= value }
    }
}

fn finish_submission(
    error: &SharedError,
    waited: &[Arc<Event>],
    signaled: &[(Arc<Event>, u64)],
    failure: Option<String>,
    completion: Option<Box<dyn FnOnce() + Send>>,
) {
    if let Some(completion) = completion {
        completion();
    }
    for event in waited {
        let source = event.error.lock().unwrap().clone();
        // Snapshot first, then release the source lock before replacing ours.
        // Different queues may complete on different Metal callback threads.
        let message = source
            .filter(|source| !Arc::ptr_eq(source, error))
            .and_then(|source| source.lock().unwrap().clone());
        if let Some(message) = message {
            *error.lock().unwrap() = Some(message);
            break;
        }
    }
    if let Some(failure) = &failure {
        error
            .lock()
            .unwrap()
            .get_or_insert_with(|| format!("[METAL] Command buffer execution failed: {failure}."));
    }
    let has_error = error.lock().unwrap().is_some();
    if has_error {
        for (event, _) in signaled {
            *event.error.lock().unwrap() = Some(error.clone());
        }
    }
    // A GPU error skips Metal's encoded signal. Unblock downstream queues only
    // after associating their events with the shared, consumable error state.
    if failure.is_some() {
        for (event, value) in signaled {
            event.signal(*value);
        }
    }
}

fn remove_completed(producers: &Producers, outputs: &HashSet<usize>, fence: &Arc<Fence>) {
    let mut producers = producers.lock().unwrap();
    for resource in outputs {
        if producers
            .get(resource)
            .is_some_and(|f| Arc::ptr_eq(f, fence))
        {
            producers.remove(resource);
        }
    }
}

fn on_completed(command: &Native, callback: impl FnOnce(Option<String>) + Send + 'static) {
    type Callback = Box<dyn FnOnce(Option<String>) + Send>;
    unsafe extern "C" fn completed(context: *mut c_void, error: *const c_char) {
        let callback = unsafe { Box::from_raw(context.cast::<Callback>()) };
        let error = (!error.is_null()).then(|| {
            unsafe { CStr::from_ptr(error) }
                .to_string_lossy()
                .into_owned()
        });
        callback(error);
    }
    let callback: Callback = Box::new(callback);
    unsafe {
        ds4_gpu_mtplx_command_completed(
            command.raw(),
            Box::into_raw(Box::new(callback)).cast(),
            completed,
        )
    };
}

pub(super) fn arch_suffix() -> Option<char> {
    static ARCH: OnceLock<Option<char>> = OnceLock::new();
    *ARCH.get_or_init(|| {
        let mut suffix = unsafe { ds4_gpu_mtplx_arch_suffix() };
        let arch = unsafe { libc::getenv(c"MLX_METAL_GPU_ARCH".as_ptr()) };
        if !arch.is_null()
            && let Some(last) = unsafe { CStr::from_ptr(arch) }.to_bytes().last()
        {
            suffix = u32::from(*last);
        }
        char::from_u32(suffix)
    })
}

fn limits() -> [i32; 2] {
    static LIMITS: OnceLock<[i32; 2]> = OnceLock::new();
    *LIMITS.get_or_init(|| {
        let defaults = match arch_suffix() {
            Some('p') => [20, 40],
            Some('s' | 'd') => [50, 50],
            _ => [40, 40],
        };
        std::array::from_fn(|i| {
            let name = [c"MLX_MAX_OPS_PER_BUFFER", c"MLX_MAX_MB_PER_BUFFER"][i];
            let value = unsafe { libc::getenv(name.as_ptr()) };
            let default = defaults[i];
            if value.is_null() {
                default
            } else {
                unsafe { libc::atoi(value) }
            }
        })
    })
}

pub(super) struct Encoder {
    command: Option<Native>,
    queue: Native,
    encoding: Option<(Native, Arc<Fence>)>,
    allocator: Arc<allocator::Allocator>,
    attached: usize,
    dependencies: Dependencies,
    temporaries: Vec<(usize, BufferDataHold)>,
    producers: Producers,
    error: SharedError,
    wait_events: Vec<Arc<Event>>,
    signal_events: Vec<(Arc<Event>, u64)>,
    limits: [i32; 2],
    // Diagnostic counts, never used for scheduling.
    barriers: usize,
    waits: usize,
    commits: usize,
    primitive_commits: usize,
}

impl Encoder {
    pub(super) fn new(allocator: Arc<allocator::Allocator>) -> Result<Self, String> {
        let queue = Native::new(unsafe { ds4_gpu_mtplx_queue_create() }, "command queue")?;
        let mut attached = 0;
        unsafe {
            allocator
                .residency
                .attach_new_sets(queue.raw(), &mut attached)
        };
        let command = Native::new(
            unsafe { ds4_gpu_mtplx_command_create(queue.raw()) },
            "command buffer",
        )?;
        Ok(Self {
            queue,
            command: Some(command),
            encoding: None,
            allocator,
            attached,
            dependencies: Dependencies::default(),
            temporaries: Vec::new(),
            producers: Arc::default(),
            error: Arc::default(),
            wait_events: Vec::new(),
            signal_events: Vec::new(),
            limits: limits(),
            barriers: 0,
            waits: 0,
            commits: 0,
            primitive_commits: 0,
        })
    }

    fn check_error(&self) -> Result<(), String> {
        self.error.lock().unwrap().take().map_or(Ok(()), Err)
    }

    fn get_encoder(&mut self) -> Result<*mut c_void, String> {
        if self.encoding.is_none() {
            self.check_error()?;
            let command = self
                .command
                .as_ref()
                .ok_or("MTPLX command buffer unavailable")?;
            let encoder = Native::new(
                unsafe { ds4_gpu_mtplx_encoder_create(command.raw()) },
                "concurrent encoder",
            )?;
            let fence = match Native::new(unsafe { ds4_gpu_mtplx_fence_create() }, "fence") {
                Ok(fence) => Arc::new(Fence(fence)),
                Err(error) => {
                    unsafe { ds4_gpu_mtplx_encoder_end(encoder.raw()) };
                    return Err(error);
                }
            };
            self.encoding = Some((encoder, fence));
        }
        Ok(self.encoding.as_ref().unwrap().0.raw())
    }

    pub(super) fn dispatch(
        &mut self,
        name: &str,
        bindings: &[Binding<'_>],
        constants: &[MetalConstant],
        grid: [u32; 3],
        threads: [u32; 3],
        individual_threads: bool,
    ) -> Result<(), String> {
        if bindings
            .iter()
            .any(|b| b.allocation.is_some() && b.access.is_none())
        {
            return Err("MTPLX encoder requires explicit input/output/raw binding roles".into());
        }
        let encoder = self.get_encoder()?;
        for binding in bindings {
            if let Some(access) = binding.access {
                let resource =
                    unsafe { ds4_gpu_mtplx_tensor_allocation(binding.raw.tensor) } as usize;
                self.dependencies.bind(resource, access);
            }
        }
        let barrier = self.dependencies.needs_barrier;
        let result = dispatch_geometry_on(
            name,
            bindings,
            constants,
            grid,
            threads,
            individual_threads,
            Some((self.command.as_ref().unwrap().raw(), encoder, barrier)),
        );
        match result {
            Ok(()) => {
                self.barriers += usize::from(barrier);
                self.dependencies.dispatched();
                Ok(())
            }
            Err(error) => Err(error),
        }
    }

    pub(super) fn barrier(&mut self) -> Result<(), String> {
        let encoder = self.get_encoder()?;
        unsafe { ds4_gpu_mtplx_encoder_barrier(encoder) };
        self.barriers += 1;
        // Explicit barriers do not change the reference's access bookkeeping.
        Ok(())
    }

    pub(super) fn register_output(&mut self, buffer: &Buffer) {
        let resource = unsafe { ds4_gpu_mtplx_tensor_allocation(buffer.raw()) } as usize;
        self.dependencies.register_output(resource);
    }

    pub(super) fn hold_until_completed(&self, hold: impl Send + 'static) -> Result<(), String> {
        let command = self
            .command
            .as_ref()
            .ok_or("MTPLX command buffer unavailable")?;
        on_completed(command, move |_| drop(hold));
        Ok(())
    }

    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    fn start_concurrent(&mut self) -> Concurrent<'_> {
        self.dependencies.concurrent = true;
        Concurrent(self)
    }

    fn add_temporary(&mut self, buffer: &Buffer) {
        let resource = unsafe { ds4_gpu_mtplx_tensor_allocation(buffer.raw()) } as usize;
        self.temporaries.push((resource, buffer.hold_data()));
    }

    /// Evaluator input/sibling Data is completion-owned, not backend scratch.
    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(super) fn defer_data(&self, buffers: &[&Buffer]) -> bool {
        let Some(command) = &self.command else {
            return false;
        };
        let data = buffers
            .iter()
            .map(|buffer| buffer.hold_data())
            .collect::<Vec<_>>();
        on_completed(command, move |_| drop(data));
        true
    }

    /// gpu::eval runs this once after the whole primitive, not per dispatch.
    pub(super) fn eval_primitive(&mut self, buffers: &[&Buffer]) -> Result<(), String> {
        let data = buffers
            .iter()
            .map(|buffer| buffer.hold_data())
            .collect::<Vec<_>>();
        if self.needs_commit() {
            self.end_encoding();
            let completion = scheduler::global().activity.start();
            self.commit_with(Some(Box::new(move || {
                // Reference task completion is announced before capture release.
                drop(completion);
                drop(data);
            })))?;
            self.primitive_commits += 1;
        } else {
            let command = self
                .command
                .as_ref()
                .ok_or("MTPLX command buffer unavailable")?;
            on_completed(command, move |_| drop(data));
        }
        Ok(())
    }

    /// The fixed indexer tape has one open GPU stream. The general evaluator
    /// must finalize all its open streams before using the same pressure gate.
    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
    pub(super) fn check_backpressure(&mut self) -> Result<(), String> {
        let activity = &scheduler::global().activity;
        if activity.needs_throttle(self.allocator.stats().active, self.allocator.memory_limit()) {
            self.end_encoding();
            self.commit()?;
            activity.wait_for_one();
            while self.allocator.stats().active > self.allocator.memory_limit()
                && activity.count() > 0
            {
                activity.wait_for_one();
            }
        }
        Ok(())
    }

    fn end_encoding(&mut self) {
        let Some((encoder, fence)) = self.encoding.take() else {
            return;
        };
        for (resource, _) in &self.temporaries {
            self.dependencies.all_inputs.remove(resource);
            self.dependencies.all_outputs.remove(resource);
        }
        let mut waiting_on = HashMap::new();
        {
            let mut producers = self.producers.lock().unwrap();
            for resource in &self.dependencies.all_inputs {
                if let Some(prior) = producers.get(resource)
                    && let std::collections::hash_map::Entry::Vacant(entry) =
                        waiting_on.entry(Arc::as_ptr(prior) as usize)
                {
                    unsafe { ds4_gpu_mtplx_encoder_fence(encoder.raw(), prior.0.raw(), 0) };
                    self.waits += 1;
                    entry.insert(prior.clone());
                }
            }
            for resource in &self.dependencies.all_outputs {
                producers.insert(*resource, fence.clone());
            }
        }
        unsafe { ds4_gpu_mtplx_encoder_fence(encoder.raw(), fence.0.raw(), 1) };
        let producers = self.producers.clone();
        let outputs = std::mem::take(&mut self.dependencies.all_outputs);
        let temporaries = std::mem::take(&mut self.temporaries);
        on_completed(self.command.as_ref().unwrap(), move |_| {
            remove_completed(&producers, &outputs, &fence);
            drop((temporaries, waiting_on));
        });
        unsafe { ds4_gpu_mtplx_encoder_end(encoder.raw()) };
        self.dependencies.end_encoding();
    }

    fn needs_commit(&self) -> bool {
        self.dependencies.needs_commit(self.limits)
    }

    pub(super) fn wait_event(&mut self, event: Arc<Event>, value: u64) -> Result<(), String> {
        self.end_encoding();
        let command = self
            .command
            .as_ref()
            .ok_or("MTPLX command buffer unavailable")?;
        unsafe { ds4_gpu_mtplx_command_event(command.raw(), event.native.raw(), value, 0) };
        self.wait_events.push(event);
        Ok(())
    }

    pub(super) fn signal_event(&mut self, event: Arc<Event>, value: u64) -> Result<(), String> {
        self.end_encoding();
        let command = self
            .command
            .as_ref()
            .ok_or("MTPLX command buffer unavailable")?;
        unsafe { ds4_gpu_mtplx_command_event(command.raw(), event.native.raw(), value, 1) };
        self.signal_events.push((event, value));
        Ok(())
    }

    fn commit(&mut self) -> Result<Native, String> {
        self.commit_with(None)
    }

    fn commit_with(
        &mut self,
        completion: Option<Box<dyn FnOnce() + Send>>,
    ) -> Result<Native, String> {
        // Reference callers end_encoding before commit; do not silently add an
        // encoder boundary here, since that would change scheduling semantics.
        if self.encoding.is_some() {
            return Err("MTPLX commit requires end_encoding".into());
        }
        unsafe {
            self.allocator
                .residency
                .attach_new_sets(self.queue.raw(), &mut self.attached)
        };
        let command = self
            .command
            .take()
            .ok_or("MTPLX command buffer unavailable")?;
        let error = self.error.clone();
        let waited = std::mem::take(&mut self.wait_events);
        let signaled = std::mem::take(&mut self.signal_events);
        on_completed(&command, move |failure| {
            finish_submission(&error, &waited, &signaled, failure, completion);
        });
        unsafe { ds4_gpu_mtplx_command_commit(command.raw()) };
        self.commits += 1;
        match Native::new(
            unsafe { ds4_gpu_mtplx_command_create(self.queue.raw()) },
            "command buffer",
        ) {
            Ok(next) => self.command = Some(next),
            Err(error) => {
                // A submitted unretained buffer must drain even if allocating
                // its replacement fails. Never reuse the submitted buffer.
                unsafe { ds4_gpu_mtplx_command_wait(command.raw()) };
                return Err(error);
            }
        }
        self.dependencies.committed();
        Ok(command)
    }

    pub(super) fn synchronize(&mut self) -> Result<(), String> {
        self.drain()?;
        self.check_error()
    }

    pub(super) fn finalize(&mut self) -> Result<(), String> {
        self.end_encoding();
        self.commit().map(drop)
    }

    fn drain(&mut self) -> Result<(), String> {
        self.end_encoding();
        let submitted = self.commit()?;
        unsafe { ds4_gpu_mtplx_command_wait(submitted.raw()) };
        Ok(())
    }
}

impl Drop for Encoder {
    fn drop(&mut self) {
        // Drain before queue, allocator and completion-owned state are released.
        // Like the reference destructor, do not rethrow an asynchronous GPU error.
        let _ = self.drain();
    }
}

#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
struct Concurrent<'a>(&'a mut Encoder);
impl Deref for Concurrent<'_> {
    type Target = Encoder;
    fn deref(&self) -> &Encoder {
        self.0
    }
}
impl DerefMut for Concurrent<'_> {
    fn deref_mut(&mut self) -> &mut Encoder {
        self.0
    }
}
impl Drop for Concurrent<'_> {
    fn drop(&mut self) {
        self.0.dependencies.end_concurrent();
    }
}

#[test]
fn mtplx_encoder_barrier_epochs_and_commit_thresholds() {
    let mut d = Dependencies::default();
    d.bind(1, Access::Input(8));
    d.bind(2, Access::Output(8));
    assert!(!d.needs_barrier);
    d.dispatched();
    assert_eq!(d.data_size, 16);
    d.bind(2, Access::Input(800)); // Same physical resource counts only once.
    d.bind(3, Access::Output(8));
    assert!(d.needs_barrier, "read after write");
    d.dispatched();
    assert_eq!(d.prev_inputs, HashSet::from([2, 3]));
    assert_eq!(d.prev_outputs, HashSet::from([3]));
    assert_eq!(d.data_size, 24);
    d.bind(1, Access::Output(8));
    assert!(!d.needs_barrier, "resource 1 is outside the current epoch");
    d.dispatched();
    d.bind(2, Access::Output(8));
    assert!(d.needs_barrier, "write after read");
    d.dispatched();

    let mut raw = Dependencies::default();
    raw.bind(9, Access::Raw);
    raw.dispatched();
    raw.bind(9, Access::Input(1024));
    assert!(!raw.needs_barrier);
    assert_eq!(raw.data_size, 0, "set_buffer already inserted all_inputs");
    assert_eq!(raw.all_outputs, HashSet::from([9]));

    let mut concurrent = Dependencies {
        concurrent: true,
        ..Default::default()
    };
    for _ in 0..2 {
        concurrent.bind(1, Access::Input(8));
        concurrent.bind(2, Access::Output(8));
        assert!(!concurrent.needs_barrier, "concurrent outputs are deferred");
        concurrent.dispatched();
    }
    assert!(concurrent.prev_outputs.is_empty());
    concurrent.end_concurrent();
    concurrent.bind(2, Access::Input(8));
    assert!(concurrent.needs_barrier);
    concurrent.dispatched();

    let mut threshold = Dependencies {
        ops: 50,
        data_size: (51 << 20) - 1,
        ..Default::default()
    };
    assert!(
        !threshold.needs_commit([50, 50]),
        "strict >, truncated Mi elements"
    );
    threshold.data_size += 1;
    assert!(threshold.needs_commit([50, 50]));
    threshold.data_size = 0;
    threshold.ops += 1;
    assert!(threshold.needs_commit([50, 50]));
    threshold.end_encoding();
    assert_eq!(
        threshold.ops, 51,
        "encoder end does not reset command counters"
    );
    threshold.committed();
    assert_eq!((threshold.ops, threshold.data_size), (0, 0));
    assert!(threshold.needs_commit([-1, -1]), "signed op threshold");
    assert!(
        !threshold.needs_commit([50, -1]),
        "reference size_t conversion"
    );
    d.all_outputs.clear(); // Moved into the completion handler in the real encoder.
    let counted = d.data_size;
    let ops = d.ops;
    d.end_encoding();
    d.bind(1, Access::Input(8));
    assert_eq!(d.data_size, counted + 8, "unique inputs reset per encoder");
    assert_eq!(
        d.ops, ops,
        "command counters persist across encoder boundaries"
    );
}

#[test]
#[ignore = "requires Apple Metal, no model weights"]
fn mtplx_encoder_commits_after_primitives_and_balances_completion() {
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    crate::engine::reset_gpu_busy_stats();
    crate::engine::gpu_canary_probe().unwrap();
    assert_eq!(crate::engine::gpu_busy_stats().command_buffers, 0);
    let scope = Scope::new().unwrap();
    scope.0.borrow_mut().limits = [1, i32::MAX];
    let input = Buffer::mtplx_bytes(16).unwrap();
    let middle = Buffer::mtplx_bytes(16).unwrap();
    let output = Buffer::mtplx_bytes(16).unwrap();
    let expected = [0x80, 0x3f].repeat(8);
    input.write(0, &expected).unwrap();
    let copy = |from: &Buffer, to: &Buffer| {
        qsa_copy_typed(
            Some(from),
            to,
            [1, 1, 8],
            [0, 0, 1],
            [0, 0, 1],
            "bfloat16",
            8,
        )
        .unwrap();
    };
    copy(&input, &middle);
    with_active(|encoder| encoder.eval_primitive(&[&input]))
        .unwrap()
        .unwrap();
    assert_eq!(scope.commits(), 0, "the threshold is strict, not >=");
    copy(&middle, &output);
    assert_eq!(
        scope.commits(),
        0,
        "dispatch must not commit a partial primitive"
    );
    with_active(|encoder| encoder.eval_primitive(&[&middle]))
        .unwrap()
        .unwrap();
    assert_eq!(scope.primitive_commits(), 1);
    assert_eq!(scope.commits(), 1);
    scope.synchronize().unwrap();
    let busy = crate::engine::gpu_busy_stats();
    assert_eq!(busy.command_buffers, scope.commits() as u64);
    assert!(busy.busy_seconds > 0.0);
    assert!(busy.max_busy_seconds > 0.0);
    assert_eq!(scheduler::global().activity.count(), 0);
    assert!(input.unique_data());
    assert!(middle.unique_data());
    let mut actual = [0; 16];
    output.read(0, &mut actual).unwrap();
    assert_eq!(actual.as_slice(), expected);

    // One primitive may encode several dependent kernels before the check.
    copy(&input, &middle);
    copy(&middle, &output);
    copy(&output, &middle);
    assert_eq!(scope.commits(), 2);
    with_active(|encoder| encoder.eval_primitive(&[&input, &output]))
        .unwrap()
        .unwrap();
    assert_eq!(scope.primitive_commits(), 2);
    scope.synchronize().unwrap();
    assert_eq!(
        crate::engine::gpu_busy_stats().command_buffers,
        scope.commits() as u64,
        "record every MTPLX command once, without counting its later waits"
    );
    assert_eq!(scheduler::global().activity.count(), 0);
    middle.read(0, &mut actual).unwrap();
    assert_eq!(actual.as_slice(), expected);

    // Fail before callback transfer, without submitting invalid GPU work.
    let mut encoder = scope.0.borrow_mut();
    let command = encoder.command.take();
    encoder.dependencies.ops = 2;
    assert_eq!(
        encoder.eval_primitive(&[]),
        Err("MTPLX command buffer unavailable".into())
    );
    assert_eq!(scheduler::global().activity.count(), 0);
    encoder.command = command;
    encoder.dependencies.committed();
}

#[test]
#[ignore = "requires Apple Metal, no model weights"]
fn mtplx_encoder_scope_routes_operators_and_allocator() {
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let scope = Scope::new().unwrap();
    assert!(
        Scope::new().is_err(),
        "nested scopes cannot replace the live stream"
    );
    let allocator = active_allocator().unwrap();
    let input = Buffer::mtplx_bytes(3 * 64 * 4).unwrap();
    input.write(0, &[0x80, 0x3f].repeat(3 * 64 * 2)).unwrap();
    let output = Buffer::mtplx_bytes(3 * 64 * 2).unwrap();
    let activated = gate_up_bf16(&input, &output, [3, 64]).unwrap();
    {
        let encoder = scope.0.borrow();
        assert_eq!(encoder.dependencies.ops, 2);
        assert_eq!(encoder.barriers, 1);
        assert_eq!(
            encoder.dependencies.data_size,
            11 * 64,
            "slice span (2*rows-1)*width plus two contiguous outputs, not full allocation bytes"
        );
        assert_eq!(encoder.commits, 0, "scope must not invent operator commits");
    }
    let bytes = [&input, &output, &activated]
        .iter()
        .map(|buffer| buffer.allocation_bytes() as usize)
        .sum::<usize>();
    assert_eq!(
        allocator.stats().active,
        bytes,
        "all operator storage uses the pooled allocator"
    );
    scope.synchronize().unwrap();
    drop((input, output, activated));
    assert_eq!(allocator.stats().active, 0);
    let failed_copy = || -> Result<(), String> {
        let _concurrent = concurrent_scope();
        assert!(scope.0.borrow().dependencies.concurrent);
        Err("synthetic concatenate encoding failure".into())
    };
    assert!(failed_copy().is_err());
    assert!(
        !scope.0.borrow().dependencies.concurrent,
        "an early return must restore the dependency join"
    );
    let source = Buffer::mtplx_bytes(32).unwrap();
    let target = Buffer::mtplx_bytes(64).unwrap();
    source.write(0, &[0x80, 0x3f].repeat(16)).unwrap();
    assert!(qsa_dynamic_copy_rows([&source, &target], [8; 2], [1, 8], [None; 2], [7, 32]).is_err());
    qsa_dynamic_copy_rows([&source, &target], [8; 2], [1, 8], [None; 2], [16, 32]).unwrap();
    assert_eq!(
        scope.0.borrow().dependencies.data_size,
        48,
        "copy counts bound array data_size, not the eight copied elements"
    );
    scope.synchronize().unwrap();
    let mut copied = [0; 16];
    target.read(0, &mut copied).unwrap();
    assert_eq!(copied.as_slice(), [0x80, 0x3f].repeat(8));
    drop((source, target));
    assert!(
        allocator.stats().cached > 0,
        "completed storage returns to the same cache"
    );
    drop(scope);
    assert!(
        active_allocator().is_none(),
        "scope teardown removes routing"
    );
    let outside = Buffer::mtplx_bytes(128).unwrap();
    assert_eq!(
        allocator.stats().active,
        0,
        "legacy work must not inherit the allocator"
    );
    drop(outside);
}

#[test]
#[ignore = "requires Apple Metal; empty original-kernel grids and nil bindings"]
fn mtplx_empty_bindings_preserve_nil_dependencies_without_dummy_storage() {
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    {
        let scope = Scope::new().unwrap();
        let allocator = active_allocator().unwrap();
        let real = Buffer::mtplx_bytes(16).unwrap();
        let before = allocator.stats();
        let empty = Buffer::mtplx_bytes(0).unwrap();
        let another = Buffer::mtplx_bytes(0).unwrap();
        let nthreads = 0_u32;
        for binding in [tensor(0, &empty).output(0), tensor(0, &another).input(0)] {
            scope
                .0
                .borrow_mut()
                .dispatch(
                    "input_coherent",
                    &[binding, bytes(1, &nthreads)],
                    &[],
                    [0, 1, 1],
                    [1024, 1, 1],
                    false,
                )
                .unwrap();
        }
        {
            let encoder = scope.0.borrow();
            assert_eq!(encoder.dependencies.data_size, 0);
            assert_eq!(encoder.dependencies.ops, 2);
            assert_eq!(encoder.dependencies.all_inputs, HashSet::from([0]));
            assert_eq!(encoder.dependencies.all_outputs, HashSet::from([0]));
            assert_eq!(
                encoder.barriers, 1,
                "null resources participate in RAW dependencies"
            );
        }
        assert_eq!(allocator.stats(), before);
        // A zero-length slice can also retain a non-null allocation and point
        // exactly one past its end. Do not substitute null Data for that view.
        let tail = real.view(16, 0).unwrap();
        assert_eq!(tail.data_identity(), real.data_identity());
        assert_eq!(tail.allocation_bytes(), 16);
        scope
            .0
            .borrow_mut()
            .dispatch(
                "input_coherent",
                &[tensor(0, &tail).input(0), bytes(1, &nthreads)],
                &[],
                [0, 1, 1],
                [1024, 1, 1],
                false,
            )
            .unwrap();
        assert_eq!(scope.0.borrow().dependencies.data_size, 0);
        drop((empty, another, tail));
        scope.synchronize().unwrap();
        assert_eq!(allocator.stats(), before);
    }
    // Also cover the native fallback owner path, which has an Objective-C
    // resource set. nil is a valid Metal binding but not an NSSet element.
    let empty = Buffer::mtplx_bytes(0).unwrap();
    let nthreads = 0_u32;
    let bindings = [tensor(0, &empty).input(0).raw, bytes(1, &nthreads).raw];
    assert_ne!(
        unsafe {
            ds4_gpu_mtplx_dispatch(
                c"input_coherent".as_ptr(),
                bindings.as_ptr(),
                2,
                std::ptr::null(),
                0,
                [0, 1, 1].as_ptr(),
                [1024, 1, 1].as_ptr(),
                0,
                std::ptr::null_mut(),
                None,
            )
        },
        0
    );
}

#[test]
#[ignore = "requires Apple Metal, no model weights"]
fn mtplx_encoder_untracked_dependencies_across_encoders_and_commands() {
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let allocator = allocator::Allocator::new().unwrap();
    allocator.set_wired_limit(8 << 20).unwrap();
    let mut enc = Encoder::new(allocator.clone()).unwrap();
    assert_ne!(enc.queue.raw(), unsafe { ds4_gpu_mtplx_command_queue() });
    assert_eq!(enc.attached, allocator.residency.num_sets());
    const N: u32 = 4096;
    let a = allocator.allocate((N * 2) as usize).unwrap().unwrap();
    let x = allocator.allocate((N * 2) as usize).unwrap().unwrap();
    let y = allocator.allocate((N * 2) as usize).unwrap().unwrap();
    let expected = (0..N)
        .flat_map(|i| (0x3f00_u16 + (i % 128) as u16).to_le_bytes())
        .collect::<Vec<_>>();
    a.write(0, &expected).unwrap();
    let copy = |enc: &mut Encoder, source: &Buffer, output: &Buffer, count: u32| {
        enc.dispatch(
            "v_copybfloat16bfloat16",
            &[
                tensor(0, source).input(count as usize),
                tensor(1, output).output(count as usize),
                bytes(2, &count),
            ],
            &[],
            [count, 1, 1],
            [count.min(1024), 1, 1],
            true,
        )
        .unwrap();
    };
    assert!(
        enc.dispatch(
            "v_copybfloat16bfloat16",
            &[tensor(0, &a), tensor(1, &x), bytes(2, &N)],
            &[],
            [N, 1, 1],
            [1024, 1, 1],
            true
        )
        .is_err()
    );
    assert!(
        enc.encoding.is_none(),
        "unclassified bindings cannot silently use serial semantics"
    );
    {
        let mut concurrent = enc.start_concurrent();
        for offset in [0, N] {
            let input = a.view(u64::from(offset), u64::from(N)).unwrap();
            let output = x.view(u64::from(offset), u64::from(N)).unwrap();
            copy(&mut concurrent, &input, &output, N / 2);
        }
    }
    assert_eq!(enc.barriers, 0);
    copy(&mut enc, &x, &y, N);
    assert_eq!(
        enc.barriers, 1,
        "join concurrent writers before their reader"
    );
    let previous = enc.dependencies.prev_outputs.clone();
    enc.barrier().unwrap();
    assert_eq!(
        enc.dependencies.prev_outputs, previous,
        "explicit barrier preserves bookkeeping"
    );
    assert_eq!(enc.barriers, 2);
    enc.end_encoding();
    assert_eq!(enc.waits, 0);
    assert_eq!(enc.dependencies.ops, 3);
    copy(&mut enc, &y, &x, N);
    enc.end_encoding();
    assert_eq!(enc.waits, 1, "two prior resources share one fence wait");
    assert_eq!(enc.dependencies.ops, 4);
    let submitted = enc.commit().unwrap(); // Deliberately no CPU wait.
    assert_eq!(enc.dependencies.ops, 0);
    assert!(!enc.needs_commit());
    for _ in 0..6 {
        copy(&mut enc, &x, &y, N);
        copy(&mut enc, &y, &x, N);
        enc.end_encoding();
        drop(enc.commit().unwrap());
    }
    enc.synchronize().unwrap();
    drop(submitted);
    assert!(enc.producers.lock().unwrap().is_empty());
    let mut actual = vec![0; expected.len()];
    x.read(0, &mut actual).unwrap();
    assert_eq!(actual, expected);

    // Completion must not erase a newer writer's entry. Exercise the exact
    // callback helper deterministically, independent of GPU completion timing.
    let old = Arc::new(Fence(
        Native::new(unsafe { ds4_gpu_mtplx_fence_create() }, "fence").unwrap(),
    ));
    let new = Arc::new(Fence(
        Native::new(unsafe { ds4_gpu_mtplx_fence_create() }, "fence").unwrap(),
    ));
    let producers: Producers = Arc::default();
    producers.lock().unwrap().insert(123, new.clone());
    remove_completed(&producers, &HashSet::from([123]), &old);
    assert!(Arc::ptr_eq(
        producers.lock().unwrap().get(&123).unwrap(),
        &new
    ));
    remove_completed(&producers, &HashSet::from([123]), &new);
    assert!(producers.lock().unwrap().is_empty());

    // Temporaries do not enter the inter-encoder producer map. They stay alive
    // through completion even after their descriptors and binding views drop.
    let temporary = allocator.allocate((N * 2) as usize).unwrap().unwrap();
    let resource = unsafe { ds4_gpu_mtplx_tensor_allocation(temporary.raw()) } as usize;
    copy(&mut enc, &a, &temporary, N);
    enc.add_temporary(&temporary);
    drop(temporary);
    enc.end_encoding();
    assert!(!enc.producers.lock().unwrap().contains_key(&resource));
    enc.synchronize().unwrap();

    // Raw buffers match set_buffer: cross-encoder input AND output, without
    // automatic intra-encoder barriers or data_size accounting.
    enc.dispatch(
        "v_copybfloat16bfloat16",
        &[
            tensor(0, &a).raw_buffer(),
            tensor(1, &x).raw_buffer(),
            bytes(2, &N),
        ],
        &[],
        [N, 1, 1],
        [1024, 1, 1],
        true,
    )
    .unwrap();
    assert_eq!(enc.dependencies.data_size, 0);
    enc.end_encoding();
    assert_eq!(enc.producers.lock().unwrap().len(), 2);
    enc.synchronize().unwrap();
    assert!(enc.producers.lock().unwrap().is_empty());

    // Destruction must submit and wait with the allocator alive.
    copy(&mut enc, &a, &y, N);
    drop(enc);
    y.read(0, &mut actual).unwrap();
    assert_eq!(actual, expected);
}

#[test]
#[ignore = "requires Apple Metal, no model weights"]
fn mtplx_encoder_shared_events_and_error_completion_order() {
    use std::sync::atomic::{AtomicUsize, Ordering};
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let allocator = allocator::Allocator::new().unwrap();
    allocator.set_wired_limit(8 << 20).unwrap();
    let input = allocator.allocate(128).unwrap().unwrap();
    let middle = allocator.allocate(128).unwrap().unwrap();
    let output = allocator.allocate(128).unwrap().unwrap();
    let expected = [0x80, 0x3f].repeat(64);
    input.write(0, &expected).unwrap();
    let mut producer = Encoder::new(allocator.clone()).unwrap();
    let mut consumer = Encoder::new(allocator).unwrap();
    let ready = Event::new().unwrap();
    let done = Event::new().unwrap();
    assert!(!ready.is_signaled(1));
    assert!(!done.is_signaled(2));
    let copy = |enc: &mut Encoder, src: &Buffer, dst: &Buffer| {
        enc.dispatch(
            "v_copybfloat16bfloat16",
            &[
                tensor(0, src).input(64),
                tensor(1, dst).output(64),
                bytes(2, &64_u32),
            ],
            &[],
            [2, 1, 1],
            [32, 1, 1],
            false,
        )
        .unwrap();
    };
    copy(&mut producer, &input, &middle);
    producer.signal_event(ready.clone(), 1).unwrap();
    assert!(
        producer.encoding.is_none(),
        "signal ends the active encoder"
    );
    let completed = Arc::new(AtomicUsize::new(0));
    let count = completed.clone();
    drop(
        producer
            .commit_with(Some(Box::new(move || {
                count.fetch_add(1, Ordering::SeqCst);
            })))
            .unwrap(),
    );
    consumer.wait_event(ready.clone(), 1).unwrap();
    copy(&mut consumer, &middle, &output);
    consumer.signal_event(done.clone(), 2).unwrap();
    consumer.synchronize().unwrap();
    producer.synchronize().unwrap();
    assert_eq!(completed.load(Ordering::SeqCst), 1);
    assert!(ready.is_signaled(1));
    assert!(done.is_signaled(2));
    done.wait(2).unwrap();
    let mut actual = vec![0; 128];
    output.read(0, &mut actual).unwrap();
    assert_eq!(
        actual, expected,
        "untracked buffer crosses queues via shared event"
    );

    // Invoke the real completion policy with a synthetic status error; never
    // deliberately hang/fault the GPU to exercise downstream unblocking.
    let first = Event::new().unwrap();
    let second = Event::new().unwrap();
    let failed = Event::new().unwrap();
    let error: SharedError = Arc::new(Mutex::new(Some("local".into())));
    *first.error.lock().unwrap() = Some(Arc::new(Mutex::new(Some("first upstream".into()))));
    *second.error.lock().unwrap() = Some(Arc::new(Mutex::new(Some("second upstream".into()))));
    let observed = error.clone();
    finish_submission(
        &error,
        &[first, second],
        &[(failed.clone(), 3)],
        Some("GPU failure".into()),
        Some(Box::new(move || {
            assert_eq!(observed.lock().unwrap().as_deref(), Some("local"))
        })),
    );
    assert!(
        failed.is_signaled(3),
        "GPU error manually signals downstream event"
    );
    assert_eq!(error.lock().unwrap().as_deref(), Some("first upstream"));
    assert_eq!(failed.wait(3).unwrap_err(), "first upstream");
    assert!(
        error.lock().unwrap().is_none(),
        "event consumes the shared encoder error"
    );
    failed.check_error().unwrap();

    // A self-referencing waited event is skipped, not recursively locked. A
    // non-GPU upstream error poisons signals but does not signal them on CPU.
    let self_event = Event::new().unwrap();
    *self_event.error.lock().unwrap() = Some(error.clone());
    let downstream = Event::new().unwrap();
    *error.lock().unwrap() = Some("existing".into());
    finish_submission(
        &error,
        &[self_event],
        &[(downstream.clone(), 4)],
        None,
        None,
    );
    assert!(!downstream.is_signaled(4));
    assert_eq!(downstream.check_error().unwrap_err(), "existing");
    finish_submission(&error, &[], &[], Some("synthetic".into()), None);
    assert_eq!(
        error.lock().unwrap().as_deref(),
        Some("[METAL] Command buffer execution failed: synthetic.")
    );
    let retained_error = producer.error.clone();
    *retained_error.lock().unwrap() = Some("destructor must preserve this".into());
    drop(producer);
    assert_eq!(
        retained_error.lock().unwrap().as_deref(),
        Some("destructor must preserve this")
    );
}