// Copyright 2022 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::vec;

use itertools::Itertools;
use risingwave_common::catalog::{DatabaseId, FragmentTypeFlag, SchemaId, TableId};
use risingwave_common::hash::VirtualNode;
use risingwave_pb::catalog::PbTable;
use risingwave_pb::common::{PbColumnOrder, PbDirection, PbNullsAre, PbOrderType};
use risingwave_pb::data::DataType;
use risingwave_pb::data::data_type::TypeName;
use risingwave_pb::ddl_service::TableJobType;
use risingwave_pb::expr::agg_call::PbKind as PbAggKind;
use risingwave_pb::expr::expr_node::RexNode;
use risingwave_pb::expr::expr_node::Type::{Add, GreaterThan};
use risingwave_pb::expr::{AggCall, ExprNode, FunctionCall, PbInputRef};
use risingwave_pb::plan_common::{ColumnCatalog, ColumnDesc, Field};
use risingwave_pb::stream_plan::stream_fragment_graph::{StreamFragment, StreamFragmentEdge};
use risingwave_pb::stream_plan::stream_node::NodeBody;
use risingwave_pb::stream_plan::{
    AggCallState, DispatchStrategy, DispatcherType, ExchangeNode, FilterNode, MaterializeNode,
    PbDispatchOutputMapping, PbStreamContext, ProjectNode, SimpleAggNode, SourceNode,
    StreamFragmentGraph as StreamFragmentGraphProto, StreamNode, StreamSource, agg_call_state,
};

use crate::MetaResult;
use crate::manager::{MetaSrvEnv, StreamingJob};
use crate::model::StreamJobFragments;
use crate::stream::{
    ActorGraphBuildResult, ActorGraphBuilder, CompleteStreamFragmentGraph, StreamFragmentGraph,
};

fn make_inputref(idx: u32) -> ExprNode {
    ExprNode {
        function_type: PbAggKind::Unspecified as i32,
        return_type: Some(DataType {
            type_name: TypeName::Int32 as i32,
            ..Default::default()
        }),
        rex_node: Some(RexNode::InputRef(idx)),
    }
}

fn make_sum_aggcall(idx: u32) -> AggCall {
    AggCall {
        kind: PbAggKind::Sum as i32,
        args: vec![PbInputRef {
            index: idx,
            r#type: Some(DataType {
                type_name: TypeName::Int64 as i32,
                ..Default::default()
            }),
        }],
        return_type: Some(DataType {
            type_name: TypeName::Int64 as i32,
            ..Default::default()
        }),
        distinct: false,
        order_by: vec![],
        filter: None,
        direct_args: vec![],
        udf: None,
        scalar: None,
    }
}

fn make_agg_call_result_state() -> AggCallState {
    AggCallState {
        inner: Some(agg_call_state::Inner::ValueState(
            agg_call_state::ValueState {},
        )),
    }
}

fn make_field(type_name: TypeName) -> Field {
    Field {
        data_type: Some(DataType {
            type_name: type_name as i32,
            ..Default::default()
        }),
        ..Default::default()
    }
}

fn make_column_order(column_index: u32) -> PbColumnOrder {
    PbColumnOrder {
        column_index,
        order_type: Some(PbOrderType {
            direction: PbDirection::Ascending as _,
            nulls_are: PbNullsAre::Largest as _,
        }),
    }
}

fn make_column(column_type: TypeName, column_id: i32) -> ColumnCatalog {
    ColumnCatalog {
        column_desc: Some(ColumnDesc {
            column_type: Some(DataType {
                type_name: column_type as i32,
                ..Default::default()
            }),
            column_id,
            ..Default::default()
        }),
        is_hidden: false,
    }
}

