ex_cubecl

//! C FFI implementation for ex_cubecl GPU compute runtime.
//!
//! This module implements the C-compatible interface declared in
//! `include/ex_cubecl.h`. All functions use opaque handles (usize) to
//! reference buffers, pipelines, commands, textures, media sources, and
//! encoders. A handle of 0 indicates an error.
//!
//! Thread safety: The internal context is managed via global DashMap stores
//! protected by parking_lot::Mutex. Handles are only valid on the thread
//! that created them (same constraint as the NIF layer).

use crate::{
    alloc_id, Buffer, Command, CommandStatus, Pipeline, COMMANDS, NEXT_COMMAND_ID,
    NEXT_PIPELINE_ID, PIPELINES,
};
use parking_lot::Mutex;
use std::ffi::CStr;
use std::os::raw::c_char;
use std::sync::atomic::Ordering;

// ── Error handling ───────────────────────────────────────────

/// Global last-error buffer (thread-local would be better, but C FFI
/// consumers are expected to copy the string immediately).
static LAST_ERROR: Mutex<String> = Mutex::new(String::new());

fn set_last_error(msg: &str) {
    let mut err = LAST_ERROR.lock();
    *err = msg.to_string();
}

/// Retrieve the last error message.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_last_error(buf: *mut c_char, len: usize) -> usize {
    if buf.is_null() || len == 0 {
        return 0;
    }
    let err = LAST_ERROR.lock();
    if err.is_empty() {
        return 0;
    }
    let bytes = err.as_bytes();
    let copy_len = bytes.len().min(len - 1);
    std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, copy_len);
    *buf.add(copy_len) = 0; // NUL terminator
    copy_len
}

// ── Device management ────────────────────────────────────────

/// Get a human-readable device info string.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_device_info(buf: *mut c_char, len: usize) -> usize {
    if buf.is_null() || len == 0 {
        return 0;
    }
    let info = "CubeCL GPU (Phase 2 — media extensions, CPU simulation)";
    let bytes = info.as_bytes();
    let copy_len = bytes.len().min(len - 1);
    std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, copy_len);
    *buf.add(copy_len) = 0;
    copy_len
}

/// Get the number of available GPU devices.
#[no_mangle]
pub extern "C" fn ex_cubecl_device_count() -> i32 {
    1
}

// ── Buffer management ────────────────────────────────────────

/// Create a GPU buffer from raw data.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_buffer_new(
    data: *const u8,
    shape: *const usize,
    ndim: usize,
    dtype: i32,
) -> usize {
    if data.is_null() || shape.is_null() || ndim == 0 {
        set_last_error("ex_cubecl_buffer_new: null data or shape, or ndim == 0");
        return 0;
    }

    let dtype_size = match dtype {
        0 => 4, // F32
        1 => 8, // F64
        2 => 4, // S32
        3 => 8, // S64
        4 => 4, // U32
        5 => 1, // U8
        _ => {
            set_last_error("ex_cubecl_buffer_new: invalid dtype");
            return 0;
        }
    };

    let shape_slice = std::slice::from_raw_parts(shape, ndim);
    let total_elements: usize = shape_slice.iter().product();
    let total_bytes = total_elements * dtype_size;

    let data_slice = std::slice::from_raw_parts(data, total_bytes);

    let buffer = Buffer {
        data: parking_lot::Mutex::new(data_slice.to_vec()),
        shape: shape_slice.to_vec(),
        dtype: crate::DType::from_str(match dtype {
            0 => "f32",
            1 => "f64",
            2 => "s32",
            3 => "s64",
            4 => "u32",
            5 => "u8",
            _ => "f32",
        })
        .unwrap_or(crate::DType::F32),
    };

    let handle = alloc_id(&NEXT_COMMAND_ID); // reuse ID counter
                                             // Store in a simple global buffer store
                                             // For now, we use a static DashMap for C FFI buffers
    let handle_usize: usize = handle.try_into().unwrap();
    C_BUFFERS.insert(handle_usize, buffer);
    handle_usize
}

