use cudarc::driver::CudaContext;
use luminal::prelude::*;
use rand::SeedableRng;

use luminal::egglog_utils::{egglog_to_llir, random_initial_choice, validate_choice_set};

use crate::kernel::KernelOp;
use crate::runtime::CudaRuntime;

/// Helper: build search space and extract all possible kernel names across many random choices.
fn extract_all_kernel_names(cx: &mut Graph) -> Vec<String> {
    cx.build_search_space::<CudaRuntime>();
    let egraph = cx.egraph().expect("egraph not built");
    let ops = cx.egglog_ops().expect("ops not built");
    let custom_ops = &cx.custom_ops;

    let mut all_names = Vec::new();
    // Try many random extractions to cover both alternatives
    for _ in 0..20 {
        let choices = random_initial_choice(egraph, &mut rand::rng());
        let mut list_cache = Default::default();
        let mut expr_cache = Default::default();
        let llir = egglog_to_llir(
            egraph,
            choices,
            ops,
            custom_ops,
            &mut list_cache,
            &mut expr_cache,
            None,
        );
        for op in llir.node_weights() {
            if let Some(k) = op.to_dialect::<dyn KernelOp>() {
                let name = k.kernel_name().to_string();
                if !all_names.contains(&name) {
                    all_names.push(name);
                }
            }
        }
    }
    all_names
}

/// When dest is NOT shared with any other op, KernelScatterNoCopy should be available.
/// The ConsumedBuffer cleanup rule should NOT fire because dest only appears inside
/// the ConsumedBuffer (not in any other ICons).
#[test]
fn test_scatter_nocopy_selected_when_dest_unshared() {
    let ctx = CudaContext::new(0).unwrap();
    ctx.bind_to_thread().unwrap();

    let mut cx = Graph::default();

    // dest: a 10-element buffer, src: 3 values, indexes: 3 indices
    let dest = cx.tensor(10).persist();
    let src = cx.tensor(3).persist();
    let indexes = cx.tensor(3).as_dtype(DType::Int).persist();

    // scatter src into dest at indexes
    let _result = src.scatter(indexes, dest).output();

    let names = extract_all_kernel_names(&mut cx);
    println!("All possible kernels: {:?}", names);

    // KernelScatterNoCopy should be available (dest is not shared)
    assert!(
        names.iter().any(|n| n == "ScatterNoCopy"),
        "Expected ScatterNoCopy to be available but got: {:?}",
        names
    );
}

/// When dest IS shared (used by another op besides the scatter), the ConsumedBuffer
/// cleanup rule should fire, deleting the ConsumedBuffer. This makes KernelScatterNoCopy
/// invalid, so it should NOT appear in any extraction.
#[test]
fn test_scatter_nocopy_not_selected_when_dest_shared() {
    let ctx = CudaContext::new(0).unwrap();
    ctx.bind_to_thread().unwrap();

    let mut cx = Graph::default();

    // dest: a 10-element buffer, src: 3 values, indexes: 3 indices
    let dest = cx.tensor(10).persist();
    let src = cx.tensor(3).persist();
    let indexes = cx.tensor(3).as_dtype(DType::Int).persist();

    // scatter src into dest at indexes
    let scatter_result = src.scatter(indexes, dest);

    // Also use dest directly in another op (add with itself) — this makes dest shared
    let _dest_also_used = (dest + dest).output();
    let _result = scatter_result.output();

    let names = extract_all_kernel_names(&mut cx);
    println!("All possible kernels: {:?}", names);

    // KernelScatterNoCopy should NOT be available (dest is shared with the add op)
    assert!(
        !names.iter().any(|n| n == "ScatterNoCopy"),
        "ScatterNoCopy should NOT be available when dest is shared, got: {:?}",
        names
    );
    // Regular KernelScatter should be present
    assert!(
        names.iter().any(|n| n == "Scatter"),
        "Expected Scatter but got: {:?}",
        names
    );
}

