use std::collections::{HashMap, HashSet};
use std::rc::Rc;

use ahash::{AHashMap, AHashSet};
use serde::{Deserialize, Serialize};
use smol_str::{format_smolstr, ToSmolStr};

use crate::errors::{Action, Entity, SbroadError};
use crate::executor::engine::Vshard;
use crate::executor::vtable::{VirtualTable, VirtualTableMap};
use crate::ir::node::expression::ExprOwned;
use crate::ir::node::expression::{Expression, MutExpression};
use crate::ir::node::relational::{MutRelational, RelOwned, Relational};
use crate::ir::node::{
    Alias, ArenaType, ArithmeticExpr, BoolExpr, Bound, BoundType, Case, Cast, Concat, Delete,
    GroupBy, Having, IndexExpr, Insert, Join, Like, Motion, Node136, NodeId, NodeOwned, OrderBy,
    Over, Projection, Reference, ReferenceTarget, Row, ScalarFunction, ScanRelation, Selection,
    SubQueryReference, Trim, UnaryExpr, Update, ValuesRow, Window,
};
use crate::ir::operator::{OrderByElement, OrderByEntity};
use crate::ir::relation::SpaceEngine;
use crate::ir::transformation::redistribution::{MotionOpcode, MotionPolicy};
use crate::ir::tree::traversal::{LevelNode, PostOrder, PostOrderWithFilter, REL_CAPACITY};
use crate::ir::tree::Snapshot;
use crate::ir::Plan;

/// Query type (used to parse the returned results).
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone)]
pub enum QueryType {
    /// SELECT query.
    DQL,
    /// INSERT query.
    DML,
}

/// Connection type to the Tarantool instance.
#[derive(Debug)]
pub enum ConnectionType {
    Read,
    Write,
}

impl ConnectionType {
    #[must_use]
    pub fn is_readonly(&self) -> bool {
        match self {
            ConnectionType::Read => true,
            ConnectionType::Write => false,
        }
    }
}

/// Wrapper over `Plan` containing `vtables` map.
#[derive(Debug, Default, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct ExecutionPlan {
    request_id: String,
    pub(crate) plan: Plan,
    /// Virtual tables for `Motion` nodes.
    /// Map of { `Motion` node_id -> it's corresponding data }
    vtables: VirtualTableMap,
}

impl From<Plan> for ExecutionPlan {
    fn from(plan: Plan) -> Self {
        ExecutionPlan {
            request_id: uuid::Uuid::new_v4().to_string(),
            plan,
            vtables: VirtualTableMap::new(),
        }
    }
}

/// Translates the original plan's node id to the new sub-plan one.
struct SubtreeMap {
    inner: AHashMap<NodeId, NodeId>,
}

impl SubtreeMap {
    fn with_capacity(capacity: usize) -> Self {
        SubtreeMap {
            inner: AHashMap::with_capacity(capacity),
        }
    }

    fn get_id(&self, expr_id: NodeId) -> NodeId {
        *self
            .inner
            .get(&expr_id)
            .unwrap_or_else(|| panic!("Could not find node with id {expr_id:?} in subtree map"))
    }

    fn get_id_if_exist(&self, expr_id: NodeId) -> Option<&NodeId> {
        self.inner.get(&expr_id)
    }

    fn contains_key(&self, expr_id: NodeId) -> bool {
        self.inner.contains_key(&expr_id)
    }

    fn insert(&mut self, old_id: NodeId, new_id: NodeId) {
        self.inner.insert(old_id, new_id);
    }
}

/// Expected number of plan nodes which are part of Subquery/CTE subtrees.
const SQ_IDS_CAPACITY: usize = 100;

impl ExecutionPlan {
    #[must_use]
    pub fn get_request_id(&self) -> &str {
        &self.request_id
    }

    #[must_use]
    pub fn get_ir_plan(&self) -> &Plan {
        &self.plan
    }

    #[must_use]
    pub fn get_sql_motion_row_max(&self) -> u64 {
        self.plan.effective_options.sql_motion_row_max as u64
    }

    #[allow(dead_code)]
    pub fn get_mut_ir_plan(&mut self) -> &mut Plan {
        &mut self.plan
    }

    #[must_use]
    pub fn get_vtables(&self) -> &VirtualTableMap {
        &self.vtables
    }

    pub fn get_mut_vtables(&mut self) -> &mut VirtualTableMap {
        &mut self.vtables
    }

    pub fn set_vtables(&mut self, vtables: VirtualTableMap) {
        self.vtables = vtables;
    }

