use core::ops::ControlFlow;

use miden_mast_package::debug_info::DebugSourceNodeId;

use crate::{
    BaseHost, BreakReason, Stopper,
    continuation_stack::{Continuation, ContinuationStack},
    errors::PackageSourceDebugContext,
    execution::{
        ExecutionState, InternalBreakReason, execute_op, finalize_clock_cycle_with_continuation,
        finalize_clock_cycle_with_continuation_and_op_helpers,
    },
    mast::{BasicBlockNode, ExecutableMastForest, MastNodeId},
    operation::Operation,
    option_map_break_reason,
    processor::Processor,
    tracer::Tracer,
};

// BASIC BLOCK PROCESSING
// ================================================================================================

/// Execute the given basic block node.
#[inline]
pub(super) fn execute_basic_block_node_from_start<P, H, S, T, F>(
    state: &mut ExecutionState<'_, P, H, S, T, F>,
    basic_block_node: &BasicBlockNode,
    node_id: MastNodeId,
    current_forest: &F,
) -> ControlFlow<InternalBreakReason<F>>
where
    P: Processor,
    H: BaseHost,
    S: Stopper<Processor = P, Forest = F>,
    T: Tracer<Processor = P, Forest = F>,
    F: ExecutableMastForest + Clone,
{
    state.tracer.start_clock_cycle(
        state.processor,
        Continuation::StartNode(node_id),
        state.continuation_stack,
        current_forest,
    );

    // Finalize the clock cycle corresponding to the SPAN operation.
    let source_node_id = state.current_source_node_id();
    finalize_clock_cycle_with_continuation(
        state.processor,
        state.tracer,
        state.stopper,
        state.continuation_stack,
        || {
            Some((
                Continuation::ResumeBasicBlock {
                    node_id,
                    batch_index: 0,
                    op_idx_in_batch: 0,
                },
                source_node_id,
            ))
        },
        current_forest,
    )
    .map_break(InternalBreakReason::from)?;

    // Execute the first batch separately, since `execute_basic_block_node_from_batch` executes
    // starting from the RESPAN preceding the batch (and there is no such RESPAN before the first
    // batch).
    if !basic_block_node.op_batches().is_empty() {
        execute_op_batch(state, basic_block_node, node_id, 0, 0, 0, current_forest)?;
    }

    // Execute the rest of the batches.
    execute_basic_block_node_from_batch(state, basic_block_node, node_id, 1, current_forest)
}

/// Executes the give basic block node starting from the specified operation index within the
/// specified batch.
#[inline]
pub(super) fn execute_basic_block_node_from_op_idx<P, H, S, T, F>(
    state: &mut ExecutionState<'_, P, H, S, T, F>,
    basic_block_node: &BasicBlockNode,
    node_id: MastNodeId,
    start_batch_index: usize,
    start_op_idx_in_batch: usize,
    current_forest: &F,
) -> ControlFlow<InternalBreakReason<F>>
where
    P: Processor,
    H: BaseHost,
    S: Stopper<Processor = P, Forest = F>,
    T: Tracer<Processor = P, Forest = F>,
    F: ExecutableMastForest + Clone,
{
    let batch_offset_in_block = basic_block_node
        .op_batches()
        .iter()
        .take(start_batch_index)
        .map(|batch| batch.ops().len())
        .sum();

    // Finish executing the specified batch from the given op index
    execute_op_batch(
        state,
        basic_block_node,
        node_id,
        start_batch_index,
        start_op_idx_in_batch,
        batch_offset_in_block,
        current_forest,
    )?;

    // Execute the rest of the batches
    execute_basic_block_node_from_batch(
        state,
        basic_block_node,
        node_id,
        start_batch_index + 1,
        current_forest,
    )
}

