// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{convert::TryInto, mem};

use kvproto::{
    coprocessor::{KeyRange, Request},
    kvrpcpb::Context,
};
use protobuf::Message;
use tidb_query_datatype::codec::{Datum, datum};
use tikv::coprocessor::REQ_TYPE_DAG;
use tikv_util::codec::number::NumberEncoder;
use tipb::{
    Aggregation, ByItem, Chunk, ColumnInfo, DagRequest, ExecType, Executor, Expr, ExprType,
    IndexLookUp, IndexScan, IntermediateOutputChannel, Limit, Selection, TableScan, TopN,
};

use super::*;

pub struct DagSelect {
    pub execs: Vec<Executor>,
    pub cols: Vec<ColumnInfo>,
    pub order_by: Vec<ByItem>,
    pub limit: Option<u64>,
    pub aggregate: Vec<Expr>,
    pub group_by: Vec<Expr>,
    pub key_ranges: Vec<KeyRange>,
    pub output_offsets: Option<Vec<u32>>,
    pub paging_size: Option<u64>,
    pub paging_size_bytes: Option<u64>,
    pub start_ts: Option<u64>,
    pub intermediate_outputs: Vec<IntermediateOutputChannel>,
}

impl DagSelect {
    pub fn from(table: &Table) -> DagSelect {
        let mut exec = Executor::default();
        exec.set_tp(ExecType::TypeTableScan);
        let mut tbl_scan = TableScan::default();
        let mut table_info = table.table_info();
        tbl_scan.set_table_id(table_info.get_table_id());
        let columns_info = table_info.take_columns();
        tbl_scan.set_columns(columns_info);
        exec.set_tbl_scan(tbl_scan);

        DagSelect {
            execs: vec![exec],
            cols: table.columns_info(),
            order_by: vec![],
            limit: None,
            aggregate: vec![],
            group_by: vec![],
            key_ranges: vec![table.get_record_range_all()],
            output_offsets: None,
            paging_size: None,
            paging_size_bytes: None,
            start_ts: None,
            intermediate_outputs: vec![],
        }
    }

    pub fn from_index(table: &Table, index: &Column) -> DagSelect {
        let idx = index.index;
        let mut exec = Executor::default();
        exec.set_tp(ExecType::TypeIndexScan);
        let mut scan = IndexScan::default();
        let mut index_info = table.index_info(idx, true);
        scan.set_table_id(index_info.get_table_id());
        scan.set_index_id(idx);

        let columns_info = index_info.take_columns();
        scan.set_columns(columns_info.clone());
        exec.set_idx_scan(scan);

        let range = table.get_index_range_all(idx);
        DagSelect {
            execs: vec![exec],
            cols: columns_info.to_vec(),
            order_by: vec![],
            limit: None,
            aggregate: vec![],
            group_by: vec![],
            key_ranges: vec![range],
            output_offsets: None,
            paging_size: None,
            paging_size_bytes: None,
            start_ts: None,
            intermediate_outputs: vec![],
        }
    }

    pub fn index_lookup(mut self, table: &Table, handle_offset: Vec<u32>) -> DagSelect {
        if let Some(l) = self.limit.take() {
            let mut exec = Executor::default();
            exec.set_tp(ExecType::TypeLimit);
            exec.set_limit({
                let mut limit = Limit::default();
                limit.set_limit(l);
                limit
            });
            self.execs.push(exec);
        }
        let exec_len = self.execs.len() + 2;
        self.execs[exec_len - 3].set_parent_idx(exec_len as u32 - 1);
        let table_scan_dag = Self::from(table);
        self.execs.push(table_scan_dag.execs[0].clone());
        self.execs.push({
            let mut exec = Executor::default();
            exec.set_tp(ExecType::TypeIndexLookUp);
            exec.set_index_lookup({
                let mut index_lookup = IndexLookUp::default();
                index_lookup.set_index_handle_offsets(handle_offset);
                index_lookup
            });
            exec
        });
        self.intermediate_outputs.push({
            let mut channel = IntermediateOutputChannel::default();
            channel.set_executor_idx(exec_len as u32 - 1);
            if let Some(output_offsets) = self.output_offsets.take() {
                channel.set_output_offsets(output_offsets);
            } else {
                channel.set_output_offsets((0..self.cols.len() as u32).collect());
            }
            channel
        });
        self.cols = table_scan_dag.cols;
        self.output_offsets = table_scan_dag.output_offsets;
        self
    }