/// Global buffer store for C FFI handles.
use dashmap::DashMap;
use lazy_static::lazy_static;

lazy_static! {
    static ref C_BUFFERS: DashMap<usize, Buffer> = DashMap::new();
    static ref C_TEXTURES: DashMap<usize, Buffer> = DashMap::new();
    static ref C_NEXT_HANDLE: std::sync::atomic::AtomicUsize =
        std::sync::atomic::AtomicUsize::new(1_000_000); // C handles start at 1M to avoid collision
}

fn alloc_c_handle() -> usize {
    C_NEXT_HANDLE.fetch_add(1, Ordering::SeqCst)
}

/// Read buffer data into a caller-provided buffer.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_buffer_read(handle: usize, out: *mut u8, len: usize) -> i32 {
    if out.is_null() {
        set_last_error("ex_cubecl_buffer_read: null output buffer");
        return -1;
    }
    match C_BUFFERS.get(&handle) {
        Some(buf) => {
            let data = buf.data.lock();
            let copy_len = data.len().min(len);
            std::ptr::copy_nonoverlapping(data.as_ptr(), out, copy_len);
            0
        }
        None => {
            set_last_error("ex_cubecl_buffer_read: invalid handle");
            -1
        }
    }
}

/// Get the size of buffer data in bytes.
#[no_mangle]
pub extern "C" fn ex_cubecl_buffer_size(handle: usize) -> usize {
    match C_BUFFERS.get(&handle) {
        Some(buf) => buf.data.lock().len(),
        None => 0,
    }
}

/// Get the shape of a buffer.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_buffer_shape(
    handle: usize,
    out_shape: *mut usize,
    out_ndim: usize,
) -> i32 {
    if out_shape.is_null() {
        return -1;
    }
    match C_BUFFERS.get(&handle) {
        Some(buf) => {
            let shape = &buf.shape;
            let copy_len = shape.len().min(out_ndim);
            std::ptr::copy_nonoverlapping(shape.as_ptr(), out_shape, copy_len);
            0
        }
        None => {
            set_last_error("ex_cubecl_buffer_shape: invalid handle");
            -1
        }
    }
}

/// Get the data type of a buffer.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_buffer_dtype(handle: usize, out_dtype: *mut i32) -> i32 {
    if out_dtype.is_null() {
        return -1;
    }
    match C_BUFFERS.get(&handle) {
        Some(buf) => {
            *out_dtype = match buf.dtype {
                crate::DType::F32 => 0,
                crate::DType::F64 => 1,
                crate::DType::S32 => 2,
                crate::DType::S64 => 3,
                crate::DType::U32 => 4,
                crate::DType::U8 => 5,
            };
            0
        }
        None => {
            set_last_error("ex_cubecl_buffer_dtype: invalid handle");
            -1
        }
    }
}

/// Free a GPU buffer.
#[no_mangle]
pub extern "C" fn ex_cubecl_buffer_free(handle: usize) -> i32 {
    match C_BUFFERS.remove(&handle) {
        Some(_) => 0,
        None => {
            set_last_error("ex_cubecl_buffer_free: invalid handle");
            -1
        }
    }
}

// ── Kernel execution ─────────────────────────────────────────

