//! Phase 10 device routing (CPU scaffold until GPU kernels ship).

use super::fixture;
use crate::layout::mmap_file_read;
use crate::query::{
    ExecutionDeviceHint, GPU_HOST_MATERIALIZE_RAM_FRACTION, cuda_backend_available,
    host_materialize_exceeds, metal_backend_available, plan_query_with_tet_mmap,
    resolve_device_route, validate_query,
};
use crate::utils::dtype::ElementDtype;

#[test]
fn host_materialize_exceeds_when_buffer_larger_than_host_budget() {
    let gib = 1024_u64.pow(3);
    assert!(!host_materialize_exceeds(gib, Some(8 * gib)));
    let limit = ((8 * gib) as f64 * GPU_HOST_MATERIALIZE_RAM_FRACTION) as u64;
    assert!(host_materialize_exceeds(limit + 1, Some(8 * gib)));
    assert!(!host_materialize_exceeds(gib, None));
    assert!(host_materialize_exceeds(20 * gib, None));
}

#[test]
fn execution_device_hint_parse_tokens() {
    assert_eq!(
        ExecutionDeviceHint::parse("cpu").unwrap(),
        ExecutionDeviceHint::Cpu
    );
    assert_eq!(
        ExecutionDeviceHint::parse("auto").unwrap(),
        ExecutionDeviceHint::Auto
    );
    assert_eq!(
        ExecutionDeviceHint::parse("cuda").unwrap(),
        ExecutionDeviceHint::Cuda(0)
    );
    assert_eq!(
        ExecutionDeviceHint::parse("cuda:2").unwrap(),
        ExecutionDeviceHint::Cuda(2)
    );
    assert_eq!(
        ExecutionDeviceHint::parse("metal").unwrap(),
        ExecutionDeviceHint::Metal
    );
    assert_eq!(
        ExecutionDeviceHint::parse("cuda:multi").unwrap(),
        ExecutionDeviceHint::CudaMulti
    );
    assert_eq!(
        ExecutionDeviceHint::parse("rocm").unwrap(),
        ExecutionDeviceHint::Rocm(0)
    );
    assert_eq!(
        ExecutionDeviceHint::parse("rocm:multi").unwrap(),
        ExecutionDeviceHint::RocmMulti
    );
    assert_eq!(
        ExecutionDeviceHint::parse("rocm:2").unwrap(),
        ExecutionDeviceHint::Rocm(2)
    );
}

#[test]
fn parse_query_json_execution_device() {
    let doc = crate::query::parse_query_json(
        r#"{"dataset":"a","mean":[],"execution":{"device":"auto"}}"#,
    )
    .unwrap();
    validate_query(&doc).unwrap();
    assert_eq!(
        doc.execution.as_ref().unwrap().device,
        Some(ExecutionDeviceHint::Auto)
    );
}

#[test]
fn resolve_device_auto_small_selection_stays_cpu() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("small.tet");
    fixture::write_multichunk_2x3_tiles(&path, "a");
    let doc = fixture::query_files::json("mean_a");
    let mmap = mmap_file_read(&path).unwrap();
    let response = plan_query_with_tet_mmap(&doc, None, &mmap, Some(0)).unwrap();
    let plan = response.read_plan.as_ref().unwrap();
    let route = resolve_device_route(
        Some(&crate::query::ExecutionHints {
            device: Some(ExecutionDeviceHint::Auto),
            ..Default::default()
        }),
        plan,
        ElementDtype::F32,
        doc.operation.as_ref(),
    );
    assert_eq!(route.used, "cpu");
    assert_eq!(route.fallback_reason, Some("auto_below_size_threshold"));
}

#[test]
fn resolve_device_cuda_falls_back_without_cuda_feature() {
    if cuda_backend_available() {
        return;
    }
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("big.tet");
    fixture::write_verify_large_f32(&path, "a");
    let doc = crate::query::parse_query_json(
        r#"{"dataset":"a","mean":[],"execution":{"device":"cuda:0"}}"#,
    )
    .unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let response = plan_query_with_tet_mmap(&doc, None, &mmap, Some(0)).unwrap();
    let ex = response.execution.as_ref().unwrap();
    assert_eq!(ex.device_requested.as_deref(), Some("cuda"));
    assert_eq!(ex.device_used, Some("cpu"));
    assert_eq!(
        ex.device_fallback_reason.as_deref(),
        Some("gpu_feature_disabled")
    );
    assert_eq!(ex.device_gpu_reduce, Some(false));
}

