//! Helper module with functions and structures for the IR.

use smol_str::{SmolStr, ToSmolStr};

use crate::backend::sql::tree::{SyntaxData, SyntaxPlan};
use crate::errors::{Action, Entity, SbroadError};
use crate::ir::node::{
    Alias, BoolExpr, Case, Constant, Delete, GroupBy, Having, Join, Motion, NodeId, OrderBy,
    Reference, Row, ScanCte, ScanRelation, ScanSubQuery, Selection, SubQueryReference,
    TimeParameters, Trim, UnaryExpr, Update, ValuesRow,
};
use crate::ir::operator::OrderByEntity;
use crate::ir::tree::traversal::{PostOrder, EXPR_CAPACITY};
use crate::ir::{Node, Plan};
use std::collections::hash_map::DefaultHasher;
use std::fmt::Write;
use std::hash::BuildHasher;

use super::node::expression::Expression;
use super::node::relational::Relational;
use super::node::{
    ArithmeticExpr, Like, Limit, Over, Parameter, ScalarFunction, Timestamp, Window,
};

/// Helper macros to build a hash map or set
/// from the list of arguments.
#[macro_export]
macro_rules! collection {
    // map-like
    ($($k:expr => $v:expr),* $(,)?) => {{
        use std::iter::{Iterator, IntoIterator};
        Iterator::collect(IntoIterator::into_iter([$(($k, $v),)*]))
    }};
    // set-like
    ($($v:expr),* $(,)?) => {{
        use std::iter::{Iterator, IntoIterator};
        Iterator::collect(IntoIterator::into_iter([$($v,)*]))
    }};
}

/// By default, `HashSet` uses `RandomState` to initialize the hasher.
/// Because of that, we can't get repeatable order in the set between
/// test runs that can cause flaky test results in some cases.
#[derive(Clone)]
pub struct RepeatableState;

impl BuildHasher for RepeatableState {
    type Hasher = DefaultHasher;
    fn build_hasher(&self) -> DefaultHasher {
        DefaultHasher::new()
    }
}

impl Default for RepeatableState {
    fn default() -> Self {
        Self
    }
}

/// Helper formatting function for tabulation.
fn formatted_tabulate(buf: &mut String, n: i32) -> Result<(), std::fmt::Error> {
    for _ in 0..n {
        write!(buf, "\t")?;
    }
    Ok(())
}

/// Helper formatting function for writing with tabulation.
fn write_with_tabulation(buf: &mut String, n: i32, text: &str) -> Result<(), std::fmt::Error> {
    formatted_tabulate(buf, n)?;
    write!(buf, "{text}")
}

/// Helper formatting function for writing with tabulation and new line.
fn writeln_with_tabulation(buf: &mut String, n: i32, text: &str) -> Result<(), std::fmt::Error> {
    formatted_tabulate(buf, n)?;
    writeln!(buf, "{text}")
}

