//! Executor module.
//!
//! The executor is located on the coordinator node in the cluster.
//! It collects all the intermediate results of the plan execution
//! in memory and executes the IR plan tree in the bottom-up manner.
//! It goes like this:
//!
//! 1. The executor collects all the motion nodes from the bottom layer.
//!    In theory all the motions in the same layer can be executed in parallel
//!    (this feature is yet to come).
//! 2. For every motion the executor:
//!    - inspects the IR sub-tree and detects the buckets to execute the query for.
//!    - builds a valid SQL query from the IR sub-tree.
//!    - performs map-reduce for that SQL query (we send it to the shards deduced from the buckets).
//!    - builds a virtual table with query results that correspond to the original motion.
//! 3. Moves to the next motion layer in the IR tree.
//! 4. For every motion the executor then:
//!    - links the virtual table results of the motion from the previous layer we depend on.
//!    - inspects the IR sub-tree and detects the buckets to execute the query.
//!    - builds a valid SQL query from the IR sub-tree.
//!    - performs map-reduce for that SQL query.
//!    - builds a virtual table with query results that correspond to the original motion.
//! 5. Repeats step 3 till we are done with motion layers.
//! 6. Executes the final IR top subtree and returns the final result to the user.
use crate::errors::{Action, Entity, SbroadError};
use crate::executor::engine::{Router, Vshard};
use crate::executor::ir::ExecutionPlan;
use crate::executor::vdbe::ExecutionInsight;
use crate::ir::bucket::{BucketSet, Buckets};
use crate::ir::node::block::{BlockOwned, MutBlock};
use crate::ir::node::expression::Expression;
use crate::ir::node::relational::{MutRelational, Relational};
use crate::ir::node::{AnonymousBlock, Insert, Motion, NodeId, Values, ValuesRow};
use crate::ir::transformation::redistribution::{MotionPolicy, Target};
use crate::ir::tree::traversal::{PostOrder, REL_CAPACITY};
use crate::ir::value::Value;
use crate::ir::{Plan, Slices};
use crate::BoundStatement;
use rmp::encode::write_str;
use smol_str::{format_smolstr, SmolStr};
use std::collections::HashMap;
use std::io::Write;
use std::rc::Rc;
use tarantool::msgpack;
use vdbe::{SqlError, SqlStmt};

pub mod bucket_discovery;
pub mod engine;
pub mod hash;
pub mod ir;
pub mod lru;
pub mod preemption;
pub mod protocol;
pub mod result;
pub mod vdbe;
pub mod vtable;

impl Plan {
    /// Apply optimization rules to the plan.
    ///
    /// # Errors
    /// - Failed to optimize the plan.
    pub fn optimize(self) -> Result<Self, SbroadError> {
        let top_id = self.get_top()?;
        self.optimize_subtree(top_id)
    }

    pub fn optimize_subtree(self, top_id: NodeId) -> Result<Self, SbroadError> {
        self.replace_in_operator_in_subtree(top_id)?
            .push_down_not_in_subtree(top_id)?
            // In the case if the query was not fully parameterized
            // and contains some constants, lets apply constant folding.
            .cast_constants()?
            .fold_boolean_tree()?
            .split_columns_in_subtree(top_id)?
            .set_dnf_in_subtree(top_id)?
            .derive_equalities_in_subtree(top_id)?
            .merge_tuples_in_subtree(top_id)?
            .add_motions_to_subtree(top_id)?
            .update_substring()?
            // After all transformations we can finally determine what parameters are unique.
            .mark_unique_parameters()
    }