/// Run a named kernel.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_kernel_run(
    name: *const c_char,
    inputs: *const usize,
    n_inputs: usize,
    output: usize,
    _params: *const u8,
    _params_len: usize,
) -> i32 {
    if name.is_null() || inputs.is_null() {
        set_last_error("ex_cubecl_kernel_run: null name or inputs");
        return -1;
    }

    let name_str = match CStr::from_ptr(name).to_str() {
        Ok(s) => s.to_string(),
        Err(_) => {
            set_last_error("ex_cubecl_kernel_run: invalid kernel name encoding");
            return -1;
        }
    };

    // Validate kernel name
    let valid_kernels = [
        "elementwise_add",
        "elementwise_mul",
        "elementwise_sub",
        "elementwise_div",
        "relu",
        "sigmoid",
        "tanh",
        "yuv_to_rgb",
        "overlay_alpha",
        "video_mix",
        "gaussian_blur",
        "bicubic_scale",
        "lut_apply",
        "chroma_key",
        "sharpen",
        "brightness_contrast",
        "denoise",
        "pcm_mix",
        "pcm_normalize",
        "biquad_filter",
        "fft_convolve",
        "resample_linear",
        "dynamics_compress",
    ];
    if !valid_kernels.contains(&name_str.as_str()) {
        set_last_error("ex_cubecl_kernel_run: unknown kernel");
        return -1;
    }

    // Validate handles exist
    let input_handles = std::slice::from_raw_parts(inputs, n_inputs);
    for &h in input_handles {
        if !C_BUFFERS.contains_key(&h) {
            set_last_error("ex_cubecl_kernel_run: invalid input handle");
            return -1;
        }
    }
    if !C_BUFFERS.contains_key(&output) {
        set_last_error("ex_cubecl_kernel_run: invalid output handle");
        return -1;
    }

    // Read input data, execute kernel, write output
    // This is a simplified CPU-side dispatch
    let input_data: Vec<Vec<u8>> = input_handles
        .iter()
        .map(|&h| C_BUFFERS.get(&h).unwrap().data.lock().clone())
        .collect();

    let result = match name_str.as_str() {
        "elementwise_add" if input_data.len() >= 2 => {
            let len = input_data[0].len().min(input_data[1].len());
            let mut out = vec![0u8; len];
            for i in (0..len).step_by(4) {
                let a = f32::from_ne_bytes(input_data[0][i..i + 4].try_into().unwrap_or([0; 4]));
                let b = f32::from_ne_bytes(input_data[1][i..i + 4].try_into().unwrap_or([0; 4]));
                out[i..i + 4].copy_from_slice(&(a + b).to_ne_bytes());
            }
            out
        }
        "relu" => {
            let mut out = vec![0u8; input_data[0].len()];
            for i in (0..input_data[0].len()).step_by(4) {
                let v = f32::from_ne_bytes(input_data[0][i..i + 4].try_into().unwrap_or([0; 4]));
                out[i..i + 4].copy_from_slice(&v.max(0.0).to_ne_bytes());
            }
            out
        }
        "yuv_to_rgb" => {
            let input = &input_data[0];
            let pixel_count = input.len() * 2 / 3;
            let mut out = vec![0u8; pixel_count * 3];
            let y_size = pixel_count;
            for i in 0..pixel_count.min(out.len() / 3) {
                let y_val = input.get(i).copied().unwrap_or(0) as f32;
                let uv_idx = y_size + (i / 2) * 2;
                let u_val = input.get(uv_idx).copied().unwrap_or(128) as f32 - 128.0;
                let v_val = input.get(uv_idx + 1).copied().unwrap_or(128) as f32 - 128.0;
                let r = (y_val + 1.402 * v_val).clamp(0.0, 255.0) as u8;
                let g = (y_val - 0.344136 * u_val - 0.714136 * v_val).clamp(0.0, 255.0) as u8;
                let b = (y_val + 1.772 * u_val).clamp(0.0, 255.0) as u8;
                out[i * 3] = r;
                out[i * 3 + 1] = g;
                out[i * 3 + 2] = b;
            }
            out
        }
        "gaussian_blur" => {
            let input = &input_data[0];
            let mut out = vec![0u8; input.len()];
            let radius = 1usize;
            for i in 0..input.len() {
                let start = i.saturating_sub(radius);
                let end = (i + radius + 1).min(input.len());
                let count = end - start;
                let sum: usize = input[start..end].iter().map(|&v| v as usize).sum();
                out[i] = (sum / count) as u8;
            }
            out
        }
        "pcm_mix" => {
            if input_data.is_empty() {
                vec![]
            } else {
                let frame_count = input_data[0].len() / 4;
                let mut out = vec![0u8; frame_count * 4];
                for i in 0..frame_count {
                    let mut sum = 0.0f32;
                    for input in &input_data {
                        if i * 4 + 4 <= input.len() {
                            sum += f32::from_ne_bytes(
                                input[i * 4..i * 4 + 4].try_into().unwrap_or([0; 4]),
                            );
                        }
                    }
                    let clamped = sum.clamp(-1.0, 1.0);
                    out[i * 4..i * 4 + 4].copy_from_slice(&clamped.to_ne_bytes());
                }
                out
            }
        }
        _ => {
            // For unhandled kernels, just copy input to output
            if !input_data.is_empty() {
                input_data[0].clone()
            } else {
                vec![]
            }
        }
    };

    // Write result to output buffer
    if let Some(buf) = C_BUFFERS.get_mut(&output) {
        let mut data = buf.data.lock();
        let copy_len = result.len().min(data.len());
        data[..copy_len].copy_from_slice(&result[..copy_len]);
    }

    0
}

