use std::{
    collections::HashMap,
    sync::Arc,
    time::{Duration, Instant},
};

use capitalize::Capitalize;
use common_display::tree::TreeDisplay;
use common_error::DaftResult;
use common_metrics::{
    ops::{NodeCategory, NodeInfo, NodeType},
    snapshot::StatSnapshotImpl,
};
use common_runtime::{OrderingAwareJoinSet, get_compute_pool_num_threads, get_compute_runtime};
use daft_local_plan::LocalNodeContext;
use daft_logical_plan::stats::StatsState;
use daft_micropartition::MicroPartition;
use opentelemetry::metrics::Meter;
use snafu::ResultExt;
use tracing::info_span;

use crate::{
    ExecutionRuntimeContext, ExecutionTaskSpawner, OperatorControlFlow, OperatorOutput,
    PipelineExecutionSnafu,
    buffer::RowBasedBuffer,
    channel::{Receiver, Sender, create_channel},
    dynamic_batching::{BatchManager, BatchingStrategy},
    pipeline::{BuilderContext, MorselSizeRequirement, NodeName, PipelineNode},
    runtime_stats::{DefaultRuntimeStats, RuntimeStats, RuntimeStatsManagerHandle},
};
#[cfg_attr(debug_assertions, derive(Debug))]
pub enum IntermediateOperatorResult {
    NeedMoreInput(Option<Arc<MicroPartition>>),
    #[allow(dead_code)]
    HasMoreOutput {
        input: Arc<MicroPartition>,
        output: Arc<MicroPartition>,
    },
}

pub(crate) type IntermediateOpExecuteResult<Op> = OperatorOutput<
    DaftResult<(
        <Op as IntermediateOperator>::State,
        IntermediateOperatorResult,
    )>,
>;
pub(crate) trait IntermediateOperator: Send + Sync {
    type State: Send + Sync + Unpin;
    type BatchingStrategy: BatchingStrategy + 'static;
    fn execute(
        &self,
        input: Arc<MicroPartition>,
        state: Self::State,
        task_spawner: &ExecutionTaskSpawner,
    ) -> IntermediateOpExecuteResult<Self>;
    fn name(&self) -> NodeName;
    fn op_type(&self) -> NodeType;
    fn multiline_display(&self) -> Vec<String>;
    fn make_state(&self) -> Self::State;
    fn make_runtime_stats(&self, meter: &Meter, node_info: &NodeInfo) -> Arc<dyn RuntimeStats> {
        Arc::new(DefaultRuntimeStats::new(meter, node_info))
    }
    /// The maximum number of concurrent workers that can be spawned for this operator.
    /// Each worker will has its own IntermediateOperatorState.
    /// This method should be overridden if the operator needs to limit the number of concurrent workers, i.e. UDFs with resource requests.
    fn max_concurrency(&self) -> usize {
        get_compute_pool_num_threads()
    }

    fn morsel_size_requirement(&self) -> Option<MorselSizeRequirement> {
        None
    }

    fn batching_strategy(&self) -> DaftResult<Self::BatchingStrategy>;
}

pub struct IntermediateNode<Op: IntermediateOperator> {
    intermediate_op: Arc<Op>,
    child: Box<dyn PipelineNode>,
    runtime_stats: Arc<dyn RuntimeStats>,
    plan_stats: StatsState,
    morsel_size_requirement: MorselSizeRequirement,
    node_info: Arc<NodeInfo>,
}

type StateId = usize;

struct ExecutionTaskResult<S> {
    state_id: StateId,
    state: S,
    result: IntermediateOperatorResult,
    elapsed: Duration,
}

struct ExecutionContext<Op: IntermediateOperator> {
    op: Arc<Op>,
    task_spawner: ExecutionTaskSpawner,
    task_set: OrderingAwareJoinSet<DaftResult<ExecutionTaskResult<Op::State>>>,
    state_pool: HashMap<StateId, Op::State>,
    output_sender: Sender<Arc<MicroPartition>>,
    batch_manager: Arc<BatchManager<Op::BatchingStrategy>>,
    runtime_stats: Arc<dyn RuntimeStats>,
}