/// Formatting helper debug functions
impl Plan {
    /// Helper function for printing Expression node.
    #[allow(clippy::too_many_lines)]
    fn write_expr(
        &self,
        buf: &mut String,
        tabulation_number: i32,
        node_id: NodeId,
    ) -> Result<(), std::fmt::Error> {
        let expr_try = self.get_expression_node(node_id);
        if let Ok(expr) = expr_try {
            write!(buf, "expression: ")?;
            match expr {
                Expression::Window(Window {
                    partition,
                    ordering,
                    frame,
                }) => {
                    writeln!(buf, "Window")?;
                    if let Some(partition) = partition {
                        writeln_with_tabulation(buf, tabulation_number + 1, "Partition")?;
                        for part in partition {
                            self.formatted_arena_node(buf, tabulation_number + 2, *part)?;
                        }
                    }
                    if let Some(ordering) = ordering {
                        writeln_with_tabulation(buf, tabulation_number + 1, "Ordering")?;
                        for order in ordering {
                            writeln_with_tabulation(
                                buf,
                                tabulation_number + 2,
                                &format!("{order:?}"),
                            )?;
                        }
                    }
                    if let Some(frame) = frame {
                        writeln_with_tabulation(
                            buf,
                            tabulation_number + 2,
                            &format!("Frame: {frame:?}"),
                        )?;
                    }
                }
                Expression::Over(Over {
                    stable_func,
                    filter,
                    window,
                }) => {
                    writeln!(buf, "Over")?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "StableFunc")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *stable_func)?;
                    if let Some(filter) = filter {
                        writeln_with_tabulation(buf, tabulation_number + 1, "Filter")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *filter)?;
                    }
                    writeln_with_tabulation(buf, tabulation_number + 1, "Window")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *window)?;
                }
                Expression::Alias(Alias { name, child }) => {
                    writeln!(buf, "Alias [name = {name}]")?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Child:")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *child)?;
                }
                Expression::Case(Case {
                    search_expr,
                    when_blocks,
                    else_expr,
                }) => {
                    writeln!(buf, "Case")?;
                    if let Some(search_expr) = search_expr {
                        writeln_with_tabulation(buf, tabulation_number + 1, "Search_expr")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *search_expr)?;
                    }
                    for (cond_expr, res_expr) in when_blocks {
                        writeln_with_tabulation(buf, tabulation_number + 1, "WHEN")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *cond_expr)?;
                        writeln_with_tabulation(buf, tabulation_number + 1, "THEN")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *res_expr)?;
                    }
                    if let Some(else_expr) = else_expr {
                        writeln_with_tabulation(buf, tabulation_number + 1, "Else expr")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *else_expr)?;
                    }
                }
                Expression::Bool(BoolExpr { op, left, right }) => {
                    writeln!(buf, "Bool [op: {op}]")?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Left child")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *left)?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Right child")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *right)?;
                }
                Expression::Constant(Constant { value }) => {
                    writeln!(buf, "Constant [value = {value}]")?;
                }
                Expression::CountAsterisk(_) => writeln!(buf, "CountAsterisk")?,
                Expression::Reference(Reference {
                    target,
                    position,
                    col_type,
                    ..
                }) => {
                    writeln!(buf, "Reference")?;
                    if !target.is_leaf() {
                        let alias = self.get_alias_from_reference_node(&expr);
                        if let Result::Ok(alias) = alias {
                            writeln_with_tabulation(
                                buf,
                                tabulation_number + 1,
                                format!("Alias: {alias}").as_str(),
                            )?;
                        }
                    }

                    // See explain logic for Reference node
                    let rel_id = self.get_relational_from_reference_node(node_id);
                    if let Ok(rel_id) = rel_id {
                        let rel_node = self.get_relation_node(rel_id);
                        if rel_node.is_ok() {
                            if let Ok(Some(name)) = self.scan_name(rel_id, *position) {
                                writeln_with_tabulation(
                                    buf,
                                    tabulation_number + 1,
                                    format!("Referenced table name (or alias): {name}").as_str(),
                                )?;
                            }
                        }
                    }

                    for target_id in target.iter() {
                        writeln_with_tabulation(
                            buf,
                            tabulation_number + 1,
                            format!("target_id: {target_id}").as_str(),
                        )?;
                    }

                    writeln_with_tabulation(
                        buf,
                        tabulation_number + 1,
                        format!("Position: {position:?}").as_str(),
                    )?;

                    let col_type_str = col_type.to_string();
                    writeln_with_tabulation(
                        buf,
                        tabulation_number + 1,
                        format!("Column type: {col_type_str}").as_str(),
                    )?;
                }
                Expression::SubQueryReference(SubQueryReference {
                    rel_id,
                    position,
                    col_type,
                    ..
                }) => {
                    writeln!(buf, "SubQueryReference")?;
                    let alias = self.get_alias_from_reference_node(&expr);
                    if let Result::Ok(alias) = alias {
                        writeln_with_tabulation(
                            buf,
                            tabulation_number + 1,
                            format!("Alias: {alias}").as_str(),
                        )?;
                    }

                    // See explain logic for Reference node
                    let rel_node = self.get_relation_node(*rel_id);
                    if rel_node.is_ok() {
                        if let Ok(Some(name)) = self.scan_name(*rel_id, *position) {
                            writeln_with_tabulation(
                                buf,
                                tabulation_number + 1,
                                format!("Referenced table name (or alias): {name}").as_str(),
                            )?;
                        }
                    }

                    writeln_with_tabulation(
                        buf,
                        tabulation_number + 1,
                        format!("rel_id: {rel_id}").as_str(),
                    )?;

                    writeln_with_tabulation(
                        buf,
                        tabulation_number + 1,
                        format!("Position: {position:?}").as_str(),
                    )?;

                    let col_type_str = col_type.to_string();
                    writeln_with_tabulation(
                        buf,
                        tabulation_number + 1,
                        format!("Column type: {col_type_str}").as_str(),
                    )?;
                }
                Expression::Row(Row { list, distribution }) => {
                    write!(buf, "Row")?;
                    if let Some(distribution) = distribution {
                        writeln!(buf, " [distribution = {distribution:?}]")?;
                    } else {
                        writeln!(buf)?;
                    }
                    writeln_with_tabulation(buf, tabulation_number + 1, "List:")?;
                    for value in list {
                        self.formatted_arena_node(buf, tabulation_number + 1, *value)?;
                    }
                }
                Expression::Index(_) => writeln!(buf, "Index")?,
                Expression::Cast(_) => writeln!(buf, "Cast")?,
                Expression::Trim(Trim {
                    kind,
                    pattern,
                    target,
                }) => {
                    writeln!(buf, "Trim [kind: {kind:?}]")?;
                    if let Some(pattern) = pattern {
                        writeln_with_tabulation(buf, tabulation_number + 1, "Pattern")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *pattern)?;
                    }
                    writeln_with_tabulation(buf, tabulation_number + 1, "Target")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *target)?;
                }
                Expression::Concat(_) => writeln!(buf, "Concat")?,
                Expression::Like(Like {
                    left,
                    right,
                    escape,
                }) => {
                    writeln!(buf, "Like")?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Left child")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *left)?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Right child")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *right)?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Escape")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *escape)?;
                }
                Expression::ScalarFunction(ScalarFunction { children, .. }) => {
                    writeln!(buf, "ScalarFunction")?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Arguments:")?;
                    for arg in children {
                        self.formatted_arena_node(buf, tabulation_number + 1, *arg)?;
                    }
                }
                Expression::Unary(UnaryExpr { op, child }) => {
                    writeln!(buf, "Unary [op: {op}]")?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Child")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *child)?;
                }
                Expression::Arithmetic(ArithmeticExpr { left, right, op }) => {
                    writeln!(buf, "Arithmetic: [op: {op}]")?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Child")?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Left child")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *left)?;
                    writeln_with_tabulation(buf, tabulation_number + 1, "Right child")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, *right)?;
                }
                Expression::Timestamp(timestamp) => match *timestamp {
                    Timestamp::Date => {
                        writeln!(buf, "Date")?;
                    }
                    Timestamp::DateTime(TimeParameters {
                        precision,
                        include_timezone,
                    }) => {
                        writeln!(buf, "DateTime [precision = {precision}, include_timezone = {include_timezone}]")?;
                    }
                },
                Expression::Parameter(Parameter {
                    param_type,
                    index,
                    unique,
                }) => {
                    writeln!(
                        buf,
                        "Parameter [type = {param_type}, index = {index}, unique: {unique}]"
                    )?;
                }
            };
        }
        Ok(())
    }

    /// Debug function for formatting working arena node.
    #[allow(clippy::too_many_lines)]
    fn formatted_arena_node(
        &self,
        buf: &mut String,
        tabulation_number: i32,
        node_id: NodeId,
    ) -> Result<(), std::fmt::Error> {
        if tabulation_number == 0 {
            writeln!(buf, "---------------------------------------------")?;
        }
        write_with_tabulation(buf, tabulation_number, format!("[id: {node_id}] ").as_str())?;
        let node = self.get_node(node_id).expect("Plan must be valid");
        match node {
            Node::Relational(relation) => {
                write!(buf, "relation: ")?;
                // Print relation name and specific info.
                match relation {
                    Relational::ScanRelation(ScanRelation {
                        alias, relation, ..
                    }) => {
                        writeln!(buf, "ScanRelation")?;
                        writeln_with_tabulation(
                            buf,
                            tabulation_number + 1,
                            format!("Relation: {relation}").as_str(),
                        )?;
                        if let Some(alias) = alias {
                            writeln_with_tabulation(
                                buf,
                                tabulation_number + 1,
                                format!("Alias: {alias}").as_str(),
                            )?;
                        }
                    }
                    Relational::Join(Join { condition, .. }) => {
                        writeln!(buf, "InnerJoin")?;
                        writeln_with_tabulation(buf, tabulation_number + 1, "Condition:")?;
                        self.formatted_arena_node(buf, tabulation_number + 2, *condition)?;
                    }
                    Relational::SelectWithoutScan(_) => {
                        writeln!(buf, "SelectWithoutScan")?;
                    }
                    Relational::Projection(_) => {
                        writeln!(buf, "Projection")?;
                    }
                    Relational::ScanCte(ScanCte { alias, .. }) => {
                        writeln!(buf, "ScanCte")?;
                        if !alias.is_empty() {
                            writeln_with_tabulation(
                                buf,
                                tabulation_number + 1,
                                format!("Alias: {alias}").as_str(),
                            )?;
                        }
                    }
                    Relational::ScanSubQuery(ScanSubQuery { alias, .. }) => {
                        writeln!(buf, "ScanSubQuery")?;
                        if let Some(alias) = alias {
                            if !alias.is_empty() {
                                writeln_with_tabulation(
                                    buf,
                                    tabulation_number + 1,
                                    format!("Alias: {alias}").as_str(),
                                )?;
                            }
                        }
                    }
                    Relational::Selection(Selection {
                        child: _,
                        subqueries: _,
                        filter,
                        output: _,
                    }) => {
                        writeln!(buf, "Selection")?;
                        writeln_with_tabulation(buf, tabulation_number + 1, "Filter")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *filter)?;
                    }
                    Relational::Having(Having { filter, .. }) => {
                        writeln!(buf, "Having")?;
                        writeln_with_tabulation(buf, tabulation_number + 1, "Filter")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *filter)?;
                    }
                    Relational::GroupBy(GroupBy { gr_exprs, .. }) => {
                        writeln!(buf, "GroupBy")?;
                        writeln_with_tabulation(buf, tabulation_number + 1, "Gr_cols:")?;
                        for expr_id in gr_exprs {
                            self.formatted_arena_node(buf, tabulation_number + 2, *expr_id)?;
                        }
                    }
                    Relational::OrderBy(OrderBy {
                        order_by_elements, ..
                    }) => {
                        writeln!(buf, "OrderBy")?;
                        writeln_with_tabulation(buf, tabulation_number + 1, "Order_by_elements:")?;
                        for element in order_by_elements {
                            let order_by_entity_str = match element.entity {
                                OrderByEntity::Expression { expr_id } => {
                                    let order_by_expr = self.get_expression_node(expr_id);
                                    if let Ok(order_by_expr) = order_by_expr {
                                        format!("{order_by_expr:?}")
                                    } else {
                                        "?".to_string()
                                    }
                                }
                                OrderByEntity::Index { value } => format!("{value}"),
                            };
                            let order_by_type = element.order_type.clone();
                            writeln_with_tabulation(buf, tabulation_number + 2, format!("Order_by_element: {order_by_entity_str} [order_type = {order_by_type:?}]").as_str())?;
                        }
                    }
                    Relational::Values { .. } => writeln!(buf, "Values")?,
                    Relational::ValuesRow(ValuesRow { data, .. }) => {
                        writeln!(buf, "ValuesRow")?;
                        writeln_with_tabulation(buf, tabulation_number + 1, "Data")?;
                        self.formatted_arena_node(buf, tabulation_number + 1, *data)?;
                    }
                    Relational::Motion(Motion { policy, alias, .. }) => {
                        write!(buf, "Motion [policy = {policy:?}, alias = ")?;
                        if let Some(alias) = alias {
                            write!(buf, "{alias}")?;
                        } else {
                            write!(buf, "None")?;
                        }
                        writeln!(buf, "]")?;
                    }
                    Relational::Union { .. } => writeln!(buf, "Union")?,
                    Relational::UnionAll { .. } => writeln!(buf, "UnionAll")?,
                    Relational::Update(Update {
                        relation,
                        update_columns_map,
                        ..
                    }) => {
                        writeln!(buf, "Update")?;
                        writeln_with_tabulation(buf, tabulation_number + 1, "Update columns map:")?;
                        for (rel_pos, proj_pos) in update_columns_map {
                            writeln_with_tabulation(buf, tabulation_number + 2, format!("Update {relation} column on pos {rel_pos} to child projection column on pos {proj_pos}").as_str())?;
                        }
                    }
                    Relational::Delete(_) => writeln!(buf, "Delete")?,
                    Relational::Insert(_) => writeln!(buf, "Insert")?,
                    Relational::Intersect(_) => writeln!(buf, "Intersect")?,
                    Relational::Except(_) => writeln!(buf, "Except")?,
                    Relational::Limit(Limit { limit, .. }) => writeln!(buf, "Limit {limit}")?,
                }
                // Print children.
                match relation {
                    Relational::Join(_)
                    | Relational::Projection(_)
                    | Relational::Except(_)
                    | Relational::Delete(_)
                    | Relational::Insert(_)
                    | Relational::Intersect(_)
                    | Relational::ScanSubQuery(_)
                    | Relational::Selection(_)
                    | Relational::SelectWithoutScan(_)
                    | Relational::Values(_)
                    | Relational::OrderBy(_)
                    | Relational::Limit(_)
                    | Relational::Motion(_)
                    | Relational::Union(_)
                    | Relational::UnionAll(_)
                    | Relational::Update(_)
                    | Relational::Having(_)
                    | Relational::GroupBy(_)
                    | Relational::ValuesRow(_) => {
                        writeln_with_tabulation(buf, tabulation_number + 1, "Children:")?;
                        for child in &relation.children() {
                            writeln_with_tabulation(
                                buf,
                                tabulation_number + 2,
                                format!("Child_id = {child}").as_str(),
                            )?;
                        }
                    }
                    Relational::ScanCte(ScanCte { child, .. }) => {
                        writeln_with_tabulation(buf, tabulation_number + 1, "Children:")?;
                        writeln_with_tabulation(
                            buf,
                            tabulation_number + 2,
                            format!("Child_id = {child}").as_str(),
                        )?;
                    }
                    Relational::ScanRelation { .. } => {
                        writeln_with_tabulation(buf, tabulation_number + 1, "[No children]")?;
                    }
                }
                // Print subqueries.
                let subqueries = relation.subqueries();
                if !subqueries.is_empty() {
                    writeln_with_tabulation(buf, tabulation_number + 1, "Subqueries:")?;
                    for subquery in &subqueries {
                        writeln_with_tabulation(
                            buf,
                            tabulation_number + 2,
                            format!("Subquery_id = {subquery}").as_str(),
                        )?;
                    }
                }
                // Print output.
                if let Relational::Delete(Delete { output: None, .. }) = relation {
                    writeln_with_tabulation(buf, tabulation_number + 1, "Output: NONE")?;
                } else {
                    let output = relation.output();
                    write_with_tabulation(buf, tabulation_number + 1, "Output:")?;
                    self.formatted_arena_node(buf, tabulation_number + 1, output)?;
                }
                writeln!(buf, "---------------------------------------------")?;
            }
            Node::Expression(_) => {
                self.write_expr(buf, tabulation_number, node_id)?;
            }
            Node::Invalid(_) => {
                writeln!(buf, "INVALID")?;
                writeln!(buf, "---------------------------------------------")?;
            }
            _ => {}
        }
        Ok(())
    }

    /// Pretty-printing plan `arena` subtree.
    ///
    /// # Errors
    /// Fail to format one of the relational node.
    pub fn formatted_arena_subtree(&self, node_id: NodeId) -> Result<String, SbroadError> {
        let ir_tree = PostOrder::new(|node| self.nodes.rel_iter(node), EXPR_CAPACITY);

        let mut buf = String::new();
        for level_node in ir_tree.traverse_into_iter(node_id) {
            let id = level_node.1;
            if self.formatted_arena_node(&mut buf, 0, id).is_err() {
                return Err(SbroadError::FailedTo(
                    Action::Serialize,
                    Some(Entity::Plan),
                    "Unable to get formatted arena string".to_smolstr(),
                ));
            }
        }
        Ok(buf)
    }

    /// Pretty-printing plan `arena`.
    ///
    /// # Errors
    /// Fail to format one of the relational node.
    pub fn formatted_arena(&self) -> Result<String, SbroadError> {
        let top_id = self.get_top()?;
        self.formatted_arena_subtree(top_id)
    }
}

