use cubecl_core::prelude::*;
use cubecl_core::server::Allocation;
use cubecl_core::{CubeElement, server};
use cubecl_std::CubeOptionArgs;

use crate::components::args::{TensorArgs, TensorInputsLaunch};
use crate::components::batch::BatchAttentionConfig;
use crate::components::batch::BatchAttentionFamily;
use crate::components::{AttentionElems, AttentionIdent, AvailableLineSizes};
use crate::components::{AttentionProblem, AttentionSelection};
use crate::kernels::Algorithm;
use crate::tests::test_utils::Sampleable;
use crate::tests::test_utils::TestPrecision;

#[derive(Debug)]
pub struct TensorRawParts<N: Numeric + CubeElement> {
    pub handle: server::Handle,
    pub shape: Vec<usize>,
    pub strides: Vec<usize>,
    pub original_data: Option<Vec<N>>,
}

enum TestInputType<T> {
    Random(u64),
    Data(Vec<T>),
    Zeros,
}

pub fn test_attention_algorithm<A, P, R>(
    client: ComputeClient<R>,
    problem: AttentionProblem,
    selection: AttentionSelection,
) where
    A: Algorithm,
    P: TestPrecision,
    R: Runtime,
{
    let query = tensor_raw_parts::<P::EG, R>(
        &client,
        &problem,
        AttentionIdent::Query,
        TestInputType::Random(12),
    );
    let key = tensor_raw_parts::<P::EG, R>(
        &client,
        &problem,
        AttentionIdent::Key,
        TestInputType::Random(34),
    );
    let value = tensor_raw_parts::<P::EG, R>(
        &client,
        &problem,
        AttentionIdent::Value,
        TestInputType::Random(56),
    );
    let mask = if problem.masked {
        Some(tensor_raw_parts::<P::EM, R>(
            &client,
            &problem,
            AttentionIdent::Mask,
            TestInputType::Random(78),
        ))
    } else {
        None
    };
    let out =
        tensor_raw_parts::<P::EG, R>(&client, &problem, AttentionIdent::Out, TestInputType::Zeros);

    test_attention_algorithm_raw::<A, P, R>(
        client, problem, selection, query, key, value, mask, out,
    );
}

#[allow(unused)]
pub fn test_attention_algorithm_explicit<A, P, R>(
    client: ComputeClient<R>,
    problem: AttentionProblem,
    selection: AttentionSelection,
    query_data: Vec<P::EG>,
    key_data: Vec<P::EG>,
    value_data: Vec<P::EG>,
    mask_data: Option<Vec<P::EM>>,
) where
    A: Algorithm,
    P: TestPrecision,
    R: Runtime,
{
    let query = tensor_raw_parts::<P::EG, R>(
        &client,
        &problem,
        AttentionIdent::Query,
        TestInputType::Data(query_data),
    );
    let key = tensor_raw_parts::<P::EG, R>(
        &client,
        &problem,
        AttentionIdent::Key,
        TestInputType::Data(key_data),
    );
    let value = tensor_raw_parts::<P::EG, R>(
        &client,
        &problem,
        AttentionIdent::Value,
        TestInputType::Data(value_data),
    );
    let mask = mask_data.map(|m| {
        tensor_raw_parts::<P::EM, R>(
            &client,
            &problem,
            AttentionIdent::Mask,
            TestInputType::Data(m),
        )
    });
    let out =
        tensor_raw_parts::<P::EG, R>(&client, &problem, AttentionIdent::Out, TestInputType::Zeros);

    test_attention_algorithm_raw::<A, P, R>(
        client, problem, selection, query, key, value, mask, out,
    );
}