// ========== IntermediateNode Implementation ==========

impl<Op: IntermediateOperator + 'static> IntermediateNode<Op> {
    pub(crate) fn new(
        intermediate_op: Arc<Op>,
        child: Box<dyn PipelineNode>,
        plan_stats: StatsState,
        ctx: &BuilderContext,
        context: &LocalNodeContext,
    ) -> Self {
        let name: Arc<str> = intermediate_op.name().into();
        let info = ctx.next_node_info(
            name,
            intermediate_op.op_type(),
            NodeCategory::Intermediate,
            context,
        );
        let runtime_stats = intermediate_op.make_runtime_stats(&ctx.meter, &info);
        let morsel_size_requirement = intermediate_op
            .morsel_size_requirement()
            .unwrap_or_default();
        Self {
            intermediate_op,
            child,
            runtime_stats,
            plan_stats,
            morsel_size_requirement,
            node_info: Arc::new(info),
        }
    }

    pub(crate) fn boxed(self) -> Box<dyn PipelineNode> {
        Box::new(self)
    }

    // ========== Helper Functions ==========

    fn spawn_execution_task(
        ctx: &mut ExecutionContext<Op>,
        input: Arc<MicroPartition>,
        state: Op::State,
        state_id: StateId,
    ) {
        let op = ctx.op.clone();
        let task_spawner = ctx.task_spawner.clone();
        ctx.task_set.spawn(async move {
            let now = Instant::now();
            let (new_state, result) = op.execute(input, state, &task_spawner).await??;
            let elapsed = now.elapsed();

            Ok(ExecutionTaskResult {
                state_id,
                state: new_state,
                result,
                elapsed,
            })
        });
    }

    async fn handle_task_completion(
        result: ExecutionTaskResult<Op::State>,
        ctx: &mut ExecutionContext<Op>,
    ) -> DaftResult<OperatorControlFlow> {
        match result.result {
            IntermediateOperatorResult::NeedMoreInput(Some(mp)) => {
                // Record execution stats
                ctx.batch_manager.record_execution_stats(
                    ctx.runtime_stats.as_ref(),
                    mp.len(),
                    result.elapsed,
                );

                // Send output
                ctx.runtime_stats.add_rows_out(mp.len() as u64);
                if ctx.output_sender.send(mp).await.is_err() {
                    return Ok(OperatorControlFlow::Break);
                }

                // Return state to pool
                ctx.state_pool.insert(result.state_id, result.state);
            }
            IntermediateOperatorResult::NeedMoreInput(None) => {
                // No output, just return state to pool
                ctx.state_pool.insert(result.state_id, result.state);
            }
            IntermediateOperatorResult::HasMoreOutput { input, output } => {
                // Record execution stats
                ctx.batch_manager.record_execution_stats(
                    ctx.runtime_stats.as_ref(),
                    output.len(),
                    result.elapsed,
                );

                // Send output
                ctx.runtime_stats.add_rows_out(output.len() as u64);
                if ctx.output_sender.send(output).await.is_err() {
                    return Ok(OperatorControlFlow::Break);
                }

                // Spawn another execution with same input and state (don't return state)
                Self::spawn_execution_task(ctx, input, result.state, result.state_id);
            }
        }
        Ok(OperatorControlFlow::Continue)
    }

    fn spawn_ready_batches(
        buffer: &mut RowBasedBuffer,
        ctx: &mut ExecutionContext<Op>,
    ) -> DaftResult<()> {
        // Check buffer for ready batches and spawn tasks while states available
        while !ctx.state_pool.is_empty()
            && let Some(batch) = buffer.next_batch_if_ready()?
        {
            let state_id = *ctx
                .state_pool
                .keys()
                .next()
                .expect("State pool should have states when has_available() returns true");
            let state = ctx
                .state_pool
                .remove(&state_id)
                .expect("State pool should have states when has_available() returns true");

            Self::spawn_execution_task(ctx, batch, state, state_id);
        }
        Ok(())
    }

    async fn process_input(
        node_id: usize,
        mut receiver: Receiver<Arc<MicroPartition>>,
        ctx: &mut ExecutionContext<Op>,
        stats_manager: &RuntimeStatsManagerHandle,
    ) -> DaftResult<()> {
        let (lower, upper) = ctx.batch_manager.initial_requirements().values();
        let mut buffer = RowBasedBuffer::new(lower, upper);
        let mut input_closed = false;
        let mut node_initialized = false;
        // Main processing loop
        while !input_closed || !ctx.task_set.is_empty() {
            tokio::select! {
                biased;

                // Branch 1: Join completed task (only if tasks exist)
                Some(join_result) = ctx.task_set.join_next(), if !ctx.task_set.is_empty() => {
                    let control = Self::handle_task_completion(join_result??, ctx).await?;
                    if !control.should_continue() {
                        return Ok(()); // Break
                    }

                    // After completing a task, update bounds and try to spawn more tasks
                    let new_requirements = ctx.batch_manager.calculate_batch_size();
                    let (lower, upper) = new_requirements.values();
                    buffer.update_bounds(lower, upper);

                    Self::spawn_ready_batches(&mut buffer, ctx)?;
                }

                // Branch 2: Receive input (only if states available and receiver open)
                morsel = receiver.recv(), if !ctx.state_pool.is_empty() && !input_closed => {
                    match morsel {
                        Some(morsel) => {
                            if !node_initialized {
                                stats_manager.activate_node(node_id);
                                node_initialized = true;
                            }
                            ctx.runtime_stats.add_rows_in(morsel.len() as u64);
                            buffer.push(morsel);
                            Self::spawn_ready_batches(&mut buffer, ctx)?;
                        }
                        None => {
                            input_closed = true;
                        }
                    }
                }
            }
        }

        // After loop exits, verify invariants
        debug_assert_eq!(ctx.task_set.len(), 0, "TaskSet should be empty after loop");
        debug_assert!(input_closed, "Receiver should be closed after loop");

        // Handle remaining buffered data
        if let Some(last_batch) = buffer.pop_all()? {
            // Since task_set is empty, all states should be back in the pool
            let state_id = *ctx
                .state_pool
                .keys()
                .next()
                .expect("State pool should have states after all tasks completed");
            let state = ctx
                .state_pool
                .remove(&state_id)
                .expect("State pool should have states after all tasks completed");

            Self::spawn_execution_task(ctx, last_batch, state, state_id);

            // Wait for final task to complete
            while let Some(join_result) = ctx.task_set.join_next().await {
                let control = Self::handle_task_completion(join_result??, ctx).await?;
                if !control.should_continue() {
                    return Ok(()); // Break
                }
            }
        }

        debug_assert_eq!(ctx.task_set.len(), 0, "TaskSet should be empty after loop");

        Ok(())
    }
}