    pub fn contains_vtable_for_motion(&self, motion_id: NodeId) -> bool {
        self.get_vtables().contains_key(&motion_id)
    }

    /// Get motion virtual table
    pub fn get_motion_vtable(&self, motion_id: NodeId) -> Result<Rc<VirtualTable>, SbroadError> {
        if let Some(result) = self.get_vtables().get(&motion_id) {
            return Ok(Rc::clone(result));
        }
        let motion_node = self.get_ir_plan().get_relation_node(motion_id)?;
        panic!("Virtual table for motion {motion_node:?} with id {motion_id} not found.")
    }

    /// Add materialize motion result to map of virtual tables.
    ///
    /// # Errors
    /// - invalid motion node
    pub fn set_motion_vtable(
        &mut self,
        motion_id: &NodeId,
        vtable: VirtualTable,
        runtime: &impl Vshard,
    ) -> Result<(), SbroadError> {
        let mut vtable = vtable;
        let program_len = if let Relational::Motion(Motion { program, .. }) =
            self.get_ir_plan().get_relation_node(*motion_id)?
        {
            program.0.len()
        } else {
            return Err(SbroadError::Invalid(
                Entity::Node,
                Some("invalid motion node".into()),
            ));
        };
        for op_idx in 0..program_len {
            let plan = self.get_ir_plan();
            let opcode = plan.get_motion_opcode(*motion_id, op_idx)?;
            match opcode {
                MotionOpcode::RemoveDuplicates => {
                    vtable.remove_duplicates();
                }
                MotionOpcode::ReshardIfNeeded => {
                    // Resharding must be done before applying projection
                    // to the virtual table. Otherwise projection can
                    // remove sharding columns.
                    match plan.get_motion_policy(*motion_id)? {
                        MotionPolicy::Segment(shard_key)
                        | MotionPolicy::LocalSegment(shard_key) => {
                            vtable.reshard(shard_key, runtime)?;
                        }
                        MotionPolicy::Full | MotionPolicy::Local | MotionPolicy::None => {}
                    }
                }
                MotionOpcode::PrimaryKey(positions) => {
                    vtable.set_primary_key(positions)?;
                }
                MotionOpcode::RearrangeForShardedUpdate {
                    update_id,
                    old_shard_columns_len,
                    new_shard_columns_positions,
                } => {
                    let update_id = *update_id;

                    if let Some(v) = vtable.rearrange_for_update(
                        runtime,
                        *old_shard_columns_len,
                        new_shard_columns_positions,
                    )? {
                        let plan = self.get_mut_ir_plan();
                        plan.set_update_delete_tuple_len(update_id, v)?;
                    }
                }
                MotionOpcode::AddMissingRowsForLeftJoin { motion_id, .. } => {
                    let motion_id = *motion_id;
                    let Some(from_vtable) = self.vtables.get(&motion_id) else {
                        return Err(SbroadError::UnexpectedNumberOfValues(format_smolstr!(
                            "expected virtual table for motion {motion_id:?}"
                        )));
                    };
                    vtable.add_missing_rows(from_vtable)?;
                }
                MotionOpcode::SerializeAsEmptyTable(_) => {}
            }
        }

        self.get_mut_vtables().insert(*motion_id, Rc::new(vtable));

        Ok(())
    }

    #[must_use]
    pub fn has_segmented_tables(&self) -> bool {
        self.vtables
            .values()
            .any(|t| !t.get_bucket_index().is_empty())
    }

    /// Return true if plan needs to be customized for each storage.
    /// I.e we can't send the same plan to all storages.
    ///
    /// The check is done by iterating over the plan nodes arena,
    /// and checking whether motion node contains serialize as empty
    /// opcode. Be sure there are no dead nodes in the plan arena:
    /// nodes that are not referenced by actual plan tree.
    #[must_use]
    pub fn has_customization_opcodes(&self) -> bool {
        for node in self.get_ir_plan().nodes.iter136() {
            if let Node136::Motion(Motion { program, .. }) = node {
                if program
                    .0
                    .iter()
                    .any(|op| matches!(op, MotionOpcode::SerializeAsEmptyTable(_)))
                {
                    return true;
                }
            }
        }
        false
    }