    pub fn optimize_block(mut self) -> Result<Self, SbroadError> {
        // Eliminate motions from the plan so it can be executed locally within a block.
        fn eliminate_motions_in_subtree(
            plan: &mut Plan,
            top_id: NodeId,
        ) -> Result<(), SbroadError> {
            let top = plan.get_relation_node(top_id)?;
            if matches!(top, Relational::Insert(_)) {
                return eliminate_insert_motion(plan, top_id);
            }
            if top.is_dml() {
                // Ensure motion is local.
                let top_children = plan.get_relation_children(top_id)?;
                let Some(motion_id) = top_children.get(0) else {
                    return Ok(());
                };
                let motion = plan.get_relation_node(*motion_id)?;
                if !matches!(
                    motion,
                    Relational::Motion(Motion {
                        policy: MotionPolicy::Local,
                        ..
                    })
                ) {
                    return Ok(());
                }

                // Unlink the DML node's motion child.
                let motion_child_id = plan.get_rel_child(*motion_id, 0)?;
                match plan.get_mut_relation_node(top_id)? {
                    MutRelational::Update(update) => update.child = motion_child_id,
                    MutRelational::Delete(delete) => delete.child = Some(motion_child_id),
                    _ => {}
                }
            }

            Ok(())
        }

        // Eliminate the motion under an INSERT in a transaction.
        //
        // Expected shape: Insert -> Motion(Segment | LocalSegment) -> Values -> ValuesRow*.
        // Each ValuesRow cell at a sharding-key position must be a Constant or Parameter
        // (parameters become constants after binding). Other cells may hold arbitrary
        // expressions — they will be evaluated by the storage's local SQL.
        fn eliminate_insert_motion(plan: &mut Plan, insert_id: NodeId) -> Result<(), SbroadError> {
            // INSERT into a global table gets MotionPolicy::Full and is not supported in blocks;
            // leave the motion in place so the block optimizer reports it.
            if plan.dml_node_table(insert_id)?.is_global() {
                return Ok(());
            }

            let motion_id = plan.dml_child_id(insert_id)?;
            let Relational::Motion(Motion { policy, .. }) = plan.get_relation_node(motion_id)?
            else {
                return Ok(());
            };
            let (MotionPolicy::Segment(key) | MotionPolicy::LocalSegment(key)) = policy else {
                return Ok(());
            };

            let values_id = plan.get_rel_child(motion_id, 0)?;
            let Relational::Values(Values {
                children: values_rows,
                ..
            }) = plan.get_relation_node(values_id)?
            else {
                return Ok(());
            };

            // Collect cell positions in each ValuesRow's data tuple that correspond to
            // sharding-key columns.
            let key_positions: Vec<usize> = key
                .targets
                .iter()
                .filter_map(|t| match t {
                    Target::Reference(pos) => Some(*pos),
                    Target::Value(_) => None,
                })
                .collect();

            for row_id in values_rows {
                let Relational::ValuesRow(ValuesRow {
                    data, subqueries, ..
                }) = plan.get_relation_node(*row_id)?
                else {
                    panic!("Expected ValuesRow under Values, got {row_id:?}");
                };
                if !subqueries.is_empty() {
                    return Err(SbroadError::other(
                        "INSERT in transaction does not support subqueries in VALUES",
                    ));
                }
                let row_list = plan.get_row_list(*data)?;
                for pos in &key_positions {
                    let cell_id = *row_list
                        .get(*pos)
                        .expect("sharding-key position must exist in values row");
                    let cell = plan.get_expression_node(cell_id)?;
                    if !matches!(cell, Expression::Constant(_) | Expression::Parameter(_)) {
                        return Err(SbroadError::other(
                            "INSERT in transaction requires constant or parameter values for sharding-key columns"
                        ));
                    }
                }
            }

            // Re-point Insert's child directly to Values, dropping the Motion node.
            if let MutRelational::Insert(Insert { child, .. }) =
                plan.get_mut_relation_node(insert_id)?
            {
                *child = values_id;
            }

            Ok(())
        }

        let block_id = self.get_top()?;
        let block = self.get_block_node(block_id)?.get_block_owned();
        let BlockOwned::Anonymous(AnonymousBlock { mut statements, .. }) = block else {
            panic!("can't opmitize {block:?} as block");
        };

        // Optimize subtrees and eliminate motions.
        for stmt in &mut statements {
            let query_id = *stmt.get();
            // Remember keyword for the top node before motions are added.
            let keyword = self.as_keyword(query_id);

            // Set top so `optimize_subtree()` can update it when subtree's top changes.
            self.set_top(query_id).expect("top node can't be missed");
            self = self.optimize_subtree(*stmt.get())?;
            eliminate_motions_in_subtree(&mut self, *stmt.get())?;
            let new_top = self.get_top().expect("just set");
            *stmt.get_mut() = new_top;

            // Ensure no motions.
            if self.subtree_has_motions(new_top)? {
                return Err(SbroadError::Other(format_smolstr!(
                    "{keyword} query has motions which are not allowed in transactions"
                )));
            }
        }

        // Set block as the top node back.
        self.set_top(block_id).expect("block node can't be missed");

        // Update statements since as a result of optimizations each subtree can have a new top.
        if let MutBlock::Anonymous(AnonymousBlock {
            statements: ref mut plan_statements,
            ..
        }) = self.get_mut_block_node(block_id)?
        {
            *plan_statements = statements;
        }

        Ok(self)
    }

