use std::sync::Arc;

use burn_ir::HandleContainer;

use crate::{FusionRuntime, NumOperations, Optimization, UnfusedOp, stream::Context};

/// Manage the execution of potentially multiple optimizations and operations out of order.
pub struct OrderedExecution<R: FusionRuntime> {
    operations: Vec<UnfusedOp<R>>,
    num_executed: usize,
    ordering: Option<Arc<Vec<usize>>>,
}

impl<R: FusionRuntime> OrderedExecution<R> {
    /// Returns the operation that can be executed without impacting the state of the execution.
    ///
    /// This is useful to implement fallback for optimizations.
    #[allow(clippy::borrowed_box)]
    pub fn operation_within_optimization(&self, index: usize) -> UnfusedOp<R> {
        match &self.ordering {
            Some(val) => {
                let index = val[index];
                self.operations[index].clone()
            }
            None => panic!("No ordering provided"),
        }
    }

    pub(crate) fn new(operations: Vec<UnfusedOp<R>>) -> Self {
        Self {
            operations,
            num_executed: 0,
            ordering: None,
        }
    }

    pub(crate) fn finish(mut self) -> (Vec<UnfusedOp<R>>, usize) {
        self.operations.drain(0..self.num_executed);
        (self.operations, self.num_executed)
    }

    pub(crate) fn execute_optimization(
        &mut self,
        optimization: &mut R::Optimization,
        context: &mut Context<R::FusionHandle>,
        ordering: Arc<Vec<usize>>,
    ) {
        if ordering.len() > self.operations.len() {
            panic!("Ordering is bigger than operations");
        }
        self.ordering = Some(ordering);
        let num_drained = optimization.len();
        optimization.execute(context, self);
        self.num_executed += num_drained;
    }

    pub(crate) fn execute_operations(
        &mut self,
        handles: &mut HandleContainer<R::FusionHandle>,
        ordering: &[usize],
    ) {
        self.num_executed += ordering.len();

        for id in ordering {
            let op = &self.operations[*id];
            op.execute(handles);
        }
    }
}