    /// Extract policy from motion node
    ///
    /// # Errors
    /// - node is not `Relation` type
    /// - node is not `Motion` type
    pub fn get_motion_policy(&self, node_id: NodeId) -> Result<MotionPolicy, SbroadError> {
        if let Relational::Motion(Motion { policy, .. }) = &self.plan.get_relation_node(node_id)? {
            return Ok(policy.clone());
        }

        Err(SbroadError::Invalid(
            Entity::Relational,
            Some("invalid motion".into()),
        ))
    }

    /// Get root from motion sub tree
    ///
    /// # Errors
    /// - node is not valid
    pub fn get_motion_subtree_root(&self, node_id: NodeId) -> Result<NodeId, SbroadError> {
        let top_id = &self.get_motion_child(node_id)?;
        let rel = self.get_ir_plan().get_relation_node(*top_id)?;
        match rel {
            Relational::ScanSubQuery { .. } | Relational::ScanCte { .. } => {
                self.get_ir_plan().get_first_rel_child(*top_id)
            }
            Relational::Except { .. }
            | Relational::GroupBy { .. }
            | Relational::OrderBy { .. }
            | Relational::Intersect { .. }
            | Relational::Join { .. }
            | Relational::Projection { .. }
            | Relational::ScanRelation { .. }
            | Relational::Selection { .. }
            | Relational::SelectWithoutScan { .. }
            | Relational::Union { .. }
            | Relational::UnionAll { .. }
            | Relational::Update { .. }
            | Relational::Values { .. }
            | Relational::Having { .. }
            | Relational::ValuesRow { .. }
            | Relational::Limit { .. } => Ok(*top_id),
            Relational::Motion { .. } | Relational::Insert { .. } | Relational::Delete { .. } => {
                Err(SbroadError::Invalid(
                    Entity::Relational,
                    Some(format_smolstr!("invalid motion child node: {rel:?}.")),
                ))
            }
        }
    }

    /// Extract the child from the Motion node.
    ///
    /// # Panics:
    /// - `motion_id` is not Motion IR node
    ///
    /// # Errors
    /// - `motion_id` is not `Relational` type
    /// - Node `motion_id` does not contain child
    pub(crate) fn get_motion_child(&self, motion_id: NodeId) -> Result<NodeId, SbroadError> {
        let Relational::Motion(_) = self.get_ir_plan().get_relation_node(motion_id)? else {
            unreachable!("expected Motion IR node");
        };
        self.plan.get_first_rel_child(motion_id)
    }

    /// Unlink the subtree of the motion node.
    /// This logic is applied as an optimization: when we dispatch plan subtrees to the storages,
    /// we don't need to dispatch a subtrees of `Motion` nodes as soon as they are already replaced
    /// with vtables. So we replace its `children` field with an empty vec.
    ///
    /// # Errors
    /// - not a motion node
    pub fn unlink_motion_subtree(&mut self, motion_id: NodeId) -> Result<(), SbroadError> {
        let motion = self.get_mut_ir_plan().get_mut_relation_node(motion_id)?;
        if let MutRelational::Motion(Motion { child, .. }) = motion {
            *child = None;
        } else {
            return Err(SbroadError::Invalid(
                Entity::Relational,
                Some(format_smolstr!("node ({motion_id:?}) is not motion")),
            ));
        }
        Ok(())
    }

