use super::*;

#[repr(C)]
pub(super) struct GpuTensor {
    _private: [u8; 0],
}

#[repr(C)]
pub(super) struct StreamExpertTable {
    pub(super) model_map: *const c_void,
    pub(super) model_size: u64,
    pub(super) layer: u32,
    pub(super) total_experts: u32,
    pub(super) gate_offset: u64,
    pub(super) up_offset: u64,
    pub(super) down_offset: u64,
    pub(super) gate_expert_bytes: u64,
    pub(super) down_expert_bytes: u64,
}

#[derive(Clone, Copy, Default)]
#[repr(C)]
pub(super) struct StreamExpertCacheStats {
    pub(super) configured_count: u32,
    pub(super) current_count: u32,
    pub(super) hits: u64,
    pub(super) misses: u64,
    pub(super) evictions: u64,
    pub(super) wraps: u64,
    pub(super) buffer_allocs: u64,
    pub(super) buffer_reuses: u64,
    pub(super) evict_advise_bytes: u64,
    pub(super) willneed_advise_bytes: u64,
    pub(super) pread_bytes: u64,
    pub(super) pread_ms: f64,
}

unsafe extern "C" {
    pub(super) fn ds4_gpu_init() -> i32;
    pub(super) fn ds4_gpu_cleanup();
    pub(super) fn ds4_gpu_set_model_map_range(
        model_map: *const c_void,
        model_size: u64,
        map_offset: u64,
        map_size: u64,
        max_tensor_bytes: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_set_model_map_spans(
        model_map: *const c_void,
        model_size: u64,
        offsets: *const u64,
        sizes: *const u64,
        count: u32,
        max_tensor_bytes: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_set_quality(quality: bool);
    pub(super) fn ds4_gpu_set_glm_model(enabled: bool);
    pub(super) fn ds4_gpu_set_ssd_streaming(enabled: bool);
    pub(super) fn ds4_gpu_set_model_fd(fd: i32) -> i32;
    pub(super) fn ds4_gpu_set_streaming_expert_cache_budget(experts: u32);
    pub(super) fn ds4_gpu_set_streaming_expert_cache_expert_bytes(bytes: u64);
    pub(super) fn ds4_gpu_recommended_working_set_size() -> u64;
    pub(super) fn ds4_gpu_stream_expert_cache_get_stats(stats: *mut StreamExpertCacheStats);
    pub(super) fn ds4_gpu_stream_expert_cache_seed_experts(
        table: *const StreamExpertTable,
        expert_ids: *const i32,
        expert_priorities: *const u32,
        experts: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_stream_expert_cache_begin_selected_load(
        table: *const StreamExpertTable,
        selected_ids: *const i32,
        selected: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_stream_expert_cache_note_service_thread();
    pub(super) fn ds4_gpu_stream_expert_cache_reset_route_hotness();
    pub(super) fn ds4_gpu_stream_prefill_batch_selected_addr_enabled(
        n_tokens: u32,
        n_total_expert: u32,
        n_expert: u32,
        gate_type: u32,
        down_type: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
        table: *const StreamExpertTable,
        selected: *const GpuTensor,
        count: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_flush_commands() -> i32;
    pub(super) fn ds4_gpu_tensor_alloc(bytes: u64) -> *mut GpuTensor;
    pub(super) fn ds4_gpu_tensor_view(
        base: *const GpuTensor,
        offset: u64,
        bytes: u64,
    ) -> *mut GpuTensor;
    pub(super) fn ds4_gpu_tensor_free(tensor: *mut GpuTensor);
    pub(super) fn ds4_gpu_tensor_fill_f32(tensor: *mut GpuTensor, value: f32, count: u64) -> i32;
    pub(super) fn ds4_gpu_tensor_read(
        tensor: *const GpuTensor,
        offset: u64,
        data: *mut c_void,
        bytes: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_tensor_write(
        tensor: *mut GpuTensor,
        offset: u64,
        data: *const c_void,
        bytes: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_tensor_copy(
        dst: *mut GpuTensor,
        dst_offset: u64,
        src: *const GpuTensor,
        src_offset: u64,
        bytes: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_tensor_copy_f32_to_f16(
        dst: *mut GpuTensor,
        dst_offset: u64,
        src: *const GpuTensor,
        src_offset: u64,
        count: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_pack_slot_rows_f32_tensor(
        out: *mut GpuTensor,
        slots: *const GpuTensor,
        rows: u32,
        width: u32,
        slot_count: u32,
        slot_stride: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_begin_commands() -> i32;
    pub(super) fn ds4_gpu_end_commands() -> i32;
    pub(super) fn ds4_gpu_signal_selected_readback_ready(event: *mut u64) -> i32;
    pub(super) fn ds4_gpu_wait_selected_readback_ready(
        event: u64,
        label: *const std::ffi::c_char,
    ) -> i32;
    pub(super) fn ds4_gpu_routed_moe_set_selected_override(selected: *const i32, count: u32)
    -> i32;

    pub(super) fn ds4_gpu_embed_tokens_hc_tensor(
        out: *mut GpuTensor,
        tokens: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        vocab: u32,
        rows: u32,
        embd: u32,
        hc: u32,
    ) -> i32;

    pub(super) fn ds4_gpu_embed_token_hc_tensor(
        out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        vocab: u32,
        token: u32,
        embd: u32,
        hc: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_rms_norm_plain_tensor(
        out: *mut GpuTensor,
        x: *const GpuTensor,
        n: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_rms_norm_plain_rows_tensor(
        out: *mut GpuTensor,
        x: *const GpuTensor,
        n: u32,
        rows: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_rms_norm_weight_tensor(
        out: *mut GpuTensor,
        x: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        n: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_rms_norm_weight_rows_tensor(
        out: *mut GpuTensor,
        x: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        n: u32,
        rows: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_repeat_hc_tensor(
        out: *mut GpuTensor,
        x: *const GpuTensor,
        embd: u32,
        hc: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
        out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        sinks: u64,
        q: *const GpuTensor,
        raw_cache: *const GpuTensor,
        rows: u32,
        visible_rows: u32,
        cache_cap: u32,
        raw_start: u32,
        heads: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_hc_rms_scale_project_f16_tensor(
        out: *mut GpuTensor,
        scale: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        input: u32,
        output: u32,
        x: *const GpuTensor,
        rows: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_f16_tensor(
        out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        rows: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_q8_0_tensor(
        out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        rows: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_quant_tensor(
        out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        kind: u32,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        rows: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
        out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        kind: u32,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        rows: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_f32_tensor(
        out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        rows: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_embed_token_quant_tensor(
        out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        kind: u32,
        vocab: u32,
        token: u32,
        embd: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_embed_tokens_quant_tensor(
        out: *mut GpuTensor,
        tokens: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        kind: u32,
        vocab: u32,
        rows: u32,
        embd: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_rope_tail_tensor(
        x: *mut GpuTensor,
        tokens: u32,
        heads: u32,
        head_dim: u32,
        rot: u32,
        pos: u32,
        original: u32,
        freq_base: f32,
        freq_scale: f32,
        ext: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_qkv_norm_store_compact_kv_tensor(
        q_out: *mut GpuTensor,
        q: *const GpuTensor,
        map: *const c_void,
        size: u64,
        q_weight: u64,
        q_n: u32,
        kv_cache: *mut GpuTensor,
        rope_cache: *mut GpuTensor,
        kv_raw: *const GpuTensor,
        kv_weight: u64,
        pos: u32,
        tokens: u32,
        cache_cap: u32,
        kv_raw_dim: u32,
        kv_lora: u32,
        rot: u32,
        cache_f16: bool,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_kv_lora_rms_norm_tensor(
        out: *mut GpuTensor,
        kv_raw: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        rows: u32,
        kv_raw_dim: u32,
        kv_lora: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_store_compact_kv_tensor(
        kv_cache: *mut GpuTensor,
        rope_cache: *mut GpuTensor,
        kv_norm: *const GpuTensor,
        kv_raw: *const GpuTensor,
        pos: u32,
        rows: u32,
        cache_cap: u32,
        kv_raw_dim: u32,
        kv_lora: u32,
        rot: u32,
        cache_f16: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_store_indexer_k_tensor(
        cache: *mut GpuTensor,
        raw: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        bias: u64,
        pos: u32,
        tokens: u32,
        cache_cap: u32,
        head_dim: u32,
        rot: u32,
        original: u32,
        eps: f32,
        freq_base: f32,
        freq_scale: f32,
        ext: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
        cache_f16: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_fill_selected_range_tensor(
        selected: *mut GpuTensor,
        count: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_fill_selected_range_batch_tensor(
        selected: *mut GpuTensor,
        rows: u32,
        pos: u32,
        count: u32,
        pad_row: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_indexer_rope_tail_tensor(
        x: *mut GpuTensor,
        tokens: u32,
        heads: u32,
        head_dim: u32,
        rot: u32,
        pos: u32,
        original: u32,
        freq_base: f32,
        freq_scale: f32,
        ext: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_indexer_score_one_tensor(
        scores: *mut GpuTensor,
        q: *const GpuTensor,
        weights: *const GpuTensor,
        cache: *const GpuTensor,
        rows: u32,
        heads: u32,
        head_dim: u32,
        scale: f32,
        cache_f16: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_indexer_scores_batch_tensor(
        scores: *mut GpuTensor,
        q: *const GpuTensor,
        weights: *const GpuTensor,
        cache: *const GpuTensor,
        visible: u32,
        rows: u32,
        pos: u32,
        heads: u32,
        head_dim: u32,
        scale: f32,
        cache_f16: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_qk_lowrank_typed_tensor(
        out: *mut GpuTensor,
        q: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        kind: u32,
        heads: u32,
        kv_lora: u32,
        q_nope: u32,
        q_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_qk_lowrank_typed_batch_tensor(
        out: *mut GpuTensor,
        q: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        kind: u32,
        rows: u32,
        heads: u32,
        kv_lora: u32,
        q_nope: u32,
        q_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_value_project_typed_batch_heads_tensor(
        heads: *mut GpuTensor,
        lora: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        kind: u32,
        rows: u32,
        n_head: u32,
        kv_lora: u32,
        value_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_attention_indexed_decode_typed_tensor(
        heads_out: *mut GpuTensor,
        q: *const GpuTensor,
        qk_low: *const GpuTensor,
        kv_cache: *const GpuTensor,
        rope_cache: *const GpuTensor,
        map: *const c_void,
        size: u64,
        value_weight: u64,
        value_kind: u32,
        selected: *const GpuTensor,
        selected_count: u32,
        cache_cap: u32,
        cache_f16: bool,
        heads: u32,
        kv_lora: u32,
        q_nope: u32,
        rot: u32,
        value_dim: u32,
        original: u32,
        freq_base: f32,
        freq_scale: f32,
        ext: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
        out: *mut GpuTensor,
        q: *const GpuTensor,
        qk_low: *const GpuTensor,
        kv_cache: *const GpuTensor,
        rope_cache: *const GpuTensor,
        rows: u32,
        pos: u32,
        selected: u32,
        cache_cap: u32,
        cache_f16: bool,
        heads: u32,
        kv_lora: u32,
        q_nope: u32,
        rot: u32,
        original: u32,
        freq_base: f32,
        freq_scale: f32,
        ext: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
        out: *mut GpuTensor,
        q: *const GpuTensor,
        qk_low: *const GpuTensor,
        kv_cache: *const GpuTensor,
        rope_cache: *const GpuTensor,
        selected: *const GpuTensor,
        rows: u32,
        selected_count: u32,
        cache_cap: u32,
        cache_f16: bool,
        heads: u32,
        kv_lora: u32,
        q_nope: u32,
        rot: u32,
        original: u32,
        freq_base: f32,
        freq_scale: f32,
        ext: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_router_select_batch_tensor(
        selected: *mut GpuTensor,
        weights: *mut GpuTensor,
        probs: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        bias: u64,
        logits: *const GpuTensor,
        experts: u32,
        experts_used: u32,
        scale: f32,
        rows: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_routed_moe_batch_tensor(
        out: *mut GpuTensor,
        mid: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        gate_weight: u64,
        up_weight: u64,
        down_weight: u64,
        gate_type: u32,
        up_type: u32,
        down_type: u32,
        gate_expert_bytes: u64,
        gate_row_bytes: u64,
        up_expert_bytes: u64,
        up_row_bytes: u64,
        down_expert_bytes: u64,
        down_row_bytes: u64,
        input: u32,
        middle: u32,
        output: u32,
        selected: *const GpuTensor,
        weights: *const GpuTensor,
        total_experts: u32,
        used_experts: u32,
        layer: u32,
        x: *const GpuTensor,
        rows: u32,
        mid_token_stride: u32,
        force_resident: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_router_select_tensor(
        selected: *mut GpuTensor,
        weights: *mut GpuTensor,
        probs: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        bias: u64,
        logits: *const GpuTensor,
        experts: u32,
        used: u32,
        scale: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_glm_routed_moe_one_tensor(
        out: *mut GpuTensor,
        mid: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        gate: u64,
        up: u64,
        down: u64,
        gate_kind: u32,
        up_kind: u32,
        down_kind: u32,
        gate_expert_bytes: u64,
        gate_row_bytes: u64,
        up_expert_bytes: u64,
        up_row_bytes: u64,
        down_expert_bytes: u64,
        down_row_bytes: u64,
        input: u32,
        hidden: u32,
        output: u32,
        selected: *const GpuTensor,
        weights: *const GpuTensor,
        total_experts: u32,
        used: u32,
        layer: u32,
        x: *const GpuTensor,
        force_resident: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_q8_0_pair_tensor(
        out_a: *mut GpuTensor,
        out_b: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight_a: u64,
        weight_b: u64,
        input: u64,
        output_a: u64,
        output_b: u64,
        x: *const GpuTensor,
        rows: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_f16_pair_tensor(
        out_a: *mut GpuTensor,
        out_b: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight_a: u64,
        weight_b: u64,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        rows: u64,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_f16_pair_compressor_store_tensor(
        out_kv: *mut GpuTensor,
        out_score: *mut GpuTensor,
        state_kv: *mut GpuTensor,
        state_score: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight_kv: u64,
        weight_score: u64,
        ape: u64,
        ape_type: u32,
        input: u64,
        width: u32,
        x: *const GpuTensor,
        ratio: u32,
        pos: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_hc_split_weighted_sum_norm_tensor(
        out: *mut GpuTensor,
        norm: *mut GpuTensor,
        split: *mut GpuTensor,
        mix: *const GpuTensor,
        residual: *const GpuTensor,
        map: *const c_void,
        size: u64,
        scale: u64,
        base: u64,
        norm_weight: u64,
        embd: u32,
        hc: u32,
        iterations: u32,
        eps: f32,
        norm_eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
        q_out: *mut GpuTensor,
        q: *const GpuTensor,
        map: *const c_void,
        size: u64,
        q_weight: u64,
        q_n: u32,
        kv_out: *mut GpuTensor,
        kv: *const GpuTensor,
        kv_weight: u64,
        kv_n: u32,
        rows: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
        out: *mut GpuTensor,
        half: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        rows: u32,
        heads: u32,
        head_dim: u32,
        rot: u32,
        pos: u32,
        original_context: u32,
        inverse: bool,
        freq_base: f32,
        freq_scale: f32,
        ext_factor: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_head_rms_norm_tensor(
        x: *mut GpuTensor,
        rows: u32,
        heads: u32,
        head_dim: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_head_rms_norm_rope_tail_tensor(
        x: *mut GpuTensor,
        rows: u32,
        heads: u32,
        head_dim: u32,
        rot: u32,
        pos: u32,
        original_context: u32,
        inverse: bool,
        freq_base: f32,
        freq_scale: f32,
        ext_factor: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_rope_tail_tensor(
        x: *mut GpuTensor,
        rows: u32,
        heads: u32,
        head_dim: u32,
        rot: u32,
        pos: u32,
        original_context: u32,
        inverse: bool,
        freq_base: f32,
        freq_scale: f32,
        ext_factor: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_decode_heads_tensor(
        heads_out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        sinks: u64,
        q: *const GpuTensor,
        raw_kv: *const GpuTensor,
        n_raw: u32,
        raw_cap: u32,
        raw_start: u32,
        comp_kv: *const GpuTensor,
        comp_f16: u32,
        n_comp: u32,
        comp_mask: *const GpuTensor,
        use_mask: u32,
        heads: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_decode_raw_batch_heads_tensor(
        heads: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        sinks: u64,
        q: *const GpuTensor,
        raw_kv: *const GpuTensor,
        tokens: u32,
        pos: u32,
        n_raw: u32,
        raw_cap: u32,
        raw_start: u32,
        window: u32,
        heads_count: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_decode_mixed_batch_heads_tensor(
        heads: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        sinks: u64,
        q: *const GpuTensor,
        raw_kv: *const GpuTensor,
        compressed: *const GpuTensor,
        compressed_f16: u32,
        compressed_mask: *const GpuTensor,
        use_mask: u32,
        tokens: u32,
        pos: u32,
        n_raw: u32,
        raw_cap: u32,
        raw_start: u32,
        n_comp: u32,
        window: u32,
        ratio: u32,
        heads_count: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
        heads: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        sinks: u64,
        q: *const GpuTensor,
        raw_kv: *const GpuTensor,
        comp_kv: *const GpuTensor,
        comp_f16: u32,
        topk: *const GpuTensor,
        tokens: u32,
        pos: u32,
        n_raw: u32,
        raw_cap: u32,
        raw_start: u32,
        n_comp: u32,
        top_k: u32,
        window: u32,
        ratio: u32,
        heads: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_indexer_score_one_tensor(
        scores: *mut GpuTensor,
        q: *const GpuTensor,
        weights: *const GpuTensor,
        index_comp: *const GpuTensor,
        n_comp: u32,
        heads: u32,
        head_dim: u32,
        scale: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_indexer_scores_decode_batch_tensor(
        scores: *mut GpuTensor,
        q: *const GpuTensor,
        weights: *const GpuTensor,
        index_comp: *const GpuTensor,
        n_comp: u32,
        tokens: u32,
        pos: u32,
        heads: u32,
        head_dim: u32,
        ratio: u32,
        scale: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_indexer_topk_tensor(
        selected: *mut GpuTensor,
        scores: *const GpuTensor,
        n_comp: u32,
        tokens: u32,
        top_k: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_dsv4_indexer_qat_tensor(
        x: *mut GpuTensor,
        rows: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_compressor_update_tensor(
        kv: *const GpuTensor,
        score: *const GpuTensor,
        state_kv: *mut GpuTensor,
        state_score: *mut GpuTensor,
        cache: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        ape: u64,
        ape_type: u32,
        norm: u64,
        norm_type: u32,
        head_dim: u32,
        ratio: u32,
        pos: u32,
        row: u32,
        rot: u32,
        original_context: u32,
        freq_base: f32,
        freq_scale: f32,
        ext_factor: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
        rms_eps: f32,
        state_already_stored: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_compressor_prefill_state_ratio4_tensor(
        state_kv: *mut GpuTensor,
        state_score: *mut GpuTensor,
        kv_tail: *const GpuTensor,
        score_tail: *const GpuTensor,
        map: *const c_void,
        size: u64,
        ape: u64,
        ape_type: u32,
        head_dim: u32,
        pos: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_dsv4_fp8_kv_quantize_tensor(
        x: *mut GpuTensor,
        rows: u32,
        head_dim: u32,
        rot: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_kv_fp8_store_raw_tensor(
        kv: *mut GpuTensor,
        raw_cache: *mut GpuTensor,
        raw_cap: u32,
        row: u32,
        head_dim: u32,
        rot: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_store_raw_kv_batch_tensor(
        raw_cache: *mut GpuTensor,
        kv: *const GpuTensor,
        raw_cap: u32,
        pos: u32,
        rows: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_compressor_prefill_tensor(
        cache: *mut GpuTensor,
        state_kv: *mut GpuTensor,
        state_score: *mut GpuTensor,
        kv: *const GpuTensor,
        score: *const GpuTensor,
        map: *const c_void,
        size: u64,
        ape: u64,
        ape_type: u32,
        norm: u64,
        norm_type: u32,
        head_dim: u32,
        ratio: u32,
        pos: u32,
        rows: u32,
        rot: u32,
        original_context: u32,
        quantize_fp8: bool,
        freq_base: f32,
        freq_scale: f32,
        ext_factor: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
        rms_eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_compressor_prefill_ratio4_replay_tensor(
        cache: *mut GpuTensor,
        state_kv: *mut GpuTensor,
        state_score: *mut GpuTensor,
        kv: *const GpuTensor,
        score: *const GpuTensor,
        map: *const c_void,
        size: u64,
        ape: u64,
        ape_type: u32,
        norm: u64,
        norm_type: u32,
        head_dim: u32,
        pos: u32,
        rows: u32,
        rot: u32,
        original_context: u32,
        quantize_fp8: bool,
        freq_base: f32,
        freq_scale: f32,
        ext_factor: f32,
        attn_factor: f32,
        beta_fast: f32,
        beta_slow: f32,
        rms_eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_prefill_raw_heads_tensor(
        heads: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        sinks: u64,
        q: *const GpuTensor,
        raw: *const GpuTensor,
        rows: u32,
        window: u32,
        heads_count: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_prefill_static_mixed_heads_tensor(
        heads: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        sinks: u64,
        q: *const GpuTensor,
        raw: *const GpuTensor,
        compressed: *const GpuTensor,
        compressed_f16: u32,
        rows: u32,
        compressed_rows: u32,
        window: u32,
        ratio: u32,
        heads_count: u32,
        head_dim: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_indexer_scores_prefill_tensor(
        scores: *mut GpuTensor,
        q: *const GpuTensor,
        weights: *const GpuTensor,
        compressed: *const GpuTensor,
        compressed_rows: u32,
        rows: u32,
        heads: u32,
        head_dim: u32,
        ratio: u32,
        scale: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_output_q8_batch_f16_tensor(
        out_half: *mut GpuTensor,
        low: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight_a: u64,
        weight_b: u64,
        group_dim: u64,
        rank: u64,
        groups: u32,
        output: u64,
        heads: *const GpuTensor,
        rows: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_output_q8_batch_tensor(
        out: *mut GpuTensor,
        low: *mut GpuTensor,
        group_scratch: *mut GpuTensor,
        low_scratch: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight_a: u64,
        weight_b: u64,
        group_dim: u64,
        rank: u64,
        groups: u32,
        output: u64,
        heads: *const GpuTensor,
        rows: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_router_select_batch_tensor(
        selected: *mut GpuTensor,
        weights: *mut GpuTensor,
        probs: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        bias: u64,
        hash: u64,
        hash_rows: u32,
        expert_groups: u32,
        groups_used: u32,
        has_bias: bool,
        hash_mode: bool,
        logits: *const GpuTensor,
        tokens: *const GpuTensor,
        experts: u32,
        experts_used: u32,
        scale: f32,
        rows: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_routed_moe_batch_tensor(
        out: *mut GpuTensor,
        gate: *mut GpuTensor,
        up: *mut GpuTensor,
        mid: *mut GpuTensor,
        experts_out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        gate_weight: u64,
        up_weight: u64,
        down_weight: u64,
        gate_type: u32,
        down_type: u32,
        gate_expert_bytes: u64,
        gate_row_bytes: u64,
        down_expert_bytes: u64,
        down_row_bytes: u64,
        input: u32,
        middle: u32,
        output: u32,
        selected: *const GpuTensor,
        weights: *const GpuTensor,
        total_experts: u32,
        used_experts: u32,
        clamp: f32,
        x: *const GpuTensor,
        layer: u32,
        rows: u32,
        mid_f16: *mut bool,
        force_resident: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_swiglu_tensor(
        out: *mut GpuTensor,
        gate: *const GpuTensor,
        up: *const GpuTensor,
        count: u32,
        clamp: f32,
        scale: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_add_tensor(
        out: *mut GpuTensor,
        a: *const GpuTensor,
        b: *const GpuTensor,
        count: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_add3_tensor(
        out: *mut GpuTensor,
        a: *const GpuTensor,
        b: *const GpuTensor,
        c: *const GpuTensor,
        count: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_directional_steering_project_tensor(
        x: *mut GpuTensor,
        directions: *const GpuTensor,
        layer: u32,
        width: u32,
        rows: u32,
        scale: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_add_rms_norm_weight_tensor(
        norm: *mut GpuTensor,
        sum: *mut GpuTensor,
        a: *const GpuTensor,
        b: *const GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        count: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_hc_expand_split_half_tensor(
        out: *mut GpuTensor,
        block_half: *const GpuTensor,
        residual: *const GpuTensor,
        split: *const GpuTensor,
        embd: u32,
        hc: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_hc_expand_split_tensor(
        out: *mut GpuTensor,
        block: *const GpuTensor,
        residual: *const GpuTensor,
        split: *const GpuTensor,
        embd: u32,
        hc: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_hc_expand_add_split_tensor(
        out: *mut GpuTensor,
        block: *const GpuTensor,
        add: *const GpuTensor,
        residual: *const GpuTensor,
        split: *const GpuTensor,
        embd: u32,
        hc: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_attention_output_low_q8_tensor(
        low: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        group_dim: u64,
        rank: u64,
        groups: u32,
        heads: *const GpuTensor,
    ) -> i32;
    pub(super) fn ds4_gpu_matmul_q8_0_hc_expand_tensor(
        out_hc: *mut GpuTensor,
        block_out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        residual: *const GpuTensor,
        split: *const GpuTensor,
        embd: u32,
        hc: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_router_select_tensor(
        selected: *mut GpuTensor,
        weights: *mut GpuTensor,
        probs: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        bias: u64,
        hash: u64,
        hash_rows: u32,
        token: u32,
        experts: u32,
        used: u32,
        scale: f32,
        expert_groups: u32,
        groups_used: u32,
        has_bias: bool,
        hash_mode: bool,
        logits: *const GpuTensor,
    ) -> i32;
    pub(super) fn ds4_gpu_routed_moe_one_tensor(
        out: *mut GpuTensor,
        gate: *mut GpuTensor,
        up: *mut GpuTensor,
        mid: *mut GpuTensor,
        experts_out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        gate_weight: u64,
        up_weight: u64,
        down_weight: u64,
        gate_type: u32,
        down_type: u32,
        gate_expert_bytes: u64,
        gate_row_bytes: u64,
        down_expert_bytes: u64,
        down_row_bytes: u64,
        input: u32,
        middle: u32,
        output: u32,
        selected: *const GpuTensor,
        weights: *const GpuTensor,
        total_experts: u32,
        used_experts: u32,
        clamp: f32,
        x: *const GpuTensor,
        add: *const GpuTensor,
        layer: u32,
        force_resident: bool,
    ) -> i32;
    pub(super) fn ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
        gate: *mut GpuTensor,
        up: *mut GpuTensor,
        mid: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        gate_weight: u64,
        up_weight: u64,
        input: u64,
        output: u64,
        x: *const GpuTensor,
        clamp: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_shared_down_hc_expand_q8_0_tensor(
        out_hc: *mut GpuTensor,
        shared_out: *mut GpuTensor,
        map: *const c_void,
        size: u64,
        weight: u64,
        input: u64,
        output: u64,
        middle: *const GpuTensor,
        routed: *const GpuTensor,
        residual: *const GpuTensor,
        split: *const GpuTensor,
        embd: u32,
        hc: u32,
    ) -> i32;
    pub(super) fn ds4_gpu_output_hc_weights_tensor(
        out: *mut GpuTensor,
        pre: *const GpuTensor,
        map: *const c_void,
        size: u64,
        scale: u64,
        base: u64,
        hc: u32,
        eps: f32,
    ) -> i32;
    pub(super) fn ds4_gpu_hc_weighted_sum_tensor(
        out: *mut GpuTensor,
        residual: *const GpuTensor,
        weights: *const GpuTensor,
        embd: u32,
        hc: u32,
    ) -> i32;
}

pub(super) struct Context {
    _model_file: File,
}

impl Context {
    pub(super) fn open(
        model: &Model,
        quality: bool,
        ssd_streaming: bool,
        admission_bytes: u64,
        model_spans: Option<&[(u64, u64)]>,
    ) -> Result<Self, String> {
        check(unsafe { ds4_gpu_init() }, "Metal initialization")?;
        unsafe {
            ds4_gpu_set_glm_model(model.shape.family == ModelFamily::Glm);
            ds4_gpu_set_ssd_streaming(ssd_streaming);
        }
        let recommended = unsafe { ds4_gpu_recommended_working_set_size() };
        if admission_bytes != 0 && recommended != 0 && admission_bytes > recommended {
            unsafe { ds4_gpu_cleanup() };
            return Err(format!(
                "model load needs {:.1} GiB including context and scratch, but Metal recommends at most {:.1} GiB",
                admission_bytes as f64 / 1_073_741_824.0,
                recommended as f64 / 1_073_741_824.0,
            ));
        }
        let data_offset = model.main.data_offset();
        let mapped = if let Some(spans) = model_spans {
            let (offsets, sizes): (Vec<_>, Vec<_>) = spans.iter().copied().unzip();
            unsafe {
                ds4_gpu_set_model_map_spans(
                    model.main.map_ptr().cast(),
                    model.main.len(),
                    offsets.as_ptr(),
                    sizes.as_ptr(),
                    spans.len() as u32,
                    model.main.max_tensor_bytes(),
                )
            }
        } else {
            unsafe {
                ds4_gpu_set_model_map_range(
                    model.main.map_ptr().cast(),
                    model.main.len(),
                    data_offset,
                    model.main.len() - data_offset,
                    model.main.max_tensor_bytes(),
                )
            }
        };
        if let Err(error) = check(mapped, "model mapping") {
            unsafe { ds4_gpu_cleanup() };
            return Err(error);
        }
        if let Some(support) = &model.support {
            let mapped = unsafe {
                ds4_gpu_set_model_map_range(
                    support.map_ptr().cast(),
                    support.len(),
                    support.data_offset(),
                    support.len() - support.data_offset(),
                    support.max_tensor_bytes(),
                )
            };
            if let Err(error) = check(mapped, "support-model mapping") {
                unsafe { ds4_gpu_cleanup() };
                return Err(error);
            }
        }
        unsafe { ds4_gpu_set_quality(quality) };
        let model_file = File::open(model.main.path()).map_err(|error| {
            unsafe { ds4_gpu_cleanup() };
            error.to_string()
        })?;
        #[cfg(target_os = "macos")]
        {
            use std::os::fd::AsRawFd;
            if let Err(error) = check(
                unsafe { ds4_gpu_set_model_fd(model_file.as_raw_fd()) },
                "model file registration",
            ) {
                unsafe { ds4_gpu_cleanup() };
                return Err(error);
            }
        }
        Ok(Self {
            _model_file: model_file,
        })
    }
}

impl Drop for Context {
    fn drop(&mut self) {
        unsafe { ds4_gpu_cleanup() };
    }
}

pub(super) struct Commands(bool);

impl Commands {
    pub(super) fn begin() -> Result<Self, String> {
        check(
            unsafe { ds4_gpu_begin_commands() },
            "beginning Metal commands",
        )?;
        Ok(Self(true))
    }

    pub(super) fn finish(mut self) -> Result<(), String> {
        self.0 = false;
        check(
            unsafe { ds4_gpu_end_commands() },
            "executing Metal commands",
        )
    }
}

impl Drop for Commands {
    fn drop(&mut self) {
        if self.0 {
            unsafe { ds4_gpu_end_commands() };
        }
    }
}

pub(super) struct Buffer(NonNull<GpuTensor>);

impl Buffer {
    pub(super) fn floats(count: u64) -> Result<Self, String> {
        Self::bytes(count * 4)
    }

    pub(super) fn bytes(bytes: u64) -> Result<Self, String> {
        NonNull::new(unsafe { ds4_gpu_tensor_alloc(bytes) })
            .map(Self)
            .ok_or_else(|| format!("Metal could not allocate {bytes} bytes"))
    }

    pub(super) fn view(&self, offset: u64, bytes: u64) -> Result<Self, String> {
        // SAFETY: `ds4_gpu_tensor_view` in native/metal/ds4_metal.m bounds-checks
        // the view, ARC-retains base_obj.buffer, and marks the view non-owning, so
        // it may outlive and be freed independently of this wrapper.
        // Recheck those guarantees whenever the vendored Metal file is re-synced.
        NonNull::new(unsafe { ds4_gpu_tensor_view(self.raw(), offset, bytes) })
            .map(Self)
            .ok_or_else(|| "Metal could not create a tensor view".to_owned())
    }

    pub(super) fn read_f32(&self, values: &mut [f32]) -> Result<(), String> {
        check(
            unsafe {
                ds4_gpu_tensor_read(
                    self.raw(),
                    0,
                    values.as_mut_ptr().cast(),
                    std::mem::size_of_val(values) as u64,
                )
            },
            "reading Metal output",
        )
    }

    pub(super) fn read_i32(&self, values: &mut [i32]) -> Result<(), String> {
        check(
            unsafe {
                ds4_gpu_tensor_read(
                    self.raw(),
                    0,
                    values.as_mut_ptr().cast(),
                    std::mem::size_of_val(values) as u64,
                )
            },
            "reading Metal integers",
        )
    }

    pub(super) fn read(&self, offset: u64, values: &mut [u8]) -> Result<(), String> {
        check(
            unsafe {
                ds4_gpu_tensor_read(
                    self.raw(),
                    offset,
                    values.as_mut_ptr().cast(),
                    values.len() as u64,
                )
            },
            "reading a Metal buffer",
        )
    }

    pub(super) fn write(&self, offset: u64, values: &[u8]) -> Result<(), String> {
        check(
            unsafe {
                ds4_gpu_tensor_write(
                    self.raw(),
                    offset,
                    values.as_ptr().cast(),
                    values.len() as u64,
                )
            },
            "restoring a Metal buffer",
        )
    }

    pub(super) fn write_i32(&self, values: &[i32]) -> Result<(), String> {
        check(
            unsafe {
                ds4_gpu_tensor_write(
                    self.raw(),
                    0,
                    values.as_ptr().cast(),
                    std::mem::size_of_val(values) as u64,
                )
            },
            "uploading tokens",
        )
    }

    pub(super) fn write_f32(&self, values: &[f32]) -> Result<(), String> {
        check(
            unsafe {
                ds4_gpu_tensor_write(
                    self.raw(),
                    0,
                    values.as_ptr().cast(),
                    std::mem::size_of_val(values) as u64,
                )
            },
            "uploading floats",
        )
    }

    pub(super) fn fill(&self, value: f32, count: u64) -> Result<(), String> {
        call(
            unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) },
            "initializing a Metal buffer",
        )
    }

    pub(super) fn copy_from(
        &self,
        destination_offset: u64,
        source: &Self,
        source_offset: u64,
        bytes: u64,
        operation: &str,
    ) -> Result<(), String> {
        call(
            unsafe {
                ds4_gpu_tensor_copy(
                    self.raw(),
                    destination_offset,
                    source.raw(),
                    source_offset,
                    bytes,
                )
            },
            operation,
        )
    }

    pub(super) fn raw(&self) -> *mut GpuTensor {
        self.0.as_ptr()
    }
}

impl Drop for Buffer {
    fn drop(&mut self) {
        unsafe { ds4_gpu_tensor_free(self.raw()) };
    }
}