fn make_source_internal_table(id: u32) -> PbTable {
    let columns = vec![
        make_column(TypeName::Varchar, 0),
        make_column(TypeName::Varchar, 1),
    ];
    PbTable {
        id: id.into(),
        schema_id: SchemaId::placeholder(),
        database_id: DatabaseId::placeholder(),
        name: String::new(),
        columns,
        pk: vec![PbColumnOrder {
            column_index: 0,
            order_type: Some(PbOrderType {
                direction: PbDirection::Descending as _,
                nulls_are: PbNullsAre::Largest as _,
            }),
        }],
        ..Default::default()
    }
}

fn make_internal_table(id: u32, is_agg_value: bool) -> PbTable {
    let mut columns = vec![make_column(TypeName::Int64, 0)];
    if !is_agg_value {
        columns.push(make_column(TypeName::Int32, 1));
    }
    PbTable {
        id: id.into(),
        schema_id: SchemaId::placeholder(),
        database_id: DatabaseId::placeholder(),
        name: String::new(),
        columns,
        pk: vec![PbColumnOrder {
            column_index: 0,
            order_type: Some(PbOrderType {
                direction: PbDirection::Descending as _,
                nulls_are: PbNullsAre::Largest as _,
            }),
        }],
        stream_key: vec![2],
        ..Default::default()
    }
}

fn make_empty_table(id: u32) -> PbTable {
    PbTable {
        id: id.into(),
        schema_id: SchemaId::placeholder(),
        database_id: DatabaseId::placeholder(),
        name: String::new(),
        columns: vec![],
        pk: vec![],
        stream_key: vec![],
        ..Default::default()
    }
}

fn make_materialize_table(id: u32) -> PbTable {
    make_internal_table(id, true)
}