fn test_attention_algorithm_raw<A, P, R>(
    client: ComputeClient<R>,
    problem: AttentionProblem,
    selection: AttentionSelection,
    query: TensorRawParts<P::EG>,
    key: TensorRawParts<P::EG>,
    value: TensorRawParts<P::EG>,
    mask: Option<TensorRawParts<P::EM>>,
    out: TensorRawParts<P::EG>,
) where
    A: Algorithm,
    P: TestPrecision,
    R: Runtime,
{
    let env = std::env::var("ATTENTION_TEST_MODE");
    let panic_on_launch_err = env.as_deref() == Ok("panic");

    let attention_elems = AttentionElems::new::<P::AP>();
    let line_sizes = {
        let ls = AvailableLineSizes::from_elem_types(
            &client,
            attention_elems.query_global.size(),
            attention_elems.mask.size(),
            attention_elems.out_global.size(),
        );
        let ls = A::filter_line_sizes(ls)
            .filter_with_tensor(AttentionIdent::Query, &query.strides, &query.shape)
            .filter_with_tensor(AttentionIdent::Key, &key.strides, &key.shape)
            .filter_with_tensor(AttentionIdent::Value, &value.strides, &value.shape)
            .filter_with_tensor(AttentionIdent::Out, &out.strides, &out.shape);

        if let Some(mask) = mask.as_ref() {
            ls.filter_with_tensor(AttentionIdent::Mask, &mask.strides, &mask.shape)
        } else {
            ls
        }
    }
    .pick_max()
    .unwrap();

    let config = match A::setup(&client, &problem, &selection, &line_sizes, &attention_elems) {
        Ok(config) => config,
        Err(err) => {
            let msg = format!("Can't launch the test: {err}");
            if panic_on_launch_err {
                panic!("{msg}");
            } else {
                println!("{msg}");
                return;
            }
        }
    };

    let cube_count_plan = config
        .hypercube_config()
        .cube_count_plan(&problem, &selection);

    let result = unsafe {
        A::BatchAttention::launch_unchecked::<TensorArgs, R>(
            &client,
            config.cube_dim(),
            cube_count_plan.resolve(),
            TensorInputsLaunch::new(
                TensorArg::from_raw_parts::<P::EG>(
                    &query.handle,
                    &query.strides,
                    &query.shape,
                    line_sizes.query,
                ),
                TensorArg::from_raw_parts::<P::EG>(
                    &key.handle,
                    &key.strides,
                    &key.shape,
                    line_sizes.key,
                ),
                TensorArg::from_raw_parts::<P::EG>(
                    &value.handle,
                    &value.strides,
                    &value.shape,
                    line_sizes.value,
                ),
                match mask.as_ref() {
                    Some(m) => CubeOptionArgs::Some(TensorArg::from_raw_parts::<P::EM>(
                        &m.handle,
                        &m.strides,
                        &m.shape,
                        line_sizes.mask,
                    )),
                    None => CubeOptionArgs::None,
                },
            ),
            TensorArg::from_raw_parts::<P::EG>(
                &out.handle,
                &out.strides,
                &out.shape,
                line_sizes.out,
            ),
            cube_count_plan.as_args(),
            config,
            &attention_elems,
        )
    };

    if let Err(err) = result {
        println!("Skipping the test with an execution error {err:?}");
        return;
    }

    P::assert_result(
        &query.original_data.unwrap(),
        &key.original_data.unwrap(),
        &value.original_data.unwrap(),
        mask.as_ref()
            .map(|m| m.original_data.as_ref().unwrap().as_slice()),
        &problem,
        &client,
        out.handle,
        &out.shape,
        &out.strides,
    );
}

fn tensor_raw_parts<T, R: Runtime>(
    client: &ComputeClient<R>,
    problem: &AttentionProblem,
    ident: AttentionIdent,
    input: TestInputType<T>,
) -> TensorRawParts<T>
where
    T: Numeric + CubeElement + Sampleable,
{
    let tensor_shape = problem.shape(ident);
    let original_data = match input {
        TestInputType::Random(seed) => {
            let handle = T::sample(client, &tensor_shape, seed);
            let data = client.read_one(handle.handle);
            let data = T::from_bytes(&data);
            data.to_owned()
        }
        TestInputType::Data(data) => {
            assert_eq!(
                data.len(),
                tensor_shape.iter().product::<usize>(),
                "Provided data length does not match tensor shape"
            );
            data
        }
        TestInputType::Zeros => vec![T::from_int(0); tensor_shape.len()],
    };

    let data_bytes = T::as_bytes(&original_data);
    let shape = tensor_shape.as_slice();
    let elem_size = T::type_size() as usize;

    let Allocation { handle, strides } =
        client.create_tensor_from_slice(data_bytes, shape, elem_size);

    TensorRawParts {
        handle,
        shape: tensor_shape.to_vec(),
        strides,
        original_data: Some(original_data),
    }
}

pub(crate) fn strides(problem: &AttentionProblem, ident: AttentionIdent) -> Vec<usize> {
    let shape = problem.shape(ident);

    let mut strides = vec![0; shape.len()];
    let mut acc = 1;
    for i in (0..shape.len()).rev() {
        strides[i] = acc;
        acc *= shape[i];
    }

    strides
}
