cgpu 0.1.0

use std::error::Error;
use std::fmt::{self, Display, Formatter};

use crate::batch::{BatchSpan, WorkItem};

pub trait GpuJob: Send + Sync {
    type Item: WorkItem;
    type Output: Send;

    fn label(&self) -> &'static str;

    fn items(&self) -> &[Self::Item];

    fn encode_batch(&self, span: &BatchSpan) -> Result<EncodedBatch, JobError>;

    fn decode_batch(&self, bytes: &[u8], span: &BatchSpan) -> Result<Vec<Self::Output>, JobError>;

    fn cpu_fallback(&self, _span: &BatchSpan) -> Result<Vec<Self::Output>, JobError> {
        Err(JobError::FallbackNotImplemented)
    }
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct EncodedBatch {
    pub label: Option<String>,
    pub shader: Option<ComputeShader>,
    pub buffers: Vec<EncodedBuffer>,
    pub metadata: Vec<u8>,
    pub dispatch: DispatchSize,
}

impl EncodedBatch {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_label(mut self, label: impl Into<String>) -> Self {
        self.label = Some(label.into());
        self
    }

    pub fn with_shader(mut self, shader: ComputeShader) -> Self {
        self.shader = Some(shader);
        self
    }

    pub fn with_wgsl(mut self, label: impl Into<String>, source: impl Into<String>) -> Self {
        self.shader = Some(ComputeShader::new(label, source));
        self
    }

    pub fn with_metadata(mut self, metadata: impl Into<Vec<u8>>) -> Self {
        self.metadata = metadata.into();
        self
    }

    pub fn with_dispatch(mut self, x: u32, y: u32, z: u32) -> Self {
        self.dispatch = DispatchSize { x, y, z };
        self
    }

    pub fn add_buffer(mut self, buffer: EncodedBuffer) -> Self {
        self.buffers.push(buffer);
        self
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ComputeShader {
    pub label: String,
    pub source: String,
    pub entry_point: String,
}

impl ComputeShader {
    pub fn new(label: impl Into<String>, source: impl Into<String>) -> Self {
        Self {
            label: label.into(),
            source: source.into(),
            entry_point: "main".to_string(),
        }
    }

    pub fn with_entry_point(mut self, entry_point: impl Into<String>) -> Self {
        self.entry_point = entry_point.into();
        self
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct EncodedBuffer {
    pub binding: u32,
    pub role: BufferRole,
    pub bytes: Vec<u8>,
    pub label: Option<String>,
    pub read_only: bool,
    pub readback: bool,
}

impl EncodedBuffer {
    pub fn new(role: BufferRole, bytes: impl Into<Vec<u8>>) -> Self {
        Self {
            binding: 0,
            role,
            bytes: bytes.into(),
            label: None,
            read_only: matches!(role, BufferRole::Input | BufferRole::Metadata),
            readback: matches!(role, BufferRole::Output),
        }
    }

    pub fn binding(mut self, binding: u32) -> Self {
        self.binding = binding;
        self
    }

    pub fn with_label(mut self, label: impl Into<String>) -> Self {
        self.label = Some(label.into());
        self
    }

    pub fn read_only(mut self, read_only: bool) -> Self {
        self.read_only = read_only;
        self
    }

    pub fn readback(mut self, readback: bool) -> Self {
        self.readback = readback;
        self
    }

    pub fn storage_read_only(binding: u32, bytes: impl Into<Vec<u8>>) -> Self {
        Self::new(BufferRole::Input, bytes)
            .binding(binding)
            .read_only(true)
            .readback(false)
    }

    pub fn storage_read_write(binding: u32, role: BufferRole, bytes: impl Into<Vec<u8>>) -> Self {
        Self::new(role, bytes).binding(binding).read_only(false)
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct DispatchSize {
    pub x: u32,
    pub y: u32,
    pub z: u32,
}

impl Default for DispatchSize {
    fn default() -> Self {
        Self { x: 1, y: 1, z: 1 }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BufferRole {
    Input,
    Output,
    Scratch,
    Metadata,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum JobError {
    EncodingFailed(String),
    DecodingFailed(String),
    FallbackNotImplemented,
    InvalidSpan { reason: String },
}

impl Display for JobError {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::EncodingFailed(message) => write!(f, "failed to encode batch: {message}"),
            Self::DecodingFailed(message) => write!(f, "failed to decode batch: {message}"),
            Self::FallbackNotImplemented => f.write_str("CPU fallback is not implemented"),
            Self::InvalidSpan { reason } => write!(f, "invalid batch span: {reason}"),
        }
    }
}

impl Error for JobError {}