impl<Op: IntermediateOperator + 'static> TreeDisplay for IntermediateNode<Op> {
    fn id(&self) -> String {
        self.node_id().to_string()
    }

    fn display_as(&self, level: common_display::DisplayLevel) -> String {
        use std::fmt::Write;
        let mut display = String::new();

        use common_display::DisplayLevel;
        match level {
            DisplayLevel::Compact => {
                writeln!(display, "{}", self.intermediate_op.name()).unwrap();
            }
            level => {
                let multiline_display = self.intermediate_op.multiline_display().join("\n");
                writeln!(display, "{}", multiline_display).unwrap();
                if let StatsState::Materialized(stats) = &self.plan_stats {
                    writeln!(display, "Stats = {}", stats).unwrap();
                }
                writeln!(display, "Batch Size = {}", self.morsel_size_requirement).unwrap();
                if matches!(level, DisplayLevel::Verbose) {
                    writeln!(display).unwrap();
                    let rt_result = self.runtime_stats.snapshot();
                    for (name, value) in rt_result.to_stats() {
                        writeln!(display, "{} = {}", name.as_ref().capitalize(), value).unwrap();
                    }
                }
            }
        }
        display
    }

    fn repr_json(&self) -> serde_json::Value {
        let children: Vec<serde_json::Value> = self
            .get_children()
            .iter()
            .map(|child| child.repr_json())
            .collect();

        let mut json = serde_json::json!({
            "id": self.node_id(),
            "category": "Intermediate",
            "type": self.intermediate_op.op_type().to_string(),
            "name": self.name(),
            "children": children,
        });

        if let StatsState::Materialized(stats) = &self.plan_stats {
            json["approx_stats"] = serde_json::json!(stats);
        }

        json
    }

    fn get_children(&self) -> Vec<&dyn TreeDisplay> {
        vec![self.child.as_tree_display()]
    }
}