/// Actually execute the scatter and verify correctness.
/// Tests all possible extractions (both KernelScatter and KernelScatterNoCopy).
#[test]
fn test_scatter_execution_correctness() {
    let ctx = CudaContext::new(0).unwrap();
    ctx.bind_to_thread().unwrap();
    let stream = ctx.default_stream();

    let mut cx = Graph::default();

    // dest: [0.0, 1.0, 2.0, 3.0, 4.0]
    let dest = cx.tensor(5).persist();
    // src: [10.0, 20.0, 30.0]
    let src = cx.tensor(3).persist();
    // indexes: [1, 3, 4]
    let indexes = cx.tensor(3).as_dtype(DType::Int).persist();

    let result = src.scatter(indexes, dest).output();

    cx.build_search_space::<CudaRuntime>();
    let egraph = cx.egraph().expect("egraph not built");
    let ops = cx.egglog_ops().expect("ops not built");

    // Expected: [0.0, 10.0, 2.0, 20.0, 30.0]
    let expected = vec![0.0f32, 10.0, 2.0, 20.0, 30.0];

    // Try many random extractions to cover both Scatter and ScatterNoCopy
    let mut rng = rand::rng();
    let mut tested_scatter = false;
    let mut tested_nocopy = false;

    for _ in 0..50 {
        let choices = random_initial_choice(egraph, &mut rng);
        if validate_choice_set(egraph, &choices, ops).is_err() {
            continue;
        }

        let mut list_cache = Default::default();
        let mut expr_cache = Default::default();
        let llir = egglog_to_llir(
            egraph,
            choices,
            ops,
            &cx.custom_ops,
            &mut list_cache,
            &mut expr_cache,
            None,
        );

        // Check which scatter variant was selected
        let mut has_nocopy = false;
        let mut has_scatter = false;
        for op in llir.node_weights() {
            if let Some(k) = op.to_dialect::<dyn KernelOp>() {
                match k.kernel_name() {
                    "ScatterNoCopy" => has_nocopy = true,
                    "Scatter" => has_scatter = true,
                    _ => {}
                }
            }
        }

        let mut rt = CudaRuntime::initialize(stream.clone());
        rt.load_llir(&llir);
        rt.set_data(dest, vec![0.0f32, 1.0, 2.0, 3.0, 4.0]);
        rt.set_data(src, vec![10.0f32, 20.0, 30.0]);
        rt.set_data(indexes, vec![1i32, 3, 4]);
        rt.execute(&cx.dyn_map);

        let actual = rt.get_f32(result);

        let variant = if has_nocopy {
            tested_nocopy = true;
            "ScatterNoCopy"
        } else if has_scatter {
            tested_scatter = true;
            "Scatter"
        } else {
            "Unknown"
        };

        assert_eq!(
            actual, expected,
            "Scatter result mismatch with variant {variant}: got {:?}, expected {:?}",
            actual, expected
        );
    }

    println!(
        "Tested Scatter: {}, Tested ScatterNoCopy: {}",
        tested_scatter, tested_nocopy
    );
    assert!(
        tested_nocopy,
        "ScatterNoCopy was never selected in 50 attempts — can't verify correctness"
    );
}

/// Test the KV-cache round-trip pattern: scatter → remove_buffer → set_buffer → scatter again.
/// This mimics how the llama model uses scatter for KV cache updates.
#[test]
fn test_scatter_kv_cache_roundtrip() {
    let ctx = CudaContext::new(0).unwrap();
    ctx.bind_to_thread().unwrap();
    let stream = ctx.default_stream();

    let mut cx = Graph::default();

    // KV cache: [5] elements (simulating a small cache)
    let cache_in = cx.named_tensor("cache", 5).persist();
    // New value to scatter: [1] element
    let src = cx.tensor(1).persist();
    // Index: [1] element (position to write)
    let indexes = cx.tensor(1).as_dtype(DType::Int).persist();

    // scatter src into cache at index position
    let cache_out = src.scatter(indexes, cache_in);
    // Also read the scatter output (simulates attention reading from cache)
    let read_out = (cache_out + 0.0).output();
    // Return cache for round-trip
    let cache_output = cache_out.output();

    cx.build_search_space::<CudaRuntime>();

    let mut rt = CudaRuntime::initialize(stream.clone());

    // Must set input data BEFORE search (profiler needs valid buffers)
    rt.set_data(cache_in, vec![0.0f32; 5]);
    rt.set_data(src, vec![10.0f32]);
    rt.set_data(indexes, vec![0i32]);

    rt = cx.search(rt, 5);

    // Print which scatter variant was selected
    for node in rt.llir_graph().node_weights() {
        if let Some(k) = node.to_dialect::<dyn KernelOp>()
            && k.kernel_name().contains("catter")
        {
            println!("Selected: {}", k.kernel_name());
        }
    }

    // Step 1: Initialize cache to zeros, scatter 10.0 at position 0
    rt.set_data(cache_in, vec![0.0f32; 5]);
    rt.set_data(src, vec![10.0f32]);
    rt.set_data(indexes, vec![0i32]);
    rt.execute(&cx.dyn_map);

    let read1 = rt.get_f32(read_out);
    println!("After step 1 (scatter 10.0 at pos 0): {:?}", read1);
    assert_eq!(
        read1,
        vec![10.0, 0.0, 0.0, 0.0, 0.0],
        "Step 1 read_out mismatch"
    );

    // Round-trip: remove cache output buffer, set as new cache input
    let cache_buf = rt.remove_buffer(cache_output);
    rt.set_buffer(cache_in, cache_buf);

    // Step 2: Scatter 20.0 at position 1
    rt.set_data(src, vec![20.0f32]);
    rt.set_data(indexes, vec![1i32]);
    rt.execute(&cx.dyn_map);

    let read2 = rt.get_f32(read_out);
    println!("After step 2 (scatter 20.0 at pos 1): {:?}", read2);
    assert_eq!(
        read2,
        vec![10.0, 20.0, 0.0, 0.0, 0.0],
        "Step 2 read_out mismatch"
    );

    // Round-trip again
    let cache_buf = rt.remove_buffer(cache_output);
    rt.set_buffer(cache_in, cache_buf);

    // Step 3: Scatter 30.0 at position 2
    rt.set_data(src, vec![30.0f32]);
    rt.set_data(indexes, vec![2i32]);
    rt.execute(&cx.dyn_map);

    let read3 = rt.get_f32(read_out);
    println!("After step 3 (scatter 30.0 at pos 2): {:?}", read3);
    assert_eq!(
        read3,
        vec![10.0, 20.0, 30.0, 0.0, 0.0],
        "Step 3 read_out mismatch"
    );
}

