//! Queue-bound pins over the shared metadata registry. Hardware operations are
//! supplied by a privileged adapter; host traces never establish device DMA.
//! Completion retains every pin until a checked reset readback. Quarantined
//! records retain their actual Pin and backing even when this scope is dropped.

use core::sync::atomic::{AtomicU64, Ordering};
use turkos_os_frame_metadata::metadata::{
    Backend, Error as MetadataError, Handle, PAGE_BYTES, Pin, Registry,
};

static NEXT_WORKSPACE_ID: AtomicU64 = AtomicU64::new(1);
fn issue_identity() -> Result<u64, Error> {
    let mut current = NEXT_WORKSPACE_ID.load(Ordering::Relaxed);
    loop {
        let next = current.checked_add(1).ok_or(Error::Generation)?;
        match NEXT_WORKSPACE_ID.compare_exchange_weak(
            current,
            next,
            Ordering::Relaxed,
            Ordering::Relaxed,
        ) {
            Ok(_) => return Ok(current),
            Err(observed) => current = observed,
        }
    }
}

pub const MAX_BUFFERS: usize = 8;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    InvalidBinding,
    AlreadyInitialized,
    Foreign,
    Stale,
    Capacity,
    Generation,
    Bounds,
    Overflow,
    Active,
    Duplicate,
    NoSubmission,
    WouldBlock,
    Completion,
    ResetUnanswered,
    Quarantined,
    Audit,
    Metadata(MetadataError),
}
impl Error {
    pub const fn tag(self) -> &'static str {
        match self {
            Self::InvalidBinding => "dma-binding",
            Self::AlreadyInitialized => "dma-initialized",
            Self::Foreign => "dma-foreign",
            Self::Stale => "dma-stale",
            Self::Capacity => "dma-capacity",
            Self::Generation => "dma-generation",
            Self::Bounds => "dma-bounds",
            Self::Overflow => "dma-overflow",
            Self::Active => "dma-active",
            Self::Duplicate => "dma-duplicate",
            Self::NoSubmission => "dma-no-submission",
            Self::WouldBlock => "dma-would-block",
            Self::Completion => "dma-completion",
            Self::ResetUnanswered => "dma-reset-unanswered",
            Self::Quarantined => "dma-quarantined",
            Self::Audit => "dma-audit",
            Self::Metadata(_) => "dma-metadata",
        }
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Binding {
    pub device: u64,
    pub queue: u16,
    pub epoch: u64,
}
impl Binding {
    pub fn new(device: u64, queue: u16, epoch: u64) -> Result<Self, Error> {
        let value = Self {
            device,
            queue,
            epoch,
        };
        value.validate()?;
        Ok(value)
    }
    fn validate(self) -> Result<(), Error> {
        if self.device == 0
            || !self.device.is_multiple_of(512)
            || self.queue >= 1024
            || self.epoch == 0
            || self.device.checked_add(0x200).is_none()
        {
            return Err(Error::InvalidBinding);
        }
        Ok(())
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
    DeviceReads,
    DeviceWrites,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Span {
    pub address: u64,
    pub length: u32,
    pub direction: Direction,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Completion {
    pub id: u32,
    pub length: u32,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ResetWitness {
    pub status: u32,
    pub queue_ready: u32,
}

/// Guest implementations use volatile MMIO/ring accesses and actual ARM64
/// barriers. A fake implementation is useful only for host protocol tests.
pub trait Device {
    fn binding(&self) -> Binding;
    fn publish_barrier(&mut self);
    fn store_available_index(&mut self, index: u16);
    fn notify(&mut self, queue: u16);
    fn read_used_index(&mut self) -> u16;
    fn consume_barrier(&mut self);
    fn read_used_element(&mut self) -> Completion;
    fn write_reset(&mut self);
    fn reset_barrier(&mut self);
    fn read_status(&mut self) -> u32;
    fn read_queue_ready(&mut self, queue: u16) -> u32;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum State {
    Free,
    Pinned,
    Active,
    Submitted,
    Completed,
    Reset,
    Quarantined,
}
struct Record {
    handle: Option<Handle>,
    owner: u64,
    pin: Option<Pin>,
    generation: u64,
    state: State,
}
impl Record {
    const fn empty() -> Self {
        Self {
            handle: None,
            owner: 0,
            pin: None,
            generation: 0,
            state: State::Free,
        }
    }
}
#[derive(Clone, Copy)]
struct Submission {
    head: u16,
    length: u32,
    direction: Direction,
}
pub struct Workspace {
    initialized: bool,
    identity: u64,
    id: u64,
    binding: Binding,
    generation: u64,
    records: [Record; MAX_BUFFERS],
    submission: Option<Submission>,
    reset: Option<ResetWitness>,
}
impl Workspace {
    pub const fn new() -> Self {
        Self {
            initialized: false,
            identity: 0,
            id: 0,
            binding: Binding {
                device: 0,
                queue: 0,
                epoch: 0,
            },
            generation: 0,
            records: [const { Record::empty() }; MAX_BUFFERS],
            submission: None,
            reset: None,
        }
    }
}
impl Default for Workspace {
    fn default() -> Self {
        Self::new()
    }
}

/// Noncopy identity of a Pin held privately in the caller-owned workspace.
#[derive(Debug)]
pub struct Ticket {
    scope: u64,
    workspace: u64,
    binding: Binding,
    slot: usize,
    generation: u64,
    handle: Handle,
}
impl Ticket {
    pub const fn address(&self) -> u64 {
        self.handle.address()
    }
    pub const fn binding(&self) -> Binding {
        self.binding
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Snapshot {
    pub pins: usize,
    pub active: usize,
    pub quarantined: usize,
    pub completed: bool,
    pub reset: Option<ResetWitness>,
    pub registry_pins: usize,
}
pub struct Dma<'a, B, const N: usize> {
    workspace: &'a mut Workspace,
    registry: &'a mut Registry<B, N>,
}
impl<'a, B: Backend, const N: usize> Dma<'a, B, N> {
    pub fn initialize(
        workspace: &'a mut Workspace,
        registry: &'a mut Registry<B, N>,
        id: u64,
        binding: Binding,
    ) -> Result<Self, Error> {
        binding.validate()?;
        if id == 0 {
            return Err(Error::Foreign);
        }
        if workspace.initialized {
            return Err(Error::AlreadyInitialized);
        }
        if workspace
            .records
            .iter()
            .any(|record| record.pin.is_some() || record.state != State::Free)
        {
            return Err(Error::Audit);
        }
        if workspace.identity == 0 {
            workspace.identity = issue_identity()?;
        }
        workspace.initialized = true;
        workspace.id = id;
        workspace.binding = binding;
        workspace.submission = None;
        workspace.reset = None;
        Ok(Self {
            workspace,
            registry,
        })
    }
    pub fn resume(
        workspace: &'a mut Workspace,
        registry: &'a mut Registry<B, N>,
        id: u64,
        binding: Binding,
    ) -> Result<Self, Error> {
        if !workspace.initialized || workspace.id != id || workspace.binding != binding {
            return Err(Error::Foreign);
        }
        let value = Self {
            workspace,
            registry,
        };
        value.audit()?;
        Ok(value)
    }
    pub fn audit(&self) -> Result<Snapshot, Error> {
        if !self.workspace.initialized || self.workspace.id == 0 || self.workspace.identity == 0 {
            return Err(Error::Audit);
        }
        self.workspace.binding.validate()?;
        let mut pins = 0;
        let mut active = 0;
        let mut quarantined = 0;
        for (index, record) in self.workspace.records.iter().enumerate() {
            if record.state == State::Free {
                if record.handle.is_some() || record.pin.is_some() || record.owner != 0 {
                    return Err(Error::Audit);
                }
                continue;
            }
            if record.generation == 0
                || record.generation > self.workspace.generation
                || record.pin.is_none()
            {
                return Err(Error::Audit);
            }
            let handle = record.handle.ok_or(Error::Audit)?;
            self.registry
                .read_byte(handle, record.owner, 0)
                .map_err(Error::Metadata)?;
            if self.workspace.records[..index].iter().any(|prior| {
                prior
                    .handle
                    .is_some_and(|old| old.address() == handle.address())
            }) {
                return Err(Error::Audit);
            }
            let observed = self
                .registry
                .owned_records()
                .find(|owned| owned.address == handle.address())
                .ok_or(Error::Audit)?;
            if observed.owner != record.owner
                || observed.generation != handle.generation()
                || observed.pins == 0
            {
                return Err(Error::Audit);
            }
            pins += 1;
            if matches!(
                record.state,
                State::Active | State::Submitted | State::Completed
            ) {
                active += 1;
            }
            if record.state == State::Quarantined {
                quarantined += 1;
            }
        }
        let registry_pins = self.registry.snapshot().pins;
        if pins > registry_pins {
            return Err(Error::Audit);
        }
        Ok(Snapshot {
            pins,
            active,
            quarantined,
            completed: self
                .workspace
                .records
                .iter()
                .any(|record| record.state == State::Completed),
            reset: self.workspace.reset,
            registry_pins,
        })
    }
    pub fn pin(&mut self, handle: Handle, owner: u64) -> Result<Ticket, Error> {
        self.audit()?;
        if self.workspace.submission.is_some()
            || self.workspace.reset.is_some()
            || self
                .workspace
                .records
                .iter()
                .any(|record| record.state == State::Active)
        {
            return Err(Error::Active);
        }
        if self.workspace.records.iter().any(|record| {
            record
                .handle
                .is_some_and(|old| old.address() == handle.address())
        }) {
            return Err(Error::Duplicate);
        }
        let slot = self
            .workspace
            .records
            .iter()
            .position(|record| record.state == State::Free)
            .ok_or(Error::Capacity)?;
        let generation = self
            .workspace
            .generation
            .checked_add(1)
            .ok_or(Error::Generation)?;
        let pin = self.registry.pin(handle, owner).map_err(Error::Metadata)?;
        self.workspace.records[slot] = Record {
            handle: Some(handle),
            owner,
            pin: Some(pin),
            generation,
            state: State::Pinned,
        };
        self.workspace.generation = generation;
        Ok(Ticket {
            scope: self.workspace.id,
            workspace: self.workspace.identity,
            binding: self.workspace.binding,
            slot,
            generation,
            handle,
        })
    }
    fn record(&self, ticket: &Ticket, binding: Binding) -> Result<&Record, Error> {
        if ticket.workspace != self.workspace.identity
            || binding != self.workspace.binding
            || ticket.binding != binding
            || ticket.scope != self.workspace.id
        {
            return Err(Error::Foreign);
        }
        let record = self
            .workspace
            .records
            .get(ticket.slot)
            .ok_or(Error::Stale)?;
        if record.state == State::Free
            || record.generation != ticket.generation
            || record.handle != Some(ticket.handle)
        {
            return Err(Error::Stale);
        }
        if record.state == State::Quarantined {
            return Err(Error::Quarantined);
        }
        Ok(record)
    }
    pub fn descriptor(
        &self,
        ticket: &Ticket,
        binding: Binding,
        offset: u64,
        length: u64,
        direction: Direction,
    ) -> Result<Span, Error> {
        self.audit()?;
        self.record(ticket, binding)?;
        let end = offset.checked_add(length).ok_or(Error::Overflow)?;
        let address = ticket
            .address()
            .checked_add(offset)
            .ok_or(Error::Overflow)?;
        address.checked_add(length).ok_or(Error::Overflow)?;
        if length == 0 || end > PAGE_BYTES as u64 || length > u32::MAX as u64 {
            return Err(Error::Bounds);
        }
        Ok(Span {
            address,
            length: length as u32,
            direction,
        })
    }
    pub fn ordinary_release(&mut self, handle: Handle, owner: u64) -> Result<(), MetadataError> {
        self.registry.release(handle, owner)
    }
    /// Inert evidence access for the privileged physical authority audit.
    /// No mutable registry or backend reference escapes this scope.
    pub fn observe_registry<R>(&self, observe: impl FnOnce(&Registry<B, N>) -> R) -> R {
        observe(self.registry)
    }
    fn validate_set(&self, tickets: &[&Ticket], state: State) -> Result<(), Error> {
        if tickets.is_empty() || tickets.len() > MAX_BUFFERS {
            return Err(Error::Bounds);
        }
        if tickets.len()
            != self
                .workspace
                .records
                .iter()
                .filter(|record| record.state != State::Free)
                .count()
        {
            return Err(Error::Bounds);
        }
        for (index, ticket) in tickets.iter().enumerate() {
            if tickets[..index]
                .iter()
                .any(|prior| prior.slot == ticket.slot)
            {
                return Err(Error::Duplicate);
            }
            if self.record(ticket, self.workspace.binding)?.state != state {
                return Err(Error::Active);
            }
        }
        Ok(())
    }
    pub fn activate(&mut self, tickets: &[&Ticket]) -> Result<(), Error> {
        self.audit()?;
        self.validate_set(tickets, State::Pinned)?;
        if self.workspace.submission.is_some() {
            return Err(Error::Active);
        }
        for ticket in tickets {
            self.workspace.records[ticket.slot].state = State::Active;
        }
        Ok(())
    }
    pub fn publish(
        &mut self,
        tickets: &[&Ticket],
        head: u16,
        span: Span,
        device: &mut impl Device,
    ) -> Result<(), Error> {
        self.audit()?;
        self.validate_set(tickets, State::Active)?;
        if device.binding() != self.workspace.binding {
            return Err(Error::Foreign);
        }
        if self.workspace.submission.is_some() || head >= 8 || span.length == 0 {
            return Err(Error::Bounds);
        }
        let end = span
            .address
            .checked_add(u64::from(span.length))
            .ok_or(Error::Overflow)?;
        if !tickets.iter().any(|ticket| {
            span.address >= ticket.address() && end <= ticket.address() + PAGE_BYTES as u64
        }) {
            return Err(Error::Bounds);
        }
        device.publish_barrier();
        device.store_available_index(1);
        device.publish_barrier();
        device.notify(self.workspace.binding.queue);
        self.workspace.submission = Some(Submission {
            head,
            length: span.length,
            direction: span.direction,
        });
        for ticket in tickets {
            self.workspace.records[ticket.slot].state = State::Submitted;
        }
        Ok(())
    }
    pub fn consume(&mut self, device: &mut impl Device) -> Result<Completion, Error> {
        self.audit()?;
        if device.binding() != self.workspace.binding {
            return Err(Error::Foreign);
        }
        let submission = self.workspace.submission.ok_or(Error::NoSubmission)?;
        if self
            .workspace
            .records
            .iter()
            .any(|record| record.state == State::Completed)
        {
            return Err(Error::Completion);
        }
        if self
            .workspace
            .records
            .iter()
            .any(|record| record.state == State::Quarantined)
        {
            return Err(Error::Quarantined);
        }
        let index = device.read_used_index();
        if index == 0 {
            return Err(Error::WouldBlock);
        }
        if index != 1 {
            return Err(Error::Completion);
        }
        // faulty acquire barrier omitted
        let used = device.read_used_element();
        if used.id != u32::from(submission.head)
            || match submission.direction {
                Direction::DeviceReads => used.length != 0,
                Direction::DeviceWrites => used.length == 0 || used.length > submission.length,
            }
        {
            return Err(Error::Completion);
        }
        for record in &mut self.workspace.records {
            if record.state == State::Submitted {
                record.state = State::Completed;
            }
        }
        Ok(used)
    }
    pub fn reset(&mut self, device: &mut impl Device) -> Result<ResetWitness, Error> {
        self.audit()?;
        if device.binding() != self.workspace.binding {
            return Err(Error::Foreign);
        }
        if self
            .workspace
            .records
            .iter()
            .any(|record| record.state == State::Quarantined)
        {
            return Err(Error::Quarantined);
        }
        device.write_reset();
        device.reset_barrier();
        let witness = ResetWitness {
            status: device.read_status(),
            queue_ready: device.read_queue_ready(self.workspace.binding.queue),
        };
        self.workspace.reset = Some(witness);
        if witness.status != 0 || witness.queue_ready != 0 {
            for record in &mut self.workspace.records {
                if record.state != State::Free {
                    record.state = State::Quarantined;
                }
            }
            return Err(Error::ResetUnanswered);
        }
        for record in &mut self.workspace.records {
            if record.state != State::Free {
                record.state = State::Reset;
            }
        }
        Ok(witness)
    }
    pub fn retire(&mut self, ticket: Ticket) -> Result<(), (Error, Ticket)> {
        let result = self.retire_checked(&ticket);
        match result {
            Ok(()) => Ok(()),
            Err(error) => Err((error, ticket)),
        }
    }
    fn retire_checked(&mut self, ticket: &Ticket) -> Result<(), Error> {
        self.audit()?;
        let state = self.record(ticket, self.workspace.binding)?.state;
        if !matches!(state, State::Pinned | State::Reset) {
            return Err(Error::Active);
        }
        let record = &mut self.workspace.records[ticket.slot];
        self.registry
            .unpin(record.pin.as_mut().ok_or(Error::Audit)?, record.owner)
            .map_err(Error::Metadata)?;
        record.pin = None;
        record.handle = None;
        record.owner = 0;
        record.state = State::Free;
        Ok(())
    }
    pub fn close(self) -> Result<(), (Error, Self)> {
        match self.audit() {
            Ok(snapshot) if snapshot.pins == 0 => {}
            Ok(_) => return Err((Error::Active, self)),
            Err(error) => return Err((error, self)),
        }
        self.workspace.initialized = false;
        self.workspace.id = 0;
        self.workspace.submission = None;
        self.workspace.reset = None;
        Ok(())
    }
}