use crate::{
    engine::{
        codegen::ir::{FuseArg, FuseBlockConfig, GlobalArgsLaunch, RefLayout},
        launch::runner::{TraceRunner, Vectorization},
    },
    optim::reduce_broadcasted::unit::{
        ElemwiseFuseBlockLaunch, ReduceFuseBlockLaunch, reduce_kernel_broadcasted,
    },
};
use cubecl::{
    Runtime,
    ir::{ElemType, FloatKind, StorageType},
    prelude::*,
    server::LaunchError,
};
use cubek::reduce::{
    ReduceDtypes, VectorizationMode,
    components::instructions::ReduceOperationConfig,
    launch::RoutineStrategy,
    output_vectorization_axis,
    routines::{
        BlueprintStrategy, GlobalReduceBlueprint, ReduceProblem, ReduceVectorSettings, Routine,
        unit::{UnitRoutine, UnitStrategy},
    },
};
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct ReduceBroadcastedFuseBlock {
    pub(crate) op: ReduceOperationConfig,
    pub(crate) input: FuseArg,
    pub(crate) output: FuseArg,
}

#[derive(new)]
pub struct FusedReduceBroadcastedLaunch<'a> {
    blocks: &'a Vec<ReduceBroadcastedFuseBlock>,
    reduce_axis: usize,
    // TODO: Support multiple strategies.
    _strategy: RoutineStrategy,
}

impl<R: Runtime> Vectorization<R> for FusedReduceBroadcastedLaunch<'_> {}

impl<R: Runtime> TraceRunner<R> for FusedReduceBroadcastedLaunch<'_> {
    type Error = LaunchError;

    fn run<'a>(
        &'a self,
        client: &'a ComputeClient<R>,
        inputs: GlobalArgsLaunch<R>,
        outputs: GlobalArgsLaunch<R>,
        configs: &'a [FuseBlockConfig],
    ) -> Result<(), Self::Error> {
        let routine = UnitRoutine;
        let first_config = &configs[0];

        let shape = match &first_config.ref_layout {
            RefLayout::Concrete(FuseArg::Output(..)) => {
                outputs.shape_ref(&first_config.ref_layout, first_config.rank)
            }
            _ => inputs.shape_ref(&first_config.ref_layout, first_config.rank),
        };

        let reduce_len = shape[self.reduce_axis];
        let reduce_count = shape.iter().product::<usize>() / reduce_len;
        let address_type = inputs
            .required_address_type()
            .max(outputs.required_address_type());

        let (blueprint, settings) = routine
            .prepare::<R>(
                client,
                ReduceProblem {
                    reduce_len,
                    reduce_count,
                    axis: self.reduce_axis,
                    dtypes: ReduceDtypes {
                        input: StorageType::Scalar(ElemType::Float(FloatKind::F32)),
                        output: StorageType::Scalar(ElemType::Float(FloatKind::F32)),
                        accumulation: StorageType::Scalar(ElemType::Float(FloatKind::F32)),
                    },
                    address_type,
                },
                ReduceVectorSettings {
                    vectorization_mode: VectorizationMode::Parallel,
                    vector_size_input: first_config.width,
                    vector_size_output: 1,
                },
                BlueprintStrategy::Inferred(UnitStrategy),
            )
            .unwrap();

        assert_eq!(blueprint.vectorization_mode, VectorizationMode::Parallel);

        let mut blocks = SequenceArg::new();
        let mut index = 0;

        for block in self.blocks {
            let arg = ReduceFuseBlockLaunch::new(
                block.op,
                configs[index].clone(),
                configs[index + 1].clone(),
                block.input.clone(),
                block.output.clone(),
                match blueprint.global {
                    GlobalReduceBlueprint::Unit(bpt) => bpt,
                    _ => panic!(),
                },
            );
            index += 2;
            blocks.push(arg);
        }

        let block_end = match configs.len() > index {
            true => ComptimeOptionArgs::Some(ElemwiseFuseBlockLaunch::new(
                configs.last().cloned().unwrap(),
            )),
            false => ComptimeOptionArgs::None,
        };

        let out_vec_axis = output_vectorization_axis(
            &inputs.strides_ref(&first_config.ref_layout, first_config.rank),
            self.reduce_axis,
            VectorizationMode::Parallel,
        );

        // TODO: Ensure parallel is selected.

        unsafe {
            reduce_kernel_broadcasted::launch_unchecked::<R>(
                client,
                settings.cube_count,
                settings.cube_dim,
                settings.address_type,
                inputs,
                outputs,
                self.reduce_axis,
                out_vec_axis,
                blocks,
                block_end,
            );
        }

        Ok(())
    }
}