/// Executes the give basic block node starting from the RESPAN preceding the specified batch.
#[inline]
pub(super) fn execute_basic_block_node_from_batch<P, H, S, T, F>(
    state: &mut ExecutionState<'_, P, H, S, T, F>,
    basic_block_node: &BasicBlockNode,
    node_id: MastNodeId,
    start_batch_index: usize,
    current_forest: &F,
) -> ControlFlow<InternalBreakReason<F>>
where
    P: Processor,
    H: BaseHost,
    S: Stopper<Processor = P, Forest = F>,
    T: Tracer<Processor = P, Forest = F>,
    F: ExecutableMastForest + Clone,
{
    let mut batch_offset_in_block = basic_block_node
        .op_batches()
        .iter()
        .take(start_batch_index)
        .map(|batch| batch.ops().len())
        .sum();

    for (batch_index, op_batch) in
        basic_block_node.op_batches().iter().enumerate().skip(start_batch_index)
    {
        {
            // Start clock cycle corresponding to the RESPAN operation before the batch.
            state.tracer.start_clock_cycle(
                state.processor,
                Continuation::Respan { node_id, batch_index },
                state.continuation_stack,
                current_forest,
            );

            // Finalize the clock cycle corresponding to the RESPAN operation.
            //
            // Note: in the continuation closure, the continuation encodes resuming from the start
            // of the batch *after* the RESPAN operation. This is because the continuation encodes
            // what happens *after* the clock is incremented. For example, if we were to put a
            // `Continuation::Respan` here instead, and execution was stopped after this RESPAN,
            // then the next call to `Processor::execute_impl()` would re-execute the RESPAN.
            let source_node_id = state.current_source_node_id();
            finalize_clock_cycle_with_continuation(
                state.processor,
                state.tracer,
                state.stopper,
                state.continuation_stack,
                || {
                    Some((
                        Continuation::ResumeBasicBlock { node_id, batch_index, op_idx_in_batch: 0 },
                        source_node_id,
                    ))
                },
                current_forest,
            )
            .map_break(InternalBreakReason::from)?;
        }

        // Execute the batch.
        execute_op_batch(
            state,
            basic_block_node,
            node_id,
            batch_index,
            0,
            batch_offset_in_block,
            current_forest,
        )?;
        batch_offset_in_block += op_batch.ops().len();
    }

    finish_basic_block(state, node_id, current_forest).map_break(InternalBreakReason::from)
}

/// Execute the finish phase of a basic block node.
#[inline]
pub(super) fn finish_basic_block<P, H, S, T, F>(
    state: &mut ExecutionState<'_, P, H, S, T, F>,
    node_id: MastNodeId,
    current_forest: &F,
) -> ControlFlow<BreakReason<F>>
where
    P: Processor,
    H: BaseHost,
    S: Stopper<Processor = P, Forest = F>,
    T: Tracer<Processor = P, Forest = F>,
    F: ExecutableMastForest + Clone,
{
    state.tracer.start_clock_cycle(
        state.processor,
        Continuation::FinishBasicBlock(node_id),
        state.continuation_stack,
        current_forest,
    );

    // Finalize the clock cycle corresponding to the END operation.
    finalize_clock_cycle_with_continuation(
        state.processor,
        state.tracer,
        state.stopper,
        state.continuation_stack,
        || None,
        current_forest,
    )
}

// HELPERS
// ================================================================================================

/// Executes a single operation batch within a basic block node, starting from the operation
/// index `start_op_idx`.
#[inline(always)]
fn execute_op_batch<P, H, S, T, F>(
    state: &mut ExecutionState<'_, P, H, S, T, F>,
    basic_block: &BasicBlockNode,
    node_id: MastNodeId,
    batch_index: usize,
    start_op_idx: usize,
    batch_offset_in_block: usize,
    current_forest: &F,
) -> ControlFlow<InternalBreakReason<F>>
where
    P: Processor,
    H: BaseHost,
    S: Stopper<Processor = P, Forest = F>,
    T: Tracer<Processor = P, Forest = F>,
    F: ExecutableMastForest + Clone,
{
    let batch = option_map_break_reason(
        basic_block.op_batches().get(batch_index),
        "batch index out of range for basic block",
    )
    .map_break(InternalBreakReason::from)?;

    // Execute operations in the batch one by one
    for (op_idx_in_batch, op) in batch.ops().iter().enumerate().skip(start_op_idx) {
        let op_idx_in_block = batch_offset_in_block + op_idx_in_batch;

        state.tracer.start_clock_cycle(
            state.processor,
            Continuation::ResumeBasicBlock { node_id, batch_index, op_idx_in_batch },
            state.continuation_stack,
            current_forest,
        );

        // Execute the operation.
        let operation_helpers = match op {
            Operation::Emit => {
                // This is a sans-IO point: we cannot proceed with handling the Emit operation,
                // since some processors need this to be done asynchronously. Thus, we break
                // here and make the implementing processor handle the loading in the outer
                // execution loop. When done, the processor *must* call
                // `finish_emit_op_execution()` below for execution to proceed properly.
                return ControlFlow::Break(InternalBreakReason::Emit {
                    op_idx: op_idx_in_block,
                    continuation: get_continuation_after_executing_operation(
                        basic_block,
                        node_id,
                        batch_index,
                        op_idx_in_batch,
                    ),
                    source_node_id: state.current_source_node_id(),
                });
            },
            _ => {
                // If the operation is not an Emit, we execute it normally.
                let source_debug_info = state.source_debug_info.clone();
                let package_source_context = source_debug_info.as_deref().map(|debug_info| {
                    PackageSourceDebugContext::new_optional(
                        debug_info,
                        state.current_source_node_id(),
                    )
                });
                match execute_op(
                    state.processor,
                    op,
                    op_idx_in_block,
                    state.host,
                    state.tracer,
                    package_source_context,
                ) {
                    Ok(operation_helpers) => operation_helpers,
                    Err(err) => {
                        return ControlFlow::Break(BreakReason::Err(err).into());
                    },
                }
            },
        };

        // Finalize the clock cycle corresponding to the operation.
        let source_node_id = state.current_source_node_id();
        finalize_clock_cycle_with_continuation_and_op_helpers(
            state.processor,
            state.tracer,
            state.stopper,
            state.continuation_stack,
            || {
                Some((
                    get_continuation_after_executing_operation(
                        basic_block,
                        node_id,
                        batch_index,
                        op_idx_in_batch,
                    ),
                    source_node_id,
                ))
            },
            operation_helpers,
            current_forest,
        )
        .map_break(InternalBreakReason::from)?;
    }

    ControlFlow::Continue(())
}