    fn subtree_has_motions(&self, top_id: NodeId) -> Result<bool, SbroadError> {
        let post_tree = PostOrder::new(|node| self.nodes.rel_iter(node), REL_CAPACITY);
        for node in post_tree.traverse_into_iter(top_id) {
            if let Relational::Motion(_) = self.get_relation_node(node.1)? {
                return Ok(true);
            }
        }
        Ok(false)
    }

    fn as_keyword(&self, top_id: NodeId) -> &'static str {
        let Ok(top) = self.get_relation_node(top_id) else {
            return "";
        };

        match top {
            Relational::Values(_) | Relational::ValuesRow(_) => "VALUES",
            Relational::Projection(_) | Relational::SelectWithoutScan(_) => "SELECT",
            Relational::ScanCte(_) | Relational::ScanSubQuery(_) => "FROM",
            Relational::ScanRelation(_) => "FROM",
            Relational::Selection(_) => "WHERE",
            Relational::UnionAll(_) => "UNION ALL",
            Relational::Union(_) => "UNION",
            Relational::Except(_) => "EXCEPT",
            Relational::Limit(_) => "LIMIT",
            Relational::Join(_) => "JOIN",
            Relational::Intersect(_) => "INTERSECT",
            Relational::Update(_) => "UPDATE",
            Relational::Delete(_) => "DELETE",
            Relational::Insert(_) => "INSERT",
            Relational::Having(_) => "HAVING",
            Relational::GroupBy(_) => "GROUP BY",
            Relational::OrderBy(_) => "ORDER BY",
            Relational::Motion(_) => "<MOTION>",
        }
    }
}

pub enum PortType {
    DispatchDql,
    DispatchDml,
    DispatchExplain,
    DispatchQueryPlan,
    ExecuteDql,
    ExecuteDml,
    ExecuteMiss,
}

pub trait Port<'p>: Write {
    fn add_mp(&mut self, data: &[u8]);

    fn process_stmt(
        &mut self,
        stmt: &mut SqlStmt,
        params: &[Value],
        max_vdbe: u64,
    ) -> Result<ExecutionInsight, SqlError>
    where
        Self: Sized;

    fn process_stmt_with_raw_params(
        &mut self,
        stmt: &mut SqlStmt,
        params: &[u8],
        max_vdbe: u64,
    ) -> Result<ExecutionInsight, SqlError>;

    fn iter(&self) -> impl Iterator<Item = &[u8]>
    where
        Self: Sized;

    fn set_type(&mut self, port_type: PortType);

    fn size(&self) -> u32;
}

/// Query to execute.
#[derive(Debug)]
pub struct ExecutingQuery<'a, C>
where
    C: Router,
{
    /// Execution plan
    exec_plan: ExecutionPlan,
    /// Coordinator runtime
    coordinator: &'a C,
    /// Bucket map of view { plan output_id (Expression::Row) -> `Buckets` }.
    /// It's supposed to denote relational nodes' output buckets destination.
    bucket_map: HashMap<NodeId, Buckets>,
}

