use std::path::PathBuf;

use half::bf16;
use xserv_kernels::{
    GemmBackend, add, matmul,
    moe::{moe_sparse_gemv_bf16, moe_topk_softmax, moe_weighted_sum_sparse},
    mul, row_scale_bf16, sigmoid, silu,
};
use xserv_model::{ModelConfig, Qwen35MoeSpec, loader};
use xserv_tensor::{Device, Tensor};

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: smoke-qwen35-moe <model-dir> [layer] [seq_len]");
        std::process::exit(1);
    }
    let model_dir = PathBuf::from(&args[1]);
    let layer: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(0);
    let seq_len: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(1);
    let device = 0;

    xserv_cuda::device::set_device(device).unwrap();
    xserv_model::init_kernels();

    let config = ModelConfig::from_file(&model_dir.join("config.json"));
    let spec = Qwen35MoeSpec::from_config(&config);
    let weight_map = read_weight_map(&model_dir);
    let keys = [
        format!("model.language_model.layers.{layer}.mlp.gate.weight"),
        format!("model.language_model.layers.{layer}.mlp.shared_expert.gate_proj.weight"),
        format!("model.language_model.layers.{layer}.mlp.shared_expert.up_proj.weight"),
        format!("model.language_model.layers.{layer}.mlp.shared_expert.down_proj.weight"),
        format!("model.language_model.layers.{layer}.mlp.shared_expert_gate.weight"),
        format!("model.language_model.layers.{layer}.mlp.experts.gate_up_proj"),
        format!("model.language_model.layers.{layer}.mlp.experts.down_proj"),
    ];
    let mut shards: Vec<String> = keys
        .iter()
        .filter_map(|k| weight_map.get(k).cloned())
        .collect();
    shards.sort();
    shards.dedup();

    let mut tensors = std::collections::HashMap::new();
    for shard in &shards {
        tensors.extend(loader::load_safetensors(
            &model_dir.join(shard),
            Device::Cpu,
        ));
    }
    let mut take = |name: &str| -> Tensor {
        tensors
            .remove(name)
            .unwrap_or_else(|| panic!("missing tensor {name}"))
            .to_device(Device::Cuda(device))
    };
    let p = format!("model.language_model.layers.{layer}.mlp");
    let router_w = take(&format!("{p}.gate.weight"))
        .transpose(0, 1)
        .contiguous();
    let shared_gate_w = take(&format!("{p}.shared_expert.gate_proj.weight"))
        .transpose(0, 1)
        .contiguous();
    let shared_up_w = take(&format!("{p}.shared_expert.up_proj.weight"))
        .transpose(0, 1)
        .contiguous();
    let shared_down_w = take(&format!("{p}.shared_expert.down_proj.weight"))
        .transpose(0, 1)
        .contiguous();
    let shared_router_w = take(&format!("{p}.shared_expert_gate.weight"))
        .transpose(0, 1)
        .contiguous();
    let expert_gate_up = take(&format!("{p}.experts.gate_up_proj")).contiguous();
    let expert_down = take(&format!("{p}.experts.down_proj")).contiguous();

    let input = make_deterministic_input(seq_len, config.hidden()).to_device(Device::Cuda(device));
    let router_logits = matmul(&input, &router_w, GemmBackend::CuBlas);
    let (topk_ids, topk_weights) =
        moe_topk_softmax(&router_logits, spec.num_experts, spec.experts_per_token);

    let shared_gate = matmul(&input, &shared_gate_w, GemmBackend::CuBlas);
    let shared_up = matmul(&input, &shared_up_w, GemmBackend::CuBlas);
    let shared_act = mul(&silu(&shared_gate), &shared_up);
    let shared_down = matmul(&shared_act, &shared_down_w, GemmBackend::CuBlas);
    let shared_router = sigmoid(&matmul(&input, &shared_router_w, GemmBackend::CuBlas));
    let shared_out = row_scale_bf16(&shared_down, &shared_router);

    let gate_up = moe_sparse_gemv_bf16(
        &input,
        &expert_gate_up,
        &topk_ids,
        spec.experts_per_token,
        0,
        spec.num_experts,
        false,
    );
    let gate = gate_up.narrow(2, 0, spec.moe_intermediate).contiguous();
    let up = gate_up
        .narrow(2, spec.moe_intermediate, spec.moe_intermediate)
        .contiguous();
    let routed_act = mul(&silu(&gate), &up);
    let down = moe_sparse_gemv_bf16(
        &routed_act.reshape(&[seq_len * spec.experts_per_token, spec.moe_intermediate]),
        &expert_down,
        &topk_ids,
        spec.experts_per_token,
        0,
        spec.num_experts,
        true,
    );
    let routed_out = moe_weighted_sum_sparse(&down, &topk_ids, &topk_weights, 0, spec.num_experts);
    let moe_out = add(&routed_out, &shared_out);

    println!("layer={layer}");
    println!("input_shape={:?}", input.shape());
    println!("router_logits_shape={:?}", router_logits.shape());
    println!("topk_ids_shape={:?}", topk_ids.shape());
    println!("topk_weights_shape={:?}", topk_weights.shape());
    println!("shared_gate_shape={:?}", shared_gate.shape());
    println!("shared_up_shape={:?}", shared_up.shape());
    println!("shared_act_shape={:?}", shared_act.shape());
    println!("shared_down_shape={:?}", shared_down.shape());
    println!("shared_router_shape={:?}", shared_router.shape());
    println!("shared_out_shape={:?}", shared_out.shape());
    println!("gate_up_shape={:?}", gate_up.shape());
    println!("routed_act_shape={:?}", routed_act.shape());
    println!("down_shape={:?}", down.shape());
    println!("routed_out_shape={:?}", routed_out.shape());
    println!("moe_out_shape={:?}", moe_out.shape());
    println!("sample_router={:?}", sample_bf16(&router_logits, 8));
    println!(
        "sample_topk_ids={:?}",
        sample_i32_raw(&topk_ids, spec.experts_per_token)
    );
    println!(
        "sample_topk_weights={:?}",
        sample_f32(&topk_weights, spec.experts_per_token)
    );
    println!("sample_shared_router={:?}", sample_bf16(&shared_router, 4));
    println!("sample_shared_out={:?}", sample_bf16(&shared_out, 8));
    println!("sample_routed_out={:?}", sample_bf16(&routed_out, 8));
    println!("sample_moe_out={:?}", sample_bf16(&moe_out, 8));
}

