cgpu 0.1.0

use cgpu::*;

#[derive(Clone)]
struct Item {
    input: usize,
    output: usize,
}

impl WorkItem for Item {
    fn bytes_in(&self) -> usize {
        self.input
    }

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

struct CpuJob {
    items: Vec<Item>,
}

impl GpuJob for CpuJob {
    type Item = Item;
    type Output = usize;

    fn label(&self) -> &'static str {
        "cpu-job"
    }

    fn items(&self) -> &[Self::Item] {
        &self.items
    }

    fn encode_batch(&self, _span: &BatchSpan) -> Result<EncodedBatch, JobError> {
        Ok(EncodedBatch::new())
    }

    fn decode_batch(
        &self,
        _bytes: &[u8],
        _span: &BatchSpan,
    ) -> Result<Vec<Self::Output>, JobError> {
        Ok(Vec::new())
    }

    fn cpu_fallback(&self, span: &BatchSpan) -> Result<Vec<Self::Output>, JobError> {
        Ok(self.items[span.range()]
            .iter()
            .map(WorkItem::estimated_total_bytes)
            .collect())
    }
}

#[cfg(feature = "gpu")]
struct DoubleGpuJob {
    items: Vec<Item>,
}

#[cfg(feature = "gpu")]
impl GpuJob for DoubleGpuJob {
    type Item = Item;
    type Output = u32;

    fn label(&self) -> &'static str {
        "double-gpu-job"
    }

    fn items(&self) -> &[Self::Item] {
        &self.items
    }

    fn encode_batch(&self, _span: &BatchSpan) -> Result<EncodedBatch, JobError> {
        let input = 21_u32.to_ne_bytes();
        let output = 0_u32.to_ne_bytes();
        Ok(EncodedBatch::new()
            .with_label("double-gpu-test")
            .with_wgsl(
                "double-gpu-test-shader",
                r#"
@group(0) @binding(0)
var<storage, read> input: array<u32>;

@group(0) @binding(1)
var<storage, read_write> output: array<u32>;

@compute @workgroup_size(1, 1, 1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    if (global_id.x == 0u) {
        output[0] = input[0] * 2u;
    }
}
"#,
            )
            .add_buffer(EncodedBuffer::storage_read_only(0, input))
            .add_buffer(
                EncodedBuffer::storage_read_write(1, BufferRole::Output, output).readback(true),
            )
            .with_dispatch(1, 1, 1))
    }

    fn decode_batch(&self, bytes: &[u8], _span: &BatchSpan) -> Result<Vec<Self::Output>, JobError> {
        Ok(vec![read_u32(bytes, 0).map_err(|error| {
            JobError::DecodingFailed(error.to_string())
        })?])
    }

    fn cpu_fallback(&self, _span: &BatchSpan) -> Result<Vec<Self::Output>, JobError> {
        Ok(vec![42])
    }
}

#[test]
fn config_defaults_match_contract() {
    let config = ExecutorConfig::default();

    assert_eq!(config.memory_fill_ratio, 0.90);
    assert_eq!(config.min_batch_bytes, 4096);
    assert_eq!(config.max_batch_bytes, 256 * 1024 * 1024);
    assert_eq!(config.execution_mode, ExecutionMode::Auto);
    assert!(config.parallel_fallback);
    assert!(config.shader_cache);
    assert!(config.enable_telemetry);
}

#[test]
fn config_loads_from_json_file() {
    let config = ExecutorConfig::from_default_json_file().unwrap();

    assert_eq!(config.memory_fill_ratio, 0.90);
    assert_eq!(config.execution_mode, ExecutionMode::Auto);
    assert_eq!(config.shader_optimization, ShaderOpt::Performance);
}

#[test]
fn executor_runs_cpu_fallback_and_reports_batches() {
    let mut config = ExecutorConfig {
        execution_mode: ExecutionMode::PreferCpu,
        min_batch_bytes: 1,
        max_batch_bytes: 100,
        memory_fill_ratio: 0.90,
        ..ExecutorConfig::default()
    };
    config.vram_override = Some(100);

    let mut executor = GpuExecutor::with_config(config).unwrap();
    let job = CpuJob {
        items: vec![
            Item {
                input: 30,
                output: 5,
            },
            Item {
                input: 30,
                output: 5,
            },
            Item {
                input: 30,
                output: 5,
            },
        ],
    };

    let report = executor.execute(&job).unwrap();

    assert_eq!(report.outputs, vec![35, 35, 35]);
    assert_eq!(report.execution_path, ExecutionPath::CpuOnly);
    assert_eq!(report.batches.len(), 2);
    assert_eq!(report.batches[0].path, BatchPath::CpuFallback);
    assert!(executor.telemetry().latest().is_some());
}

#[cfg(feature = "gpu")]
#[test]
fn executor_runs_encoded_gpu_batch_when_gpu_available() {
    let mut executor = match GpuExecutor::with_config(ExecutorConfig {
        execution_mode: ExecutionMode::GpuOnly,
        vram_override: Some(16 * 1024 * 1024),
        ..ExecutorConfig::default()
    }) {
        Ok(executor) => executor,
        Err(ExecutorInitError::NoGpuAvailable) => return,
        Err(error) => panic!("unexpected executor init error: {error}"),
    };
    let job = DoubleGpuJob {
        items: vec![Item {
            input: 4,
            output: 4,
        }],
    };

    let report = executor.execute(&job).unwrap();

    assert_eq!(report.outputs, vec![42]);
    assert_eq!(report.execution_path, ExecutionPath::GpuOnly);
    assert_eq!(report.batches[0].path, BatchPath::Gpu);
}

#[cfg(not(feature = "gpu"))]
#[test]
fn gpu_only_requires_gpu_feature() {
    let error = GpuExecutor::with_config(ExecutorConfig {
        execution_mode: ExecutionMode::GpuOnly,
        ..ExecutorConfig::default()
    })
    .unwrap_err();

    assert!(matches!(error, ExecutorInitError::NoGpuAvailable));
}