#[test]
fn resolve_device_metal_falls_back_without_metal_feature() {
    if metal_backend_available() {
        return;
    }
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("big.tet");
    fixture::write_verify_large_f32(&path, "a");
    let doc = crate::query::parse_query_json(
        r#"{"dataset":"a","mean":[],"execution":{"device":"metal"}}"#,
    )
    .unwrap();
    validate_query(&doc).unwrap();
    let mmap = mmap_file_read(&path).unwrap();
    let response = plan_query_with_tet_mmap(&doc, None, &mmap, Some(0)).unwrap();
    let ex = response.execution.as_ref().unwrap();
    assert_eq!(ex.device_requested.as_deref(), Some("metal"));
    assert_eq!(ex.device_used, Some("cpu"));
    assert_eq!(
        ex.device_fallback_reason.as_deref(),
        Some("gpu_feature_disabled")
    );
}

#[cfg(all(feature = "tetration-metal", target_os = "macos"))]
#[test]
fn resolve_device_auto_large_prefers_metal_when_enabled() {
    use crate::query::GPU_AUTO_MIN_LOGICAL_BYTES;
    use crate::query::types::ReadPlan;

    let doc = crate::query::parse_query_json(r#"{"dataset":"a","mean":[]}"#).unwrap();
    validate_query(&doc).unwrap();
    let elems = usize::try_from(GPU_AUTO_MIN_LOGICAL_BYTES / 4 + 1).unwrap();
    let plan = ReadPlan {
        chunk_touch_policy: "dense",
        chunk_count: 0,
        total_stored_bytes: 0,
        chunks: vec![],
        dataset_shape: vec![elems as u64],
        chunk_shape: vec![elems as u64],
        selection_box_start: vec![0],
        selection_box_stop_exclusive: vec![elems as u64],
        selection_step: vec![1],
        logical_selection_shape: vec![elems as u64],
        logical_f32_element_count: elems,
    };
    let route = resolve_device_route(
        Some(&crate::query::ExecutionHints {
            device: Some(ExecutionDeviceHint::Auto),
            ..Default::default()
        }),
        &plan,
        ElementDtype::F32,
        doc.operation.as_ref(),
    );
    assert_eq!(route.used, "metal");
    assert!(route.gpu_reduce);
}

#[test]
fn resolve_device_auto_huge_selection_still_routes_gpu_for_streaming_fold() {
    use crate::query::types::ReadPlan;
    use crate::utils::host_memory;

    if !metal_backend_available() && !cuda_backend_available() {
        return;
    }

    let gib = 1024_u64.pow(3);
    let logical_bytes = 20 * gib;
    if !host_materialize_exceeds(logical_bytes, host_memory::available_memory_bytes()) {
        return;
    }

    let doc = crate::query::parse_query_json(r#"{"dataset":"a","mean":[]}"#).unwrap();
    let elems = (logical_bytes / 4) as usize;
    let plan = ReadPlan {
        chunk_touch_policy: "dense",
        chunk_count: 0,
        total_stored_bytes: 0,
        chunks: vec![],
        dataset_shape: vec![elems as u64],
        chunk_shape: vec![elems as u64],
        selection_box_start: vec![0],
        selection_box_stop_exclusive: vec![elems as u64],
        selection_step: vec![1],
        logical_selection_shape: vec![elems as u64],
        logical_f32_element_count: elems,
    };
    let route = resolve_device_route(
        Some(&crate::query::ExecutionHints {
            device: Some(ExecutionDeviceHint::Auto),
            ..Default::default()
        }),
        &plan,
        ElementDtype::F32,
        doc.operation.as_ref(),
    );
    assert!(route.gpu_reduce);
    assert!(
        route.used == "metal" || route.used == "cuda",
        "expected GPU backend, got {}",
        route.used
    );
    assert!(route.fallback_reason.is_none());
}