/// Given the current operation being executed within a basic block, returns the appropriate
/// continuation to add to the continuation stack if execution is stopped right after execution the
/// operation (node_id, batch_index, op_idx_in_batch).
///
/// That is, `op_idx_in_batch` is the index of the operation that was just executed within the batch
/// `batch_index` of the basic block `basic_block_node`.
#[inline(always)]
fn get_continuation_after_executing_operation<F>(
    basic_block_node: &BasicBlockNode,
    node_id: MastNodeId,
    batch_index: usize,
    op_idx_in_batch: usize,
) -> Continuation<F> {
    let last_op_idx_in_batch = basic_block_node.op_batches()[batch_index].ops().len() - 1;
    let last_batch_idx_in_block = basic_block_node.num_op_batches() - 1;

    if op_idx_in_batch < last_op_idx_in_batch {
        // The operation that just executed was not the last one in the batch, so continue within
        // the same batch at the following operation
        Continuation::ResumeBasicBlock {
            node_id,
            batch_index,
            op_idx_in_batch: op_idx_in_batch + 1,
        }
    } else if batch_index < last_batch_idx_in_block {
        // The operation that just executed was the last one in the batch, but there are more
        // batches to execute in this basic block, so continue at the RESPAN before the next batch
        Continuation::Respan { node_id, batch_index: batch_index + 1 }
    } else {
        // The operation that just executed was the last one in the last batch, so finish the basic
        // block
        Continuation::FinishBasicBlock(node_id)
    }
}

// EXPORTS
// ================================================================================================

/// Function to be called after [`InternalBreakReason::Emit`] is handled. See the documentation of
/// that enum variant for more details.
pub fn finish_emit_op_execution<P, S, T, F>(
    post_emit_continuation: Continuation<F>,
    source_node_id: Option<DebugSourceNodeId>,
    processor: &mut P,
    continuation_stack: &mut ContinuationStack<F>,
    current_forest: &F,
    tracer: &mut T,
    stopper: &S,
) -> ControlFlow<BreakReason<F>>
where
    P: Processor,
    S: Stopper<Processor = P, Forest = F>,
    T: Tracer<Processor = P, Forest = F>,
    F: ExecutableMastForest + Clone,
{
    // When we enter here, the `continuation_stack` top contains the continuation to execute *after*
    // the basic block that contained the `Emit` operation (i.e. after all operations are executed,
    // and the finish phase of the basic block is complete). Hence, we need to add the
    // `post_emit_continuation` on top of the continuation stack so that execution resumes at the
    // operation right after the `Emit`.
    //
    // However, if the `stopper` stops execution in `finalize_clock_cycle_with_continuation()`, the
    // stopper will already include the `post_emit_continuation` in the break reason (which the
    // processor will then push onto the continuation stack). Hence, in this case, we do not need to
    // push the `post_emit_continuation` ourselves. In other words, *only if* the
    // `finalize_clock_cycle_with_continuation()` completes successfully do we need to push the
    // `post_emit_continuation` ourselves.

    finalize_clock_cycle_with_continuation(
        processor,
        tracer,
        stopper,
        continuation_stack,
        {
            let post_emit_continuation = post_emit_continuation.clone();
            || Some((post_emit_continuation, source_node_id))
        },
        current_forest,
    )?;

    continuation_stack.push_with_source_node_id(post_emit_continuation, source_node_id);

    ControlFlow::Continue(())
}