vyre-driver 0.7.2

use std::collections::BTreeMap;

use vyre_megakernel::{
    Artifact, ArtifactValueId, Digest, TargetPayload, TargetPayloadFormat, TargetProfile,
};

use super::BackendError;

/// Immutable identity of one acquired execution device generation.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct DeviceIdentity {
    /// Stable registered backend identifier.
    pub backend: &'static str,
    /// Backend-local physical or logical device identifier.
    pub device: String,
    /// Monotonic generation that changes after device loss or reacquisition.
    pub generation: u64,
}

/// Acquired device identity, target compatibility, and health.
pub trait Device: Send + Sync {
    /// Immutable identity for this acquired generation.
    fn identity(&self) -> &DeviceIdentity;
    /// Exact target payload representation admitted by this device.
    fn target_format(&self) -> &TargetPayloadFormat;
    /// Exact immutable compilation profile admitted by this device.
    fn target_profile(&self) -> &TargetProfile;
    /// Whether new materialization and submission are currently allowed.
    fn is_healthy(&self) -> bool;
}

/// Host or resident bytes bound to one canonical artifact value.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum BoundResource {
    /// Caller-owned bytes uploaded for this submission.
    Host(Vec<u8>),
    /// Backend-resident resource handle.
    Resident(super::Resource),
}

/// Complete typed bindings for one immutable artifact instance.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BindingSet {
    artifact: Digest,
    resources: BTreeMap<ArtifactValueId, BoundResource>,
    invocation_grid: Option<[u32; 3]>,
}

impl BindingSet {
    /// Construct an empty binding set tied to one artifact identity.
    #[must_use]
    pub const fn new(artifact: Digest) -> Self {
        Self {
            artifact,
            resources: BTreeMap::new(),
            invocation_grid: None,
        }
    }

    /// Artifact identity these bindings are valid for.
    #[must_use]
    pub const fn artifact(&self) -> Digest {
        self.artifact
    }

    /// Bind or replace one canonical value.
    pub fn insert(&mut self, value: ArtifactValueId, resource: BoundResource) {
        self.resources.insert(value, resource);
    }

    /// Exact canonical value bindings.
    #[must_use]
    pub const fn resources(&self) -> &BTreeMap<ArtifactValueId, BoundResource> {
        &self.resources
    }

    /// Set the runtime invocation grid without changing immutable artifact identity.
    pub fn set_invocation_grid(&mut self, grid: [u32; 3]) -> Result<(), BackendError> {
        if let Some(axis) = grid.iter().position(|extent| *extent == 0) {
            return Err(BackendError::InvalidProgram {
                fix: format!(
                    "Fix: invocation grid axis {axis} must be positive, got {}.",
                    grid[axis]
                ),
            });
        }
        self.invocation_grid = Some(grid);
        Ok(())
    }

    /// Runtime grid override for this invocation.
    #[must_use]
    pub const fn invocation_grid(&self) -> Option<[u32; 3]> {
        self.invocation_grid
    }
}

/// Completed typed submission result.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Completion {
    /// Artifact identity executed by this submission.
    pub artifact: Digest,
    /// Canonical output values keyed by artifact ABI identity.
    pub outputs: BTreeMap<ArtifactValueId, Vec<u8>>,
    /// Updated retained values keyed by artifact ABI identity.
    pub retained: BTreeMap<ArtifactValueId, Vec<u8>>,
    /// Backend-measured device duration when available.
    pub device_ns: Option<u64>,
}

/// One in-flight submission against an immutable artifact instance.
pub trait Submission: Send + Sync {
    /// Non-blocking completion probe.
    fn is_ready(&self) -> bool;
    /// Wait for completion and typed readback.
    fn wait(self: Box<Self>) -> Result<Completion, BackendError>;
}

/// Device-native immutable executable and resource layout.
pub trait ArtifactInstance: Send + Sync {
    /// Neutral artifact identity implemented by this instance.
    fn artifact(&self) -> Digest;
    /// Exact payload identity materialized into this instance.
    fn payload(&self) -> Digest;
    /// Device generation that owns every native handle.
    fn device(&self) -> &DeviceIdentity;
    /// Validate bindings and submit one invocation.
    fn submit(&self, bindings: BindingSet) -> Result<Box<dyn Submission>, BackendError>;
}

/// Device-specific admission and native-handle construction.
pub trait ArtifactMaterializer: Send + Sync {
    /// Acquired target device.
    fn device(&self) -> &dyn Device;

    /// Allocate one resource owned by this materializer's device generation.
    fn allocate_resident(&self, _byte_len: usize) -> Result<super::Resource, BackendError> {
        Err(BackendError::UnsupportedFeature {
            name: "artifact resident buffer allocation".to_string(),
            backend: self.device().identity().backend.to_string(),
        })
    }

    /// Upload bytes into one resource owned by this materializer.
    fn upload_resident(
        &self,
        _resource: &super::Resource,
        _bytes: &[u8],
    ) -> Result<(), BackendError> {
        Err(BackendError::UnsupportedFeature {
            name: "artifact resident buffer upload".to_string(),
            backend: self.device().identity().backend.to_string(),
        })
    }

    /// Upload bytes at one aligned offset into a resident resource.
    fn upload_resident_at(
        &self,
        resource: &super::Resource,
        offset_bytes: usize,
        bytes: &[u8],
    ) -> Result<(), BackendError> {
        if offset_bytes == 0 {
            return self.upload_resident(resource, bytes);
        }
        Err(BackendError::UnsupportedFeature {
            name: "artifact resident ranged upload".to_string(),
            backend: self.device().identity().backend.to_string(),
        })
    }

    /// Release one resource owned by this materializer.
    fn free_resident(&self, _resource: super::Resource) -> Result<(), BackendError> {
        Err(BackendError::UnsupportedFeature {
            name: "artifact resident buffer free".to_string(),
            backend: self.device().identity().backend.to_string(),
        })
    }

    /// Materialize authenticated immutable target bytes.
    fn materialize(
        &self,
        artifact: &Artifact,
        payload: &TargetPayload,
    ) -> Result<Box<dyn ArtifactInstance>, BackendError>;
}