    #[must_use]
    pub fn limit(mut self, n: u64) -> DagSelect {
        self.limit = Some(n);
        self
    }

    #[must_use]
    pub fn order_by(mut self, col: &Column, desc: bool) -> DagSelect {
        let col_offset = offset_for_column(&self.cols, col.id);
        let mut item = ByItem::default();
        let mut expr = Expr::default();
        expr.set_field_type(col.as_field_type());
        expr.set_tp(ExprType::ColumnRef);
        expr.mut_val().encode_i64(col_offset).unwrap();
        item.set_expr(expr);
        item.set_desc(desc);
        self.order_by.push(item);
        self
    }

    #[must_use]
    pub fn count(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::Count)
    }

    #[must_use]
    pub fn aggr_col(mut self, col: &Column, aggr_t: ExprType) -> DagSelect {
        let col_offset = offset_for_column(&self.cols, col.id);
        let mut col_expr = Expr::default();
        col_expr.set_field_type(col.as_field_type());
        col_expr.set_tp(ExprType::ColumnRef);
        col_expr.mut_val().encode_i64(col_offset).unwrap();
        let mut expr = Expr::default();
        let mut expr_ft = col.as_field_type();
        // Avg will contains two auxiliary columns (sum, count) and the sum should be a
        // `Decimal`
        if aggr_t == ExprType::Avg || aggr_t == ExprType::Sum {
            expr_ft.set_tp(0xf6); // FieldTypeTp::NewDecimal
        }
        expr.set_field_type(expr_ft);
        expr.set_tp(aggr_t);
        expr.mut_children().push(col_expr);
        self.aggregate.push(expr);
        self
    }

    #[must_use]
    pub fn first(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::First)
    }

    #[must_use]
    pub fn sum(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::Sum)
    }

    #[must_use]
    pub fn avg(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::Avg)
    }

    #[must_use]
    pub fn max(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::Max)
    }

    #[must_use]
    pub fn min(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::Min)
    }

    #[must_use]
    pub fn bit_and(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::AggBitAnd)
    }

    #[must_use]
    pub fn bit_or(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::AggBitOr)
    }

    #[must_use]
    pub fn bit_xor(self, col: &Column) -> DagSelect {
        self.aggr_col(col, ExprType::AggBitXor)
    }

    #[must_use]
    pub fn group_by(mut self, cols: &[&Column]) -> DagSelect {
        for col in cols {
            let offset = offset_for_column(&self.cols, col.id);
            let mut expr = Expr::default();
            expr.set_field_type(col.as_field_type());
            expr.set_tp(ExprType::ColumnRef);
            expr.mut_val().encode_i64(offset).unwrap();
            self.group_by.push(expr);
        }
        self
    }

    #[must_use]
    pub fn output_offsets(mut self, output_offsets: Option<Vec<u32>>) -> DagSelect {
        self.output_offsets = output_offsets;
        self
    }

    #[must_use]
    pub fn where_expr(mut self, expr: Expr) -> DagSelect {
        let mut exec = Executor::default();
        exec.set_tp(ExecType::TypeSelection);
        let mut selection = Selection::default();
        selection.mut_conditions().push(expr);
        exec.set_selection(selection);
        self.execs.push(exec);
        self
    }

    #[must_use]
    pub fn projection(mut self, exprs: Vec<Expr>) -> DagSelect {
        let mut exec = Executor::default();
        exec.set_tp(ExecType::TypeProjection);
        let mut projection = tipb::Projection::default();
        projection.set_exprs(exprs.into());
        exec.set_projection(projection);
        self.execs.push(exec);
        self
    }

    #[must_use]
    pub fn desc(mut self, desc: bool) -> DagSelect {
        self.execs[0].mut_tbl_scan().set_desc(desc);
        self
    }

    #[must_use]
    pub fn paging_size(mut self, paging_size: u64) -> DagSelect {
        assert_ne!(paging_size, 0);
        self.paging_size = Some(paging_size);
        self
    }

    #[must_use]
    pub fn paging_size_bytes(mut self, paging_size_bytes: u64) -> DagSelect {
        assert_ne!(paging_size_bytes, 0);
        self.paging_size_bytes = Some(paging_size_bytes);
        self
    }

    #[must_use]
    pub fn key_ranges(mut self, key_ranges: Vec<KeyRange>) -> DagSelect {
        self.key_ranges = key_ranges;
        self
    }

    #[must_use]
    pub fn start_ts(mut self, start_ts: txn_types::TimeStamp) -> DagSelect {
        self.start_ts = Some(start_ts.into_inner());
        self
    }

    pub fn build(self) -> Request {
        let start_ts = self
            .start_ts
            .unwrap_or_else(|| next_id().try_into().unwrap());
        self.build_impl(Context::default(), &[0], start_ts)
    }

    pub fn build_with(self, ctx: Context, flags: &[u64]) -> Request {
        let start_ts = self
            .start_ts
            .unwrap_or_else(|| next_id().try_into().unwrap());
        self.build_impl(ctx, flags, start_ts)
    }

    fn build_impl(mut self, ctx: Context, flags: &[u64], start_ts: u64) -> Request {
        if !self.aggregate.is_empty() || !self.group_by.is_empty() {
            let mut exec = Executor::default();
            exec.set_tp(ExecType::TypeAggregation);
            let mut aggr = Aggregation::default();
            if !self.aggregate.is_empty() {
                aggr.set_agg_func(self.aggregate.into());
            }

            if !self.group_by.is_empty() {
                aggr.set_group_by(self.group_by.into());
            }
            exec.set_aggregation(aggr);
            self.execs.push(exec);
        }

        if !self.order_by.is_empty() {
            let mut exec = Executor::default();
            exec.set_tp(ExecType::TypeTopN);
            let mut topn = TopN::default();
            topn.set_order_by(self.order_by.into());
            if let Some(limit) = self.limit.take() {
                topn.set_limit(limit);
            }
            exec.set_top_n(topn);
            self.execs.push(exec);
        }

        if let Some(l) = self.limit.take() {
            let mut exec = Executor::default();
            exec.set_tp(ExecType::TypeLimit);
            let mut limit = Limit::default();
            limit.set_limit(l);
            exec.set_limit(limit);
            self.execs.push(exec);
        }

        let mut dag = DagRequest::default();
        dag.set_executors(self.execs.into());
        dag.set_flags(flags.iter().fold(0, |acc, f| acc | *f));
        dag.set_collect_range_counts(true);

        let output_offsets = if self.output_offsets.is_some() {
            self.output_offsets.take().unwrap()
        } else {
            (0..self.cols.len() as u32).collect()
        };
        dag.set_output_offsets(output_offsets);
        dag.set_intermediate_output_channels(mem::take(&mut self.intermediate_outputs).into());

        let mut req = Request::default();
        req.set_start_ts(start_ts);
        req.set_tp(REQ_TYPE_DAG);
        req.set_data(dag.write_to_bytes().unwrap());
        req.set_ranges(self.key_ranges.into());
        req.set_paging_size(self.paging_size.unwrap_or(0));
        req.set_paging_size_bytes(self.paging_size_bytes.unwrap_or(0));
        req.set_context(ctx);
        req
    }
}

pub struct DagChunkSpliter {
    chunks: Vec<Chunk>,
    datums: Vec<Datum>,
    col_cnt: usize,
}

impl DagChunkSpliter {
    pub fn new(chunks: Vec<Chunk>, col_cnt: usize) -> DagChunkSpliter {
        DagChunkSpliter {
            chunks,
            col_cnt,
            datums: Vec::with_capacity(0),
        }
    }
}

impl Iterator for DagChunkSpliter {
    type Item = Vec<Datum>;

    fn next(&mut self) -> Option<Vec<Datum>> {
        loop {
            if self.chunks.is_empty() && self.datums.is_empty() {
                return None;
            } else if self.datums.is_empty() {
                let chunk = self.chunks.remove(0);
                let mut data = chunk.get_rows_data();
                self.datums = datum::decode(&mut data).unwrap();
                continue;
            }
            assert_eq!(self.datums.len() >= self.col_cnt, true);
            let mut cols = self.datums.split_off(self.col_cnt);
            std::mem::swap(&mut self.datums, &mut cols);
            return Some(cols);
        }
    }
}