impl<'a, C> ExecutingQuery<'a, C>
where
    C: Router,
{
    pub fn from_bound_statement(runtime: &'a C, statement: BoundStatement) -> Self {
        Self {
            exec_plan: ExecutionPlan::new(*statement.plan),
            coordinator: runtime,
            bucket_map: HashMap::new(),
        }
    }

    /// A shorthand to create a [`ExecutingQuery`] directly from SQL text.
    /// Equivalent to chaining [`BoundStatement::parse_and_bind`] and [`ExecutingQuery::from_bound_statement`].
    pub fn from_text_and_params(
        coordinator: &'a C,
        query_text: &str,
        params: Vec<Value>,
    ) -> Result<Self, SbroadError>
    where
        C::Cache: lru::Cache<SmolStr, Rc<Plan>>,
        C::ParseTree: crate::frontend::Ast,
    {
        let bound_statement = BoundStatement::parse_and_bind(
            coordinator,
            query_text,
            params,
            crate::ir::options::Options::default(),
        )?;

        Ok(Self::from_bound_statement(coordinator, bound_statement))
    }

    /// Get the execution plan of the query.
    #[must_use]
    pub fn get_exec_plan(&self) -> &ExecutionPlan {
        &self.exec_plan
    }

    /// Get the mutable reference to the execution plan of the query.
    #[must_use]
    pub fn get_mut_exec_plan(&mut self) -> &mut ExecutionPlan {
        &mut self.exec_plan
    }

    /// Get the coordinator runtime of the query.
    #[must_use]
    pub fn get_coordinator(&self) -> &C {
        self.coordinator
    }

    pub fn materialize_subtree<'p>(
        &mut self,
        slices: Slices,
        mut port: Option<&mut impl Port<'p>>,
    ) -> Result<(), SbroadError> {
        let tier = self.exec_plan.get_ir_plan().tier.as_ref();
        // all tables from one tier, so we can use corresponding vshard object
        let vshard = self.coordinator.get_vshard_object_by_tier(tier)?;

        for slice in slices.slices() {
            // TODO: make it work in parallel
            for motion_id in slice.positions() {
                if self.exec_plan.get_vtables().contains_key(motion_id) {
                    continue;
                }
                let motion = self.exec_plan.get_ir_plan().get_relation_node(*motion_id)?;
                if let Relational::Motion(Motion { policy, .. }) = motion {
                    match policy {
                        MotionPolicy::Segment(_)
                        // EXPLAIN(RAW) will be executed on the storage node
                        // We don't want to materialize VALUES on router and lose them in the execution plan.
                            if !self.exec_plan.get_ir_plan().is_raw_explain() =>
                        {
                            // If child is values, then we can materialize it
                            // on the router.
                            let motion_child_id = self
                                .get_exec_plan()
                                .get_ir_plan()
                                .get_motion_child(*motion_id)?;
                            let motion_child = self
                                .get_exec_plan()
                                .get_ir_plan()
                                .get_relation_node(motion_child_id)?;

                            if matches!(motion_child, Relational::Values { .. }) {
                                let virtual_table = self
                                    .coordinator
                                    .materialize_values(&mut self.exec_plan, motion_child_id)?;
                                self.exec_plan.set_motion_vtable(
                                    motion_id,
                                    virtual_table,
                                    &vshard,
                                )?;
                                self.get_mut_exec_plan().unlink_motion_subtree(*motion_id)?;
                                continue;
                            }
                        }
                        // Skip it and dispatch the query to the segments
                        // (materialization would be done on the segments). Note that we
                        // will operate with vtables for LocalSegment motions via calls like
                        // `self.exec_plan.contains_vtable_for_motion(node_id)`
                        // in order to define whether virtual table was materialized for values.
                        MotionPolicy::LocalSegment(_) => {
                            continue;
                        }
                        // Local policy should be skipped and dispatched to the segments:
                        // materialization would be done there.
                        MotionPolicy::Local => continue,
                        _ => {}
                    }
                }

                let top_id = self
                    .exec_plan
                    .get_ir_plan()
                    .get_motion_subtree_root(*motion_id)?;

                let buckets = self.bucket_discovery(top_id)?;
                let mut virtual_table = self.coordinator.materialize_motion(
                    &mut self.exec_plan,
                    motion_id,
                    &buckets,
                )?;

                if self.exec_plan.get_ir_plan().is_raw_explain() {
                    // Take the tuples from the virtual table and encode them into
                    // explain msgpack.
                    let tuples = std::mem::take(virtual_table.get_mut_tuples());
                    for tuple in tuples.into_iter() {
                        let mp = msgpack::encode(&tuple);
                        if let Some(p) = port.as_mut() {
                            p.add_mp(mp.as_slice())
                        }
                    }
                }

                self.exec_plan
                    .set_motion_vtable(motion_id, virtual_table, &vshard)?;
            }
        }

        Ok(())
    }

    /// Builds explain from current query
    ///
    /// # Errors
    /// - Failed to build explain
    pub fn produce_explain<'p>(&mut self, port: &mut impl Port<'p>) -> Result<(), SbroadError> {
        let mut mp: Vec<u8> = Vec::new();
        for line in self.as_explain()?.lines() {
            write_str(&mut mp, line).map_err(|e| {
                SbroadError::FailedTo(
                    Action::Deserialize,
                    Some(Entity::MsgPack),
                    format_smolstr!("{e}"),
                )
            })?;
            port.add_mp(&mp);
            mp.clear();
        }
        Ok(())
    }

    /// Dispatch a distributed query from coordinator to the segments.
    ///
    /// # Errors
    /// - Failed to get a motion subtree.
    /// - Failed to discover buckets.
    /// - Failed to materialize motion result and build a virtual table.
    /// - Failed to get plan top.
    pub fn dispatch<'p>(&mut self, port: &mut impl Port<'p>) -> Result<(), SbroadError> {
        let top_id = self.exec_plan.get_ir_plan().get_top()?;
        if self.exec_plan.get_ir_plan().is_block()? {
            let block = self.exec_plan.get_ir_plan().get_owned_block_node(top_id)?;
            let BlockOwned::Anonymous(block) = block else {
                unreachable!("plan.is_block() returned true, but top is {block:?}")
            };

            let mut block_buckets = None;
            for query_id in block.statements.iter().map(|stmt| *stmt.get()) {
                let buckets = self.bucket_discovery(query_id)?;
                match buckets {
                    Buckets::All => {
                        return Err(SbroadError::Other(format_smolstr!(
                            "transaction cannot be executed on all buckets"
                        )))
                    }
                    Buckets::Filtered(BucketSet::Exact(ref filtered)) => {
                        if filtered.len() != 1 {
                            return Err(SbroadError::Other(format_smolstr!(
                                "transaction can only be executed on a single bucket, got {buckets}"
                            )));
                        }

                        if let Some(block_buckets) = &block_buckets {
                            if block_buckets != &buckets {
                                return Err(SbroadError::Other(format_smolstr!(
                                    "transaction queries have different buckets: {block_buckets} and {buckets}"
                                )));
                            }
                        } else {
                            block_buckets = Some(buckets);
                        }
                    }
                    Buckets::Filtered(_) => {
                        return Err(SbroadError::Invalid(
                            Entity::Buckets,
                            Some("buckets are not discovered".into()),
                        ))
                    }
                    Buckets::Any => (),
                }
            }

            let buckets = block_buckets.unwrap_or(Buckets::Any);
            return self
                .coordinator
                .dispatch(&mut self.exec_plan, top_id, &buckets, port);
        }

        let slices = self.exec_plan.get_ir_plan().clone_slices();
        self.materialize_subtree(slices, Some(port))?;
        let ir_plan = self.exec_plan.get_ir_plan();
        if ir_plan.get_relation_node(top_id)?.is_motion() {
            let err =
                |s: &str| -> SbroadError { SbroadError::Invalid(Entity::Plan, Some(s.into())) };
            let aliases = ir_plan.get_relational_aliases(top_id)?;
            let Some(mut ref_table) = self.exec_plan.get_mut_vtables().remove(&top_id) else {
                return Err(err(&format!("no virtual table for motion id {top_id:?}")));
            };
            let Some(table) = Rc::get_mut(&mut ref_table) else {
                return Err(err("there are other references for the virtual table"));
            };

            // Skip metadata in case of `EXPLAIN (RAW)`
            if !self.exec_plan.get_ir_plan().is_raw_explain() {
                table
                    .dump_mp(aliases.iter().map(|s| s.as_str()), port)
                    .map_err(|e| {
                        SbroadError::Invalid(Entity::VirtualTable, Some(format_smolstr!("{e}")))
                    })?;
            }
            return Ok(());
        }

        let plan_id_target = self.exec_plan.get_plan_id_target()?;
        if let Some(node_id) = plan_id_target {
            self.exec_plan.set_plan_id(node_id)?;
        }
        let buckets = self.bucket_discovery(top_id)?;
        self.coordinator
            .dispatch(&mut self.exec_plan, top_id, &buckets, port)?;

        Ok(())
    }

    /// Query explain
    ///
    /// # Errors
    /// - Failed to build explain
    pub fn to_explain(&mut self) -> Result<SmolStr, SbroadError> {
        self.as_explain()
    }

    /// Checks that query is explain and have not to be executed
    pub fn is_logical_explain(&self) -> bool {
        self.exec_plan.get_ir_plan().is_logical_explain()
    }

    /// Checks that query is a statement block.
    ///
    /// # Errors
    /// - plan is invalid
    pub fn is_block(&self) -> Result<bool, SbroadError> {
        self.exec_plan.get_ir_plan().is_block()
    }

    /// Checks that query is DDL.
    ///
    /// # Errors
    /// - Plan is invalid.
    pub fn is_ddl(&self) -> Result<bool, SbroadError> {
        self.exec_plan.get_ir_plan().is_ddl()
    }

    /// Checks that query is ACL.
    ///
    /// # Errors
    /// - Plan is invalid
    pub fn is_acl(&self) -> Result<bool, SbroadError> {
        self.exec_plan.get_ir_plan().is_acl()
    }

    /// Checks that query is TCL.
    ///
    /// # Errors
    /// - Plan is invalid
    pub fn is_tcl(&self) -> Result<bool, SbroadError> {
        self.exec_plan.get_ir_plan().is_tcl()
    }

    #[cfg(test)]
    pub fn get_motion_id(&self, slice_id: usize, pos_idx: usize) -> NodeId {
        *self
            .exec_plan
            .get_ir_plan()
            .clone_slices()
            .slice(slice_id)
            .unwrap()
            .position(pos_idx)
            .unwrap()
    }

    /// Checks that query is for plugin.
    ///
    /// # Errors
    /// - Plan is invalid
    pub fn is_plugin(&self) -> Result<bool, SbroadError> {
        self.exec_plan.get_ir_plan().is_plugin()
    }

    pub fn is_backup(&self) -> Result<bool, SbroadError> {
        self.exec_plan.get_ir_plan().is_backup()
    }

    /// Checks that query is Deallocate.
    ///
    /// # Errors
    /// - Plan is invalid
    pub fn is_deallocate(&self) -> Result<bool, SbroadError> {
        self.exec_plan.get_ir_plan().is_deallocate()
    }

    /// Checks that query is an empty query.
    pub fn is_empty(&self) -> bool {
        self.exec_plan.get_ir_plan().is_empty()
    }
}

#[cfg(feature = "mock")]
#[cfg(test)]
pub mod tests;