/// [`make_stream_fragments`] build all stream fragments for SQL as follow:
/// ```sql
/// create table t (v1 int, v2 int);
/// create materialized view T_distributed as select sum(v1)+1 as V from t where v1>v2;
/// ```
fn make_stream_fragments() -> Vec<StreamFragment> {
    let mut fragments = vec![];
    // table source node
    let column_ids = [1, 2, 0];
    let columns = column_ids
        .iter()
        .map(|column_id| ColumnCatalog {
            column_desc: Some(ColumnDesc {
                column_id: *column_id,
                ..Default::default()
            }),
            ..Default::default()
        })
        .collect_vec();
    let source_node = StreamNode {
        node_body: Some(NodeBody::Source(Box::new(SourceNode {
            source_inner: Some(StreamSource {
                source_id: 1.into(),
                state_table: Some(make_source_internal_table(0)),
                columns,
                ..Default::default()
            }),
        }))),
        stream_key: vec![2],
        ..Default::default()
    };
    fragments.push(StreamFragment {
        fragment_id: 2.into(),
        node: Some(source_node),
        fragment_type_mask: FragmentTypeFlag::Source as u32,
        requires_singleton: false,
    });

    // exchange node
    let exchange_node = StreamNode {
        node_body: Some(NodeBody::Exchange(Box::new(ExchangeNode {
            strategy: Some(DispatchStrategy {
                r#type: DispatcherType::Hash as i32,
                dist_key_indices: vec![0],
                output_mapping: PbDispatchOutputMapping::simple(vec![0, 1, 2]).into(),
            }),
        }))),
        fields: vec![
            make_field(TypeName::Int32),
            make_field(TypeName::Int32),
            make_field(TypeName::Int64),
        ],
        input: vec![],
        stream_key: vec![2],
        operator_id: 1.into(),
        identity: "ExchangeExecutor".to_owned(),
        ..Default::default()
    };

    // filter node
    let function_call = FunctionCall {
        children: vec![make_inputref(0), make_inputref(1)],
    };
    let filter_node = StreamNode {
        node_body: Some(NodeBody::Filter(Box::new(FilterNode {
            search_condition: Some(ExprNode {
                function_type: GreaterThan as i32,
                return_type: Some(DataType {
                    type_name: TypeName::Boolean as i32,
                    ..Default::default()
                }),
                rex_node: Some(RexNode::FuncCall(function_call)),
            }),
        }))),
        fields: vec![], // TODO: fill this later
        input: vec![exchange_node],
        stream_key: vec![0, 1],
        operator_id: 2.into(),
        identity: "FilterExecutor".to_owned(),
        ..Default::default()
    };

    // simple agg node
    let simple_agg_node = StreamNode {
        node_body: Some(NodeBody::SimpleAgg(Box::new(SimpleAggNode {
            agg_calls: vec![make_sum_aggcall(0), make_sum_aggcall(1)],
            is_append_only: false,
            agg_call_states: vec![make_agg_call_result_state(), make_agg_call_result_state()],
            intermediate_state_table: Some(make_empty_table(1)),
            ..Default::default()
        }))),
        input: vec![filter_node],
        fields: vec![], // TODO: fill this later
        stream_key: vec![0, 1],
        operator_id: 3.into(),
        identity: "SimpleAggExecutor".to_owned(),
        ..Default::default()
    };

    fragments.push(StreamFragment {
        fragment_id: 1.into(),
        node: Some(simple_agg_node),
        fragment_type_mask: 0,
        requires_singleton: false,
    });

    // exchange node
    let exchange_node_1 = StreamNode {
        node_body: Some(NodeBody::Exchange(Box::new(ExchangeNode {
            strategy: Some(DispatchStrategy {
                r#type: DispatcherType::Simple as i32,
                ..Default::default()
            }),
        }))),
        fields: vec![make_field(TypeName::Int64), make_field(TypeName::Int64)],
        input: vec![],
        stream_key: vec![0, 1],
        operator_id: 4.into(),
        identity: "ExchangeExecutor".to_owned(),
        ..Default::default()
    };

    // agg node
    let simple_agg_node_1 = StreamNode {
        node_body: Some(NodeBody::SimpleAgg(Box::new(SimpleAggNode {
            agg_calls: vec![make_sum_aggcall(0), make_sum_aggcall(1)],
            is_append_only: false,
            agg_call_states: vec![make_agg_call_result_state(), make_agg_call_result_state()],
            intermediate_state_table: Some(make_empty_table(2)),
            ..Default::default()
        }))),
        fields: vec![], // TODO: fill this later
        input: vec![exchange_node_1],
        stream_key: vec![0, 1],
        operator_id: 5.into(),
        identity: "SimpleAggExecutor".to_owned(),
        ..Default::default()
    };

    // project node
    let function_call_1 = FunctionCall {
        children: vec![make_inputref(0), make_inputref(1)],
    };
    let project_node = StreamNode {
        node_body: Some(NodeBody::Project(Box::new(ProjectNode {
            select_list: vec![
                ExprNode {
                    rex_node: Some(RexNode::FuncCall(function_call_1)),
                    function_type: Add as i32,
                    return_type: Some(DataType {
                        type_name: TypeName::Int64 as i32,
                        ..Default::default()
                    }),
                },
                make_inputref(0),
                make_inputref(1),
            ],
            ..Default::default()
        }))),
        fields: vec![], // TODO: fill this later
        input: vec![simple_agg_node_1],
        stream_key: vec![1, 2],
        operator_id: 6.into(),
        identity: "ProjectExecutor".to_owned(),
        ..Default::default()
    };

    // mview node
    let mview_node = StreamNode {
        input: vec![project_node],
        stream_key: vec![],
        node_body: Some(NodeBody::Materialize(Box::new(MaterializeNode {
            // `table_id` and `table` are left empty when generated from frontend.
            table_id: TableId::placeholder(),
            table: None,
            column_orders: vec![make_column_order(1), make_column_order(2)],
            staging_table: None,
            refresh_progress_table: None,
            cleaned_by_ttl_watermark: false,
        }))),
        fields: vec![], // TODO: fill this later
        operator_id: 7.into(),
        identity: "MaterializeExecutor".to_owned(),
        ..Default::default()
    };

    fragments.push(StreamFragment {
        fragment_id: 0.into(),
        node: Some(mview_node),
        fragment_type_mask: FragmentTypeFlag::Mview as u32,
        requires_singleton: true,
    });

    fragments
}