fn read_weight_map(model_dir: &std::path::Path) -> std::collections::HashMap<String, String> {
    let index_path = model_dir.join("model.safetensors.index.json");
    let text = std::fs::read_to_string(&index_path)
        .unwrap_or_else(|e| panic!("failed to read {}: {e}", index_path.display()));
    let index: serde_json::Value = serde_json::from_str(&text)
        .unwrap_or_else(|e| panic!("failed to parse {}: {e}", index_path.display()));
    index
        .get("weight_map")
        .and_then(|v| v.as_object())
        .expect("index must contain weight_map")
        .iter()
        .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
        .collect()
}

fn make_deterministic_input(seq_len: usize, hidden: usize) -> Tensor {
    let mut data = Vec::with_capacity(seq_len * hidden);
    for i in 0..seq_len * hidden {
        let x = ((i % 97) as f32 - 48.0) / 128.0;
        data.push(bf16::from_f32(x));
    }
    Tensor::from_slice(&data, &[seq_len, hidden])
}

fn sample_bf16(t: &Tensor, n: usize) -> Vec<f32> {
    let cpu = t.to_device(Device::Cpu);
    cpu.as_slice::<bf16>()
        .iter()
        .take(n)
        .map(|v| v.to_f32())
        .collect()
}

fn sample_f32(t: &Tensor, n: usize) -> Vec<f32> {
    let cpu = t.to_device(Device::Cpu);
    cpu.as_slice::<f32>().iter().take(n).copied().collect()
}

fn sample_i32_raw(t: &Tensor, n: usize) -> Vec<i32> {
    let cpu = t.to_device(Device::Cpu);
    let bytes = cpu.as_raw_bytes();
    (0..n)
        .map(|i| {
            let start = i * 4;
            i32::from_ne_bytes(bytes[start..start + 4].try_into().unwrap())
        })
        .collect()
}