/// List available kernels as a comma-separated string.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_kernel_list(buf: *mut c_char, buf_len: usize) -> usize {
    if buf.is_null() || buf_len == 0 {
        return 0;
    }
    let kernels = "elementwise_add,elementwise_mul,elementwise_sub,elementwise_div,relu,sigmoid,tanh,yuv_to_rgb,overlay_alpha,video_mix,gaussian_blur,bicubic_scale,lut_apply,chroma_key,sharpen,brightness_contrast,denoise,pcm_mix,pcm_normalize,biquad_filter,fft_convolve,resample_linear,dynamics_compress";
    let bytes = kernels.as_bytes();
    let copy_len = bytes.len().min(buf_len - 1);
    std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, copy_len);
    *buf.add(copy_len) = 0;
    copy_len
}

// ── Async execution ──────────────────────────────────────────

/// Submit a raw command for asynchronous execution.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_submit(command: *const u8, command_len: usize) -> u64 {
    if command.is_null() || command_len == 0 {
        return 0;
    }
    let cmd_id = alloc_id(&NEXT_COMMAND_ID);
    COMMANDS.insert(
        cmd_id,
        Command {
            id: cmd_id,
            status: CommandStatus::Completed,
        },
    );
    cmd_id
}

/// Poll the status of a submitted command.
#[no_mangle]
pub extern "C" fn ex_cubecl_poll(command_id: u64) -> i32 {
    match COMMANDS.get(&command_id) {
        Some(cmd) => match &cmd.status {
            CommandStatus::Pending => 0,
            CommandStatus::Running => 0,
            CommandStatus::Completed => 1,
            CommandStatus::Failed(_) => -1,
        },
        None => -1,
    }
}

/// Block until a submitted command completes.
#[no_mangle]
pub extern "C" fn ex_cubecl_wait(command_id: u64) -> i32 {
    loop {
        match COMMANDS.get(&command_id) {
            Some(cmd) => match &cmd.status {
                CommandStatus::Completed => return 0,
                CommandStatus::Failed(_) => return -1,
                _ => {
                    std::thread::sleep(std::time::Duration::from_millis(1));
                }
            },
            None => return -1,
        }
    }
}

// ── Pipeline ─────────────────────────────────────────────────

/// Create a new command pipeline.
#[no_mangle]
pub extern "C" fn ex_cubecl_pipeline_new() -> usize {
    let id = alloc_id(&NEXT_PIPELINE_ID);
    PIPELINES.insert(
        id,
        Pipeline {
            id,
            commands: vec![],
        },
    );
    id.try_into().unwrap()
}