    /// Build a new execution plan from the subtree of the existing execution plan.
    /// The operation is destructive and the subtree is removed from the original plan.
    ///
    /// # Errors
    /// - the original execution plan is invalid
    ///
    /// # Panics
    /// - Plan is in invalid state
    #[allow(clippy::too_many_lines)]
    pub fn take_subtree(&mut self, subtree_id: NodeId) -> Result<Self, SbroadError> {
        // Get the subtree nodes indexes.
        let plan = self.get_ir_plan();
        let top_id = plan.get_top()?;

        // XXX For historical reasons we destructively extract the subtree
        // before sending it to a storage for execution. Ideally we would not
        // do this and instead just send it via a read-only reference, and this
        // will be done in future when we switch to the new sql-protocol.
        //
        // But for now we introduce this hack. It is needed because global DML
        // operations go through CAS API, which may require retrying the request
        // in expected situations (for example due to automatic compaction of
        // raft log). In order to support correct DQL + DML queries we need to
        // keep the original plan intact during the retry.
        let dont_mutate = plan.is_dml_on_global_table()?;

        // We don't cut CTE and subquery subtrees during plan traversal
        // as they can be reused in other slices of the plan.
        // So, we collect such subtree nodes into the set to avoid their removal.
        //
        // E.g. look at `one-sharded-one-global-filter-subquery` test to see an
        // example of a subquery which is referenced by several nodes (one
        // below Motion and one above it). Such a situation is caused by our
        // EXCEPT implementation logic.
        let mut nodes_to_save = AHashSet::default();
        if !dont_mutate {
            let filter = |node_id: NodeId| -> bool {
                if let Ok(Relational::ScanCte(_) | Relational::ScanSubQuery(_)) =
                    plan.get_relation_node(node_id)
                {
                    return true;
                }
                false
            };
            let mut rel_tree = PostOrderWithFilter::with_capacity(
                |node| plan.nodes.rel_iter(node),
                REL_CAPACITY,
                Box::new(filter),
            );
            rel_tree.populate_nodes(top_id);

            // Preallocate memory for all subqueries and CTE subtrees.
            nodes_to_save.reserve(SQ_IDS_CAPACITY * 2);

            for LevelNode(_, node_id) in rel_tree.iter() {
                let subtree = PostOrder::with_capacity(
                    |node| plan.exec_plan_subtree_iter(node, Snapshot::Oldest),
                    REL_CAPACITY,
                );
                for LevelNode(_, id) in subtree.into_iter(*node_id) {
                    nodes_to_save.insert(id);
                }
            }
        }

        let mut subtree = PostOrder::with_capacity(
            |node| plan.exec_plan_subtree_iter(node, Snapshot::Oldest),
            plan.nodes.len(),
        );
        subtree.populate_nodes(subtree_id);
        let nodes = subtree.take_nodes();

        let mut subtree_map = SubtreeMap::with_capacity(nodes.len());
        let vtables_capacity = self.get_vtables().len();
        // Map of { plan node_id -> virtual table }.
        let mut new_vtables: HashMap<NodeId, Rc<VirtualTable>> =
            HashMap::with_capacity(vtables_capacity);

        let mut new_plan = Plan::new();
        new_plan.mark_as_explain(plan.get_explain_type());
        // In case we have a Motion among rel node children (maybe not direct), we
        // need to rename rel output aliases, because Motion
        // may have changed them according to its vtable column names.
        // This map tracks outputs of rel nodes that have changed their aliases.
        let mut rel_renamed_output_lists: HashMap<NodeId, Vec<NodeId>> = HashMap::new();

        // In motion, we can eliminate children. To keep references valid, we use this set.
        // This set is also needed for debugging purposes: if we traverse in the wrong order,
        // the array will contain references that point to incorrect locations.
        let mut invalid_refs = HashSet::with_capacity(nodes.len());

        for LevelNode(_, node_id) in nodes {
            // We have already processed this node (sub-queries in BETWEEN
            // and CTEs can be referred twice).
            if subtree_map.contains_key(node_id) {
                continue;
            }

            let mut_plan = self.get_mut_ir_plan();

            // Replace the node with some invalid value.
            let node = mut_plan.get_node(node_id)?;
            let mut node: NodeOwned = if dont_mutate || nodes_to_save.contains(&node_id) {
                node.into_owned()
            } else {
                mut_plan.replace_with_stub(node_id)
            };

            let mut is_invalid_ref = false;
            let ir_plan = self.get_ir_plan();
            match node {
                NodeOwned::Relational(ref mut rel) => {
                    match rel {
                        RelOwned::Selection(Selection {
                            filter: ref mut expr_id,
                            ..
                        })
                        | RelOwned::Having(Having {
                            filter: ref mut expr_id,
                            ..
                        })
                        | RelOwned::Join(Join {
                            condition: ref mut expr_id,
                            ..
                        }) => {
                            // We transform selection's, having's filter and join's condition to DNF for a better bucket calculation.
                            // But as a result we can produce an extremely verbose SQL query from such a plan (tarantool's
                            // parser can fail to parse such SQL).

                            // XXX: UNDO operation can cause problems if we introduce more complicated transformations
                            // for filter/condition (but then the UNDO logic should be changed as well).
                            let undo_expr_id = ir_plan.undo.get_oldest(expr_id);
                            *expr_id = subtree_map.get_id(*undo_expr_id);
                        }
                        RelOwned::ScanRelation(ScanRelation {
                            relation,
                            indexed_by,
                            ..
                        }) => {
                            let table = ir_plan
                                .relations
                                .get(relation)
                                .unwrap_or_else(|| {
                                    panic!(
                                        "could not find relation {relation} in the original plan"
                                    )
                                })
                                .clone();

                            if cfg!(not(feature = "mock")) {
                                if let Some(indexed_by) = indexed_by {
                                    let index = ir_plan.indexes.get(indexed_by).ok_or(
                                        SbroadError::NotFound(
                                            Entity::Index,
                                            format_smolstr!("with name {indexed_by}"),
                                        ),
                                    )?;
                                    new_plan.index_version_map.insert(
                                        [index.table_id, index.id],
                                        *ir_plan
                                            .index_version_map
                                            .get(&[index.table_id, index.id])
                                            .expect("index id and version must exist"),
                                    );
                                }
                                new_plan.table_version_map.insert(
                                    table.id,
                                    ir_plan.table_version_map.get(&table.id).cloned().unwrap(),
                                );
                            }

                            new_plan.add_rel(table);
                        }
                        RelOwned::Insert(Insert { relation, .. })
                        | RelOwned::Delete(Delete { relation, .. })
                        | RelOwned::Update(Update { relation, .. }) => {
                            let table = ir_plan
                                .relations
                                .get(relation)
                                .expect("relation must exist")
                                .clone();

                            if cfg!(not(feature = "mock")) {
                                new_plan.table_version_map.insert(
                                    table.id,
                                    ir_plan.table_version_map.get(&table.id).cloned().unwrap(),
                                );
                            }

                            new_plan.add_rel(table);
                        }
                        RelOwned::Motion(Motion {
                            child,
                            policy,
                            output,
                            ..
                        }) => {
                            if let Some(vtable) = self.get_vtables().get(&node_id) {
                                let next_id = new_plan.nodes.next_id(ArenaType::Arena136);
                                new_vtables.insert(next_id, Rc::clone(vtable));

                                // We need to fix motion aliases based on materialized
                                // vtable. Motion's aliases will copy "COL_i" naming of
                                // vtable.
                                let output_list: Vec<NodeId> =
                                    new_plan.get_row_list(subtree_map.get_id(*output))?.to_vec();

                                if node_id != top_id {
                                    // We should rename Motion aliases only in case it's not a
                                    // top node. Otherwise, it has no effect as soon as `to_sql`
                                    // will use column names of vtable and ignore motion aliases.
                                    vtable
                                        .get_columns().iter()
                                        .zip(output_list.iter())
                                        .map(|(vtable_column, alias_id)| {
                                            let alias = new_plan.get_mut_expression_node(
                                                *alias_id
                                            )?;
                                            let MutExpression::Alias(Alias { ref mut name, .. }) = alias else {
                                                return Err(SbroadError::Invalid(
                                                    Entity::Expression,
                                                    Some(format_smolstr!("Expected Alias under Motion output, got {alias:?}"))
                                                ))
                                            };
                                            *name = vtable_column.name.clone();
                                            Ok(())
                                        })
                                        .collect::<Result<Vec<()>, SbroadError>>()?;
                                }
                            }
                            // We should not remove the child of a local motion node.
                            // The subtree is needed to compile the SQL on the storage.
                            if !policy.is_local() {
                                *child = None;

                                // We need to make invalid reference leafs
                                if !invalid_refs.is_empty() {
                                    let output_list: Vec<NodeId> = new_plan
                                        .get_row_list(subtree_map.get_id(*output))?
                                        .to_vec();

                                    for node in output_list {
                                        let node = new_plan.get_child_under_alias(node)?;

                                        if !invalid_refs.contains(&node) {
                                            continue;
                                        }

                                        invalid_refs.remove(&node);

                                        if let MutExpression::Reference(Reference {
                                            ref mut target,
                                            ..
                                        }) = new_plan.get_mut_expression_node(node)?
                                        {
                                            *target = ReferenceTarget::Leaf;
                                        }
                                    }
                                }
                            }
                        }
                        RelOwned::GroupBy(GroupBy { gr_exprs, .. }) => {
                            let mut new_cols: Vec<NodeId> = Vec::with_capacity(gr_exprs.len());
                            for col_id in gr_exprs.iter() {
                                let new_col_id = subtree_map.get_id(*col_id);
                                new_cols.push(new_col_id);
                            }
                            *gr_exprs = new_cols;
                        }
                        RelOwned::OrderBy(OrderBy {
                            order_by_elements, ..
                        }) => {
                            let mut new_elements: Vec<OrderByElement> =
                                Vec::with_capacity(order_by_elements.len());
                            for element in order_by_elements.iter() {
                                let new_entity = match element.entity {
                                    OrderByEntity::Expression { expr_id } => {
                                        let new_element_id = subtree_map.get_id(expr_id);
                                        OrderByEntity::Expression {
                                            expr_id: new_element_id,
                                        }
                                    }
                                    OrderByEntity::Index { value } => {
                                        OrderByEntity::Index { value }
                                    }
                                };
                                new_elements.push(OrderByElement {
                                    entity: new_entity,
                                    order_type: element.order_type.clone(),
                                });
                            }
                            *order_by_elements = new_elements;
                        }
                        RelOwned::ValuesRow(ValuesRow { data, .. }) => {
                            *data = subtree_map.get_id(*data);
                        }
                        RelOwned::Projection(Projection {
                            windows,
                            group_by,
                            having,
                            ..
                        }) => {
                            for window in windows {
                                *window = subtree_map.get_id(*window);
                            }
                            if let Some(group_by_id) = group_by {
                                *group_by_id = subtree_map.get_id(*group_by_id);
                            }
                            if let Some(having_id) = having {
                                *having_id = subtree_map.get_id(*having_id);
                            }
                        }
                        RelOwned::Except { .. }
                        | RelOwned::Intersect { .. }
                        | RelOwned::SelectWithoutScan { .. }
                        | RelOwned::ScanSubQuery { .. }
                        | RelOwned::ScanCte { .. }
                        | RelOwned::Union { .. }
                        | RelOwned::UnionAll { .. }
                        | RelOwned::Values { .. }
                        | RelOwned::Limit { .. } => {}
                    }

                    for child_id in rel.mut_children() {
                        *child_id = subtree_map.get_id(*child_id);

                        let child_rel_node = new_plan.get_relation_node(*child_id)?;
                        if let Relational::Motion(Motion { output, .. }) = child_rel_node {
                            let motion_output_list: Vec<NodeId> =
                                new_plan.get_row_list(*output)?.to_vec();
                            rel_renamed_output_lists.insert(*child_id, motion_output_list);
                        }
                    }

                    if rel.has_output() {
                        let output = rel.mut_output();
                        *rel.mut_output() = subtree_map.get_id(*output);

                        // If we deal with Projection we have to fix
                        // only References that have an Asterisk source.
                        // References without asterisks would be covered with aliases like
                        // "COL_1 as <alias>".
                        // If we deal with `top_id` node we shouldn't rename its aliases with motion child.
                        // F.e., `SELECT * FROM t LIMIT 1`
                        let is_projection = matches!(rel, RelOwned::Projection(_));
                        if !rel_renamed_output_lists.is_empty() {
                            let rel_output_list: Vec<NodeId> =
                                new_plan.get_row_list(*rel.mut_output())?.to_vec();

                            for output_id in &rel_output_list {
                                let ref_under_alias = new_plan.get_child_under_alias(*output_id)?;
                                let ref_expr = new_plan.get_expression_node(ref_under_alias)?;
                                let Expression::Reference(Reference {
                                    position,
                                    target,
                                    asterisk_source,
                                    ..
                                }) = ref_expr
                                else {
                                    continue;
                                };

                                if node_id == top_id || is_projection && asterisk_source.is_none() {
                                    continue;
                                }

                                if target.is_leaf() {
                                    break;
                                }

                                let ref_rel_node = target
                                    .first()
                                    .expect("Reference must have at least one target.");

                                if let Some(child_output_list) =
                                    rel_renamed_output_lists.get(ref_rel_node)
                                {
                                    let child_output_alias_id =
                                        child_output_list.get(*position).unwrap_or_else(|| {
                                            panic!(
                                                "Unable to get motion output at requested position."
                                            );
                                        });
                                    let child_alias =
                                        new_plan.get_expression_node(*child_output_alias_id)?;
                                    let child_alias_name =
                                        child_alias.get_alias_name()?.to_smolstr();

                                    let rel_alias = new_plan.get_mut_expression_node(*output_id)?;
                                    if let MutExpression::Alias(Alias { ref mut name, .. }) =
                                        rel_alias
                                    {
                                        *name = child_alias_name;
                                    } else {
                                        panic!("Expected alias under Row output list");
                                    }
                                }
                            }

                            let arena_type = rel.arena_type();
                            let next_id = new_plan.nodes.next_id(arena_type);

                            rel_renamed_output_lists.insert(next_id, rel_output_list);
                        }
                    }
                }
                NodeOwned::Expression(ref mut expr) => match expr {
                    ExprOwned::Window(Window {
                        ref mut partition,
                        ref mut ordering,
                        ref mut frame,
                        ..
                    }) => {
                        if let Some(partition) = partition {
                            for expr_id in partition {
                                *expr_id = subtree_map.get_id(*expr_id);
                            }
                        }
                        if let Some(ordering) = ordering {
                            for element in ordering {
                                match &mut element.entity {
                                    OrderByEntity::Expression { expr_id } => {
                                        *expr_id = subtree_map.get_id(*expr_id);
                                    }
                                    OrderByEntity::Index { .. } => {}
                                }
                            }
                        }
                        if let Some(frame) = frame {
                            match &mut frame.bound {
                                Bound::Single(BoundType::PrecedingOffset(expr_id))
                                | Bound::Single(BoundType::FollowingOffset(expr_id)) => {
                                    *expr_id = subtree_map.get_id(*expr_id);
                                }
                                Bound::Between(b_start, b_end) => {
                                    for b_type in [b_start, b_end].iter_mut() {
                                        match b_type {
                                            BoundType::PrecedingOffset(expr_id)
                                            | BoundType::FollowingOffset(expr_id) => {
                                                *expr_id = subtree_map.get_id(*expr_id);
                                            }
                                            _ => {}
                                        }
                                    }
                                }
                                _ => {}
                            }
                        }
                    }
                    ExprOwned::Over(Over {
                        ref mut stable_func,
                        ref mut filter,
                        ref mut window,
                        ..
                    }) => {
                        if let Some(filter) = filter {
                            *filter = subtree_map.get_id(*filter)
                        }
                        *stable_func = subtree_map.get_id(*stable_func);
                        *window = subtree_map.get_id(*window);
                    }
                    ExprOwned::Alias(Alias { ref mut child, .. })
                    | ExprOwned::Cast(Cast { ref mut child, .. })
                    | ExprOwned::Unary(UnaryExpr { ref mut child, .. }) => {
                        *child = subtree_map.get_id(*child);
                    }
                    ExprOwned::Index(IndexExpr {
                        ref mut child,
                        ref mut which,
                    }) => {
                        *child = subtree_map.get_id(*child);
                        *which = subtree_map.get_id(*which);
                    }
                    ExprOwned::Bool(BoolExpr {
                        ref mut left,
                        ref mut right,
                        ..
                    })
                    | ExprOwned::Arithmetic(ArithmeticExpr {
                        ref mut left,
                        ref mut right,
                        ..
                    })
                    | ExprOwned::Concat(Concat {
                        ref mut left,
                        ref mut right,
                        ..
                    }) => {
                        *left = subtree_map.get_id(*left);
                        *right = subtree_map.get_id(*right);
                    }
                    ExprOwned::Like(Like {
                        escape: ref mut escape_id,
                        ref mut right,
                        ref mut left,
                    }) => {
                        *left = subtree_map.get_id(*left);
                        *right = subtree_map.get_id(*right);
                        *escape_id = subtree_map.get_id(*escape_id);
                    }
                    ExprOwned::Trim(Trim {
                        ref mut pattern,
                        ref mut target,
                        ..
                    }) => {
                        if let Some(pattern) = pattern {
                            *pattern = subtree_map.get_id(*pattern);
                        }
                        *target = subtree_map.get_id(*target);
                    }
                    ExprOwned::Reference(Reference { ref mut target, .. }) => match target {
                        ReferenceTarget::Leaf => {}
                        ReferenceTarget::Single(node_id) => {
                            match subtree_map.get_id_if_exist(*node_id) {
                                Some(node) => {
                                    *target = ReferenceTarget::Single(*node);
                                }
                                None => is_invalid_ref = true,
                            }
                        }
                        ReferenceTarget::Union(left, right) => {
                            match (
                                subtree_map.get_id_if_exist(*left),
                                subtree_map.get_id_if_exist(*right),
                            ) {
                                (Some(left), Some(right)) => {
                                    *target = ReferenceTarget::Union(*left, *right);
                                }
                                _ => is_invalid_ref = true,
                            }
                        }
                        ReferenceTarget::Values(nodes) => {
                            if nodes
                                .iter()
                                .any(|node| subtree_map.get_id_if_exist(*node).is_none())
                            {
                                is_invalid_ref = true;
                            } else {
                                let new_targets = nodes
                                    .iter()
                                    .map(|node_id| subtree_map.get_id(*node_id))
                                    .collect();
                                *target = ReferenceTarget::Values(new_targets);
                            }
                        }
                    },
                    ExprOwned::SubQueryReference(SubQueryReference { rel_id, .. }) => {
                        match subtree_map.get_id_if_exist(*rel_id) {
                            Some(node) => {
                                *rel_id = *node;
                            }
                            None => is_invalid_ref = true,
                        }
                    }
                    ExprOwned::Row(Row {
                        list: ref mut children,
                        ..
                    })
                    | ExprOwned::ScalarFunction(ScalarFunction {
                        ref mut children, ..
                    }) => {
                        for child in children {
                            *child = subtree_map.get_id(*child);
                        }
                    }
                    ExprOwned::Constant { .. }
                    | ExprOwned::CountAsterisk { .. }
                    | ExprOwned::Timestamp(_)
                    | ExprOwned::Parameter { .. } => {}
                    ExprOwned::Case(Case {
                        search_expr,
                        when_blocks,
                        else_expr,
                    }) => {
                        if let Some(search_expr) = search_expr {
                            *search_expr = subtree_map.get_id(*search_expr);
                        }
                        for (cond_expr, res_expr) in when_blocks {
                            *cond_expr = subtree_map.get_id(*cond_expr);
                            *res_expr = subtree_map.get_id(*res_expr);
                        }
                        if let Some(else_expr) = else_expr {
                            *else_expr = subtree_map.get_id(*else_expr);
                        }
                    }
                },
                NodeOwned::Invalid { .. }
                | NodeOwned::Ddl { .. }
                | NodeOwned::Acl { .. }
                | NodeOwned::Tcl(_)
                | NodeOwned::Plugin { .. }
                | NodeOwned::Deallocate { .. }
                | NodeOwned::Block { .. } => {
                    panic!("Unexpected node in `take_subtree`: {node:?}")
                }
            }

            let id = new_plan.nodes.push(node.into());

            if is_invalid_ref {
                invalid_refs.insert(id);
            }

            subtree_map.insert(node_id, id);
            if subtree_id == node_id {
                new_plan.set_top(id)?;
            }
        }

        if !invalid_refs.is_empty() {
            return Err(SbroadError::Other(format_smolstr!(
                "invalid references ({}) should be resolved",
                invalid_refs.len()
            )));
        }

        new_plan.stash_constants(Snapshot::Oldest)?;
        new_plan.effective_options = self.get_ir_plan().effective_options.clone();
        new_plan.tier.clone_from(&self.get_ir_plan().tier);

        let vtables = if new_vtables.is_empty() {
            VirtualTableMap::new()
        } else {
            new_vtables
        };
        let new_exec_plan = ExecutionPlan {
            request_id: self.request_id.clone(),
            plan: new_plan,
            vtables,
        };
        Ok(new_exec_plan)
    }