/// Test scatter with TWO cache buffers and dual outputs (closer to llama K+V pattern).
#[test]
fn test_scatter_dual_cache() {
    let ctx = CudaContext::new(0).unwrap();
    ctx.bind_to_thread().unwrap();
    let stream = ctx.default_stream();

    let mut cx = Graph::default();

    // Two caches (like K and V)
    let k_cache = cx.named_tensor("k_cache", 5).persist();
    let v_cache = cx.named_tensor("v_cache", 5).persist();

    // Input values
    let k_new = cx.tensor(1).persist();
    let v_new = cx.tensor(1).persist();
    let indexes = cx.tensor(1).as_dtype(DType::Int).persist();

    // Scatter into both caches
    let k_out = k_new.scatter(indexes, k_cache);
    let v_out = v_new.scatter(indexes, v_cache);

    // Read both (simulates attention using the scattered caches)
    let k_read = k_out + 0.0;
    let v_read = v_out + 0.0;

    // Compute something from the scattered values (simulates attention output)
    let attn = k_read * v_read;

    // Output everything
    let attn_out = attn.output();
    let k_cache_out = k_out.output();
    let v_cache_out = v_out.output();

    cx.build_search_space::<CudaRuntime>();

    let mut rt = CudaRuntime::initialize(stream.clone());

    rt.set_data(k_cache, vec![0.0f32; 5]);
    rt.set_data(v_cache, vec![0.0f32; 5]);
    rt.set_data(k_new, vec![2.0f32]);
    rt.set_data(v_new, vec![3.0f32]);
    rt.set_data(indexes, vec![0i32]);

    // Use seeded search for deterministic scatter variant selection.
    // Seed 0 reliably selects Scatter (not ScatterNoCopy) for both caches.
    let mut rng = rand::rngs::SmallRng::seed_from_u64(0);
    rt = cx.search_options(rt, SearchOptions::new(5), &mut rng);

    // Print selected variants
    for node in rt.llir_graph().node_weights() {
        if let Some(k) = node.to_dialect::<dyn KernelOp>()
            && k.kernel_name().contains("catter")
        {
            println!("Dual test selected: {}", k.kernel_name());
        }
    }

    // Step 1: scatter k=2.0, v=3.0 at position 0
    rt.set_data(k_cache, vec![0.0f32; 5]);
    rt.set_data(v_cache, vec![0.0f32; 5]);
    rt.set_data(k_new, vec![2.0f32]);
    rt.set_data(v_new, vec![3.0f32]);
    rt.set_data(indexes, vec![0i32]);
    rt.execute(&cx.dyn_map);

    let attn1 = rt.get_f32(attn_out);
    println!("Attn step 1: {:?}", attn1);
    // k=[2,0,0,0,0], v=[3,0,0,0,0], attn = k*v = [6,0,0,0,0]
    assert_eq!(attn1, vec![6.0, 0.0, 0.0, 0.0, 0.0], "Step 1 attn mismatch");

    // Round-trip
    let k_buf = rt.remove_buffer(k_cache_out);
    let v_buf = rt.remove_buffer(v_cache_out);
    rt.set_buffer(k_cache, k_buf);
    rt.set_buffer(v_cache, v_buf);

    // Step 2: scatter k=4.0, v=5.0 at position 1
    rt.set_data(k_new, vec![4.0f32]);
    rt.set_data(v_new, vec![5.0f32]);
    rt.set_data(indexes, vec![1i32]);
    rt.execute(&cx.dyn_map);

    let attn2 = rt.get_f32(attn_out);
    println!("Attn step 2: {:?}", attn2);
    // k=[2,4,0,0,0], v=[3,5,0,0,0], attn = k*v = [6,20,0,0,0]
    assert_eq!(
        attn2,
        vec![6.0, 20.0, 0.0, 0.0, 0.0],
        "Step 2 attn mismatch"
    );

    // Round-trip
    let k_buf = rt.remove_buffer(k_cache_out);
    let v_buf = rt.remove_buffer(v_cache_out);
    rt.set_buffer(k_cache, k_buf);
    rt.set_buffer(v_cache, v_buf);

    // Step 3: scatter k=6.0, v=7.0 at position 2
    rt.set_data(k_new, vec![6.0f32]);
    rt.set_data(v_new, vec![7.0f32]);
    rt.set_data(indexes, vec![2i32]);
    rt.execute(&cx.dyn_map);

    let attn3 = rt.get_f32(attn_out);
    println!("Attn step 3: {:?}", attn3);
    // k=[2,4,6,0,0], v=[3,5,7,0,0], attn = k*v = [6,20,42,0,0]
    assert_eq!(
        attn3,
        vec![6.0, 20.0, 42.0, 0.0, 0.0],
        "Step 3 attn mismatch"
    );
}
