cgpu 0.1.0

use std::ops::Range;

use crate::budget::GpuMemoryBudget;
use crate::error::GpuError;
use crate::telemetry::{measure_phase, GpuPhase, GpuTelemetry};

pub trait WorkItem: Send + Sync {
    fn bytes_in(&self) -> usize;

    fn bytes_out(&self) -> usize;

    fn estimated_total_bytes(&self) -> usize {
        self.bytes_in().saturating_add(self.bytes_out())
    }
}

pub trait ByteSized: WorkItem {}

impl<T: WorkItem + ?Sized> ByteSized for T {}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct BatchSpan {
    pub index: usize,
    pub start: usize,
    pub end: usize,
    pub bytes_in: usize,
    pub bytes_out: usize,
    pub estimated_bytes: usize,
}

impl BatchSpan {
    pub fn new(index: usize, start: usize, end: usize, bytes_in: usize, bytes_out: usize) -> Self {
        Self {
            index,
            start,
            end,
            bytes_in,
            bytes_out,
            estimated_bytes: bytes_in.saturating_add(bytes_out),
        }
    }

    pub fn len(&self) -> usize {
        self.end.saturating_sub(self.start)
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn range(&self) -> Range<usize> {
        self.start..self.end
    }
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct GpuBatchPlan {
    spans: Vec<BatchSpan>,
    bytes_in: usize,
    bytes_out: usize,
    estimated_bytes: usize,
}

impl GpuBatchPlan {
    pub fn empty() -> Self {
        Self::default()
    }

    pub fn spans(&self) -> &[BatchSpan] {
        &self.spans
    }

    pub fn into_spans(self) -> Vec<BatchSpan> {
        self.spans
    }

    pub fn batch_count(&self) -> usize {
        self.spans.len()
    }

    pub fn bytes_in(&self) -> usize {
        self.bytes_in
    }

    pub fn bytes_out(&self) -> usize {
        self.bytes_out
    }

    pub fn estimated_bytes(&self) -> usize {
        self.estimated_bytes
    }

    pub fn validate_for_len(&self, item_len: usize) -> Result<(), GpuError> {
        let mut cursor = 0;

        for span in &self.spans {
            if span.start != cursor {
                return Err(GpuError::InvalidBatchPlan("batch spans must be contiguous"));
            }

            if span.end < span.start {
                return Err(GpuError::InvalidBatchPlan("batch span end precedes start"));
            }

            if span.end > item_len {
                return Err(GpuError::InvalidBatchPlan("batch span exceeds item length"));
            }

            if span.is_empty() {
                return Err(GpuError::InvalidBatchPlan("batch spans cannot be empty"));
            }

            cursor = span.end;
        }

        if cursor != item_len {
            return Err(GpuError::InvalidBatchPlan(
                "batch spans do not cover every item",
            ));
        }

        Ok(())
    }

    pub(crate) fn push_span(
        &mut self,
        start: usize,
        end: usize,
        bytes_in: usize,
        bytes_out: usize,
        estimated_bytes: usize,
    ) {
        self.spans.push(BatchSpan {
            index: self.spans.len(),
            start,
            end,
            bytes_in,
            bytes_out,
            estimated_bytes,
        });
        self.bytes_in = self.bytes_in.saturating_add(bytes_in);
        self.bytes_out = self.bytes_out.saturating_add(bytes_out);
        self.estimated_bytes = self.estimated_bytes.saturating_add(estimated_bytes);
    }
}

pub fn pack_by_bytes<T: WorkItem>(
    items: &[T],
    budget: GpuMemoryBudget,
    telemetry: &mut GpuTelemetry,
) -> Result<GpuBatchPlan, GpuError> {
    let plan = measure_phase(telemetry, GpuPhase::Pack, || {
        pack_by_bytes_inner(items, budget)
    })?;
    telemetry.record_plan(&plan);
    Ok(plan)
}

fn pack_by_bytes_inner<T: WorkItem>(
    items: &[T],
    budget: GpuMemoryBudget,
) -> Result<GpuBatchPlan, GpuError> {
    if items.is_empty() {
        return Ok(GpuBatchPlan::empty());
    }

    let target_batch_bytes = budget.target_batch_bytes();
    let safe_vram_bytes = budget.safe_vram_bytes();
    let mut plan = GpuBatchPlan::empty();

    let mut batch_start = 0;
    let mut batch_bytes_in = 0_usize;
    let mut batch_bytes_out = 0_usize;
    let mut batch_estimated_bytes = 0_usize;

    for (item_index, item) in items.iter().enumerate() {
        let item_bytes_in = item.bytes_in();
        let item_bytes_out = item.bytes_out();
        let item_estimated_bytes = item.estimated_total_bytes();

        if item_estimated_bytes > safe_vram_bytes {
            return Err(GpuError::ItemExceedsSafeBudget {
                item_index,
                item_bytes: item_estimated_bytes,
                safe_vram_bytes,
            });
        }

        let would_exceed_target = batch_estimated_bytes > 0
            && batch_estimated_bytes.saturating_add(item_estimated_bytes) > target_batch_bytes;

        if would_exceed_target {
            plan.push_span(
                batch_start,
                item_index,
                batch_bytes_in,
                batch_bytes_out,
                batch_estimated_bytes,
            );
            batch_start = item_index;
            batch_bytes_in = 0;
            batch_bytes_out = 0;
            batch_estimated_bytes = 0;
        }

        batch_bytes_in = batch_bytes_in.saturating_add(item_bytes_in);
        batch_bytes_out = batch_bytes_out.saturating_add(item_bytes_out);
        batch_estimated_bytes = batch_estimated_bytes.saturating_add(item_estimated_bytes);
    }

    if batch_estimated_bytes > 0 {
        plan.push_span(
            batch_start,
            items.len(),
            batch_bytes_in,
            batch_bytes_out,
            batch_estimated_bytes,
        );
    }

    Ok(plan)
}