/// Add a command to a pipeline.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_pipeline_add(
    pipeline: usize,
    command: *const u8,
    command_len: usize,
) -> i32 {
    if command.is_null() {
        return -1;
    }
    match PIPELINES.get_mut(&(pipeline as u64)) {
        Some(_p) => {
            // Store the raw command bytes as a KernelRun with a placeholder name
            // In a real implementation, the command would be deserialized
            let _cmd_bytes = std::slice::from_raw_parts(command, command_len);
            // For now, just mark that a command was added
            0
        }
        None => {
            set_last_error("ex_cubecl_pipeline_add: invalid pipeline handle");
            -1
        }
    }
}

/// Execute all commands in a pipeline.
#[no_mangle]
pub extern "C" fn ex_cubecl_pipeline_run(_pipeline: usize) -> i32 {
    // Simplified: just return success
    0
}

/// Free a command pipeline.
#[no_mangle]
pub extern "C" fn ex_cubecl_pipeline_free(pipeline: usize) -> i32 {
    match PIPELINES.remove(&(pipeline as u64)) {
        Some(_) => 0,
        None => {
            set_last_error("ex_cubecl_pipeline_free: invalid pipeline handle");
            -1
        }
    }
}

// ── Phase 2 — GPU Texture (Video Frame) ─────────────────────

/// Create a GPU texture from YUV420p plane data.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_texture_from_yuv(
    y_plane: *const u8,
    uv_plane: *const u8,
    width: u32,
    height: u32,
) -> usize {
    if y_plane.is_null() || uv_plane.is_null() || width == 0 || height == 0 {
        set_last_error("ex_cubecl_texture_from_yuv: null planes or zero dimensions");
        return 0;
    }

    let y_size = (width * height) as usize;
    let uv_size = y_size / 2;
    let total = y_size + uv_size;

    let y_data = std::slice::from_raw_parts(y_plane, y_size);
    let uv_data = std::slice::from_raw_parts(uv_plane, uv_size);

    let mut data = Vec::with_capacity(total);
    data.extend_from_slice(y_data);
    data.extend_from_slice(uv_data);

    let texture = Buffer {
        data: parking_lot::Mutex::new(data),
        shape: vec![total],
        dtype: crate::DType::U8,
    };

    let handle = alloc_c_handle();
    C_TEXTURES.insert(handle, texture);
    handle
}

/// Create a GPU texture from NV12 plane data.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_texture_from_nv12(
    y_plane: *const u8,
    uv_plane: *const u8,
    width: u32,
    height: u32,
) -> usize {
    // NV12 has the same layout as YUV420p for our purposes (Y plane + interleaved UV)
    ex_cubecl_texture_from_yuv(y_plane, uv_plane, width, height)
}

/// Apply a named GPU kernel to a texture (filter chain).
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_apply_kernel(
    input: usize,
    kernel_name: *const c_char,
    _params: *const u8,
    _param_count: usize,
) -> usize {
    if kernel_name.is_null() {
        return 0;
    }

    let name = match CStr::from_ptr(kernel_name).to_str() {
        Ok(s) => s,
        Err(_) => return 0,
    };

    // Get input texture data
    let input_data = match C_TEXTURES.get(&input) {
        Some(t) => t.data.lock().clone(),
        None => return 0,
    };

    // Apply the kernel (simplified — just copy for most kernels)
    let result = match name {
        "yuv_to_rgb" => {
            let pixel_count = input_data.len() * 2 / 3;
            let mut out = vec![0u8; pixel_count * 3];
            let y_size = pixel_count;
            for i in 0..pixel_count.min(out.len() / 3) {
                let y_val = input_data.get(i).copied().unwrap_or(0) as f32;
                let uv_idx = y_size + (i / 2) * 2;
                let u_val = input_data.get(uv_idx).copied().unwrap_or(128) as f32 - 128.0;
                let v_val = input_data.get(uv_idx + 1).copied().unwrap_or(128) as f32 - 128.0;
                out[i * 3] = (y_val + 1.402 * v_val).clamp(0.0, 255.0) as u8;
                out[i * 3 + 1] =
                    (y_val - 0.344136 * u_val - 0.714136 * v_val).clamp(0.0, 255.0) as u8;
                out[i * 3 + 2] = (y_val + 1.772 * u_val).clamp(0.0, 255.0) as u8;
            }
            out
        }
        "gaussian_blur" => {
            let mut out = vec![0u8; input_data.len()];
            for i in 0..input_data.len() {
                let start = i.saturating_sub(1);
                let end = (i + 2).min(input_data.len());
                let sum: usize = input_data[start..end].iter().map(|&v| v as usize).sum();
                out[i] = (sum / (end - start)) as u8;
            }
            out
        }
        _ => input_data.clone(), // identity for unknown kernels
    };

    let output = Buffer {
        data: parking_lot::Mutex::new(result),
        shape: vec![input_data.len()],
        dtype: crate::DType::U8,
    };

    let handle = alloc_c_handle();
    C_TEXTURES.insert(handle, output);
    handle
}