impl<Op: IntermediateOperator + 'static> PipelineNode for IntermediateNode<Op> {
    fn children(&self) -> Vec<&dyn PipelineNode> {
        vec![self.child.as_ref()]
    }

    fn boxed_children(&self) -> Vec<&Box<dyn PipelineNode>> {
        vec![&self.child]
    }

    fn name(&self) -> Arc<str> {
        self.node_info.name.clone()
    }

    fn propagate_morsel_size_requirement(
        &mut self,
        downstream_requirement: MorselSizeRequirement,
        default_requirement: MorselSizeRequirement,
    ) {
        let operator_morsel_size_requirement = self.intermediate_op.morsel_size_requirement();
        let combined_morsel_size_requirement = MorselSizeRequirement::combine_requirements(
            operator_morsel_size_requirement,
            downstream_requirement,
        );
        self.morsel_size_requirement = combined_morsel_size_requirement;
        self.child.propagate_morsel_size_requirement(
            combined_morsel_size_requirement,
            default_requirement,
        );
    }

    fn start(
        self: Box<Self>,
        maintain_order: bool,
        runtime_handle: &mut ExecutionRuntimeContext,
    ) -> crate::Result<Receiver<Arc<MicroPartition>>> {
        let node_id = self.node_id();
        let name = self.name();

        // 1. Start children and wrap receivers
        let child_result_receiver = self.child.start(maintain_order, runtime_handle)?;
        // 2. Setup
        let op = self.intermediate_op.clone();
        let (destination_sender, destination_receiver) = create_channel(1);

        // 3. Create task spawner
        let compute_runtime = get_compute_runtime();
        let task_spawner = ExecutionTaskSpawner::new(
            compute_runtime,
            runtime_handle.memory_manager(),
            self.runtime_stats.clone(),
            info_span!("IntermediateOp::execute"),
        );

        // 4. Spawn process_input task
        let stats_manager = runtime_handle.stats_manager();
        let runtime_stats = self.runtime_stats.clone();
        runtime_handle.spawn(
            async move {
                // Initialize state pool with max_concurrency states
                let state_pool = (0..op.max_concurrency())
                    .map(|i| (i, op.make_state()))
                    .collect();

                // Create batch manager and task set
                let batch_manager = Arc::new(BatchManager::new(op.batching_strategy().context(
                    PipelineExecutionSnafu {
                        node_name: op.name().to_string(),
                    },
                )?));
                let task_set = OrderingAwareJoinSet::new(maintain_order);

                // Process each child receiver sequentially
                let mut ctx = ExecutionContext {
                    op,
                    task_spawner,
                    task_set,
                    state_pool,
                    output_sender: destination_sender,
                    batch_manager,
                    runtime_stats,
                };
                Self::process_input(node_id, child_result_receiver, &mut ctx, &stats_manager)
                    .await?;

                // Finalize node after processing completes
                stats_manager.finalize_node(node_id);

                Ok(())
            },
            &name,
        );

        Ok(destination_receiver)
    }

    fn as_tree_display(&self) -> &dyn TreeDisplay {
        self
    }
    fn node_id(&self) -> usize {
        self.node_info.id
    }

    fn node_info(&self) -> Arc<NodeInfo> {
        self.node_info.clone()
    }

    fn runtime_stats(&self) -> Arc<dyn RuntimeStats> {
        self.runtime_stats.clone()
    }
}