/// Formatting helper debug functions
impl SyntaxPlan<'_> {
    /// Pretty-printing plan `nodes.arena`.
    ///
    /// # Errors
    /// - Failed to retrieve top node.
    /// - Failed to format one of the nodes.
    pub fn formatted(&self, plan: &Plan) -> Result<String, SbroadError> {
        let mut buf = String::new();
        let top = self.top.ok_or_else(|| {
            SbroadError::Invalid(
                Entity::SyntaxPlan,
                Some(SmolStr::from(
                    "Top id should have already be set for SyntaxPlan",
                )),
            )
        })?;
        if self.formatted_inner(plan, &mut buf, 0, top).is_err() {
            return Err(SbroadError::FailedTo(
                Action::Serialize,
                Some(Entity::SyntaxPlan),
                "Unable to get formatted arena string".to_smolstr(),
            ));
        }
        Ok(buf)
    }

    fn formatted_inner(
        &self,
        plan: &Plan,
        buf: &mut String,
        tabulation_number: i32,
        node_id: usize,
    ) -> Result<(), std::fmt::Error> {
        let node = self.nodes.arena.get(node_id);
        if let Some(node) = node {
            formatted_tabulate(buf, tabulation_number)?;
            let data = &node.data;
            if let SyntaxData::PlanId(id) = data {
                let node = plan.get_node(*id);
                if let Ok(node) = node {
                    writeln!(buf, "{node:?} [id={id:?}]")?;
                }
            } else {
                writeln!(buf, "{data:?}")?;
            }

            if let Some(left_id) = node.left {
                writeln_with_tabulation(buf, tabulation_number + 1, "Left:")?;
                self.formatted_inner(plan, buf, tabulation_number + 2, left_id)?;
            }
            if !node.right.is_empty() {
                writeln_with_tabulation(buf, tabulation_number + 1, "Right:")?;
            }
            for right in &node.right {
                self.formatted_inner(plan, buf, tabulation_number + 2, *right)?;
            }
        } else {
            writeln_with_tabulation(buf, tabulation_number, "MISSING")?;
        }
        Ok(())
    }
}

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