fn make_fragment_edges() -> Vec<StreamFragmentEdge> {
    vec![
        StreamFragmentEdge {
            dispatch_strategy: Some(DispatchStrategy {
                r#type: DispatcherType::Simple as i32,
                dist_key_indices: vec![],
                output_mapping: PbDispatchOutputMapping::identical(0).into(), /* dummy length as it's not used */
            }),
            link_id: 4,
            upstream_id: 1.into(),
            downstream_id: 0.into(),
        },
        StreamFragmentEdge {
            dispatch_strategy: Some(DispatchStrategy {
                r#type: DispatcherType::Hash as i32,
                dist_key_indices: vec![0],
                output_mapping: PbDispatchOutputMapping::identical(0).into(), /* dummy length as it's not used */
            }),
            link_id: 1,
            upstream_id: 2.into(),
            downstream_id: 1.into(),
        },
    ]
}

fn make_stream_graph() -> StreamFragmentGraphProto {
    let fragments = make_stream_fragments();
    StreamFragmentGraphProto {
        fragments: HashMap::from_iter(fragments.into_iter().map(|f| (f.fragment_id, f))),
        edges: make_fragment_edges(),
        ctx: Some(PbStreamContext::default()),
        dependent_table_ids: vec![],
        table_ids_cnt: 3,
        parallelism: None,
        backfill_parallelism: None,
        adaptive_parallelism_strategy: Default::default(),
        backfill_adaptive_parallelism_strategy: Default::default(),
        max_parallelism: VirtualNode::COUNT_FOR_TEST as _,
        backfill_order: Default::default(),
    }
}

#[tokio::test]
async fn test_graph_builder() -> MetaResult<()> {
    let env = MetaSrvEnv::for_test().await;
    let job = StreamingJob::Table(None, make_materialize_table(888), TableJobType::General);

    let graph = make_stream_graph();
    let fragment_graph = StreamFragmentGraph::new(&env, graph, &job)?;
    let internal_tables = fragment_graph.incomplete_internal_tables();

    let actor_graph_builder =
        ActorGraphBuilder::new(CompleteStreamFragmentGraph::for_test(fragment_graph))?;
    let ActorGraphBuildResult {
        graph,
        upstream_fragment_downstreams,
        downstream_fragment_relations,
        ..
    } = actor_graph_builder.generate_graph()?;

    let new_fragment_relation = || {
        upstream_fragment_downstreams
            .iter()
            .chain(downstream_fragment_relations.iter())
            .flat_map(|(fragment_id, downstreams)| {
                downstreams
                    .iter()
                    .map(|relation| (*fragment_id, relation.downstream_fragment_id))
            })
    };

    let stream_job_fragments = StreamJobFragments::for_test(0.into(), graph);
    let mview_fragment_ids = stream_job_fragments.mview_fragment_ids();

    assert_eq!(mview_fragment_ids.len(), 1);

    assert_eq!(internal_tables.len(), 3);

    for fragment in stream_job_fragments.fragments() {
        let mut node = &fragment.nodes;
        while !node.get_input().is_empty() {
            node = node.get_input().first().unwrap();
        }
        match node.get_node_body().unwrap() {
            NodeBody::Merge(merge_node) => {
                assert!(
                    new_fragment_relation().any(|(upstream_fragment_id, fragment_id)| {
                        upstream_fragment_id == merge_node.upstream_fragment_id
                            && fragment_id == fragment.fragment_id
                    })
                );
            }
            NodeBody::Source(_) => {
                // check nothing.
            }
            _ => {
                panic!("it should be MergeNode or SourceNode.");
            }
        }
    }

    Ok(())
}