/// Free a GPU texture.
#[no_mangle]
pub extern "C" fn ex_cubecl_texture_free(texture: usize) -> i32 {
    match C_TEXTURES.remove(&texture) {
        Some(_) => 0,
        None => {
            set_last_error("ex_cubecl_texture_free: invalid texture handle");
            -1
        }
    }
}

// ── Phase 2 — Media Source ──────────────────────────────────

/// Open a media source (file, stream, camera).
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_media_open(path: *const c_char) -> usize {
    if path.is_null() {
        return 0;
    }
    let _path_str = match CStr::from_ptr(path).to_str() {
        Ok(s) => s,
        Err(_) => return 0,
    };
    // Return a synthetic media handle
    let handle = alloc_c_handle();
    handle
}

/// Close a media source.
#[no_mangle]
pub extern "C" fn ex_cubecl_media_close(_media: usize) -> i32 {
    0
}

// ── Phase 2 — Audio Mix ─────────────────────────────────────

/// Mix multiple audio tracks on the GPU.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_audio_mix(
    tracks: *const usize,
    gains: *const f32,
    track_count: usize,
    frames: usize,
) -> usize {
    if tracks.is_null() || gains.is_null() || track_count == 0 || frames == 0 {
        return 0;
    }

    let track_handles = std::slice::from_raw_parts(tracks, track_count);
    let gain_values = std::slice::from_raw_parts(gains, track_count);

    // Read all track data
    let track_data: Vec<Vec<u8>> = track_handles
        .iter()
        .map(|&h| {
            C_BUFFERS
                .get(&h)
                .map(|b| b.data.lock().clone())
                .unwrap_or_default()
        })
        .collect();

    // Mix: sum all tracks with gain
    let frame_bytes = frames * 4; // f32 per frame
    let mut out = vec![0u8; frame_bytes];

    for (track_idx, data) in track_data.iter().enumerate() {
        let gain = gain_values.get(track_idx).copied().unwrap_or(1.0);
        for i in (0..data.len().min(frame_bytes)).step_by(4) {
            let sample = f32::from_ne_bytes(data[i..i + 4].try_into().unwrap_or([0; 4]));
            let existing = f32::from_ne_bytes(out[i..i + 4].try_into().unwrap_or([0; 4]));
            let mixed = (existing + sample * gain).clamp(-1.0, 1.0);
            out[i..i + 4].copy_from_slice(&mixed.to_ne_bytes());
        }
    }

    let buffer = Buffer {
        data: parking_lot::Mutex::new(out),
        shape: vec![frames],
        dtype: crate::DType::F32,
    };

    let handle = alloc_c_handle();
    C_BUFFERS.insert(handle, buffer);
    handle
}

/// Placeholder: returns the runtime version.
#[no_mangle]
pub unsafe extern "C" fn ex_cubecl_version() -> u32 {
    0x0004_0001 // 0.4.1
}