    /// # Errors
    /// - execution plan is invalid
    pub fn query_type(&self) -> Result<QueryType, SbroadError> {
        let top_id = self.get_ir_plan().get_top()?;
        let top = self.get_ir_plan().get_relation_node(top_id)?;
        if top.is_dml() {
            Ok(QueryType::DML)
        } else {
            Ok(QueryType::DQL)
        }
    }

    pub fn vtables_empty(&self) -> bool {
        self.get_vtables().is_empty()
    }

    /// Calculates an engine for the virtual tables in the plan.
    /// The main problem is that we can't use different engines
    /// for an SQL statement within transaction.
    ///
    /// # Errors
    /// - execution plan is invalid
    /// - multiple engines are used in the plan
    pub fn vtable_engine(&self) -> Result<SpaceEngine, SbroadError> {
        let mut engine: Option<SpaceEngine> = None;
        for table in self.get_ir_plan().relations.tables.values() {
            let table_engine = table.engine();
            if engine.is_none() {
                engine = Some(table_engine);
            } else if engine != Some(table_engine) {
                return Err(SbroadError::FailedTo(
                    Action::Build,
                    Some(Entity::Plan),
                    format_smolstr!(
                        "cannot build execution plan for DML query with multiple engines: {:?}",
                        self.get_ir_plan().relations.tables
                    ),
                ));
            }
        }
        Ok(engine.unwrap_or(SpaceEngine::Memtx))
    }

    /// # Errors
    /// - execution plan is invalid
    pub fn connection_type(&self) -> Result<ConnectionType, SbroadError> {
        match self.query_type()? {
            QueryType::DML => Ok(ConnectionType::Write),
            QueryType::DQL => {
                if self.vtables_empty() {
                    Ok(ConnectionType::Read)
                } else {
                    Ok(ConnectionType::Write)
                }
            }
        }
    }
}
