// 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::rc::Rc;

use risingwave_common::catalog::FragmentTypeMask;
use risingwave_pb::id::FragmentId;
use risingwave_pb::stream_plan::stream_fragment_graph::{
    StreamFragment as StreamFragmentProto, StreamFragmentEdge as StreamFragmentEdgeProto,
};
use risingwave_pb::stream_plan::{DispatchStrategy, StreamNode};

pub type LocalFragmentId = FragmentId;

/// [`StreamFragment`] represent a fragment node in fragment DAG.
#[derive(Clone, Debug)]
pub struct StreamFragment {
    /// the allocated fragment id.
    pub fragment_id: LocalFragmentId,

    /// root stream node in this fragment.
    pub node: Option<Box<StreamNode>>,

    /// Bitwise-OR of type Flags of this fragment.
    pub fragment_type_mask: FragmentTypeMask,

    /// Mark whether this fragment requires exactly one actor.
    pub requires_singleton: bool,
}

/// An edge between the nodes in the fragment graph.
#[derive(Debug, Clone)]
pub struct StreamFragmentEdge {
    /// Dispatch strategy for the fragment.
    pub dispatch_strategy: DispatchStrategy,

    /// A unique identifier of this edge. Generally it should be exchange node's operator id. When
    /// rewriting fragments into delta joins or when inserting 1-to-1 exchange, there will be
    /// virtual links generated.
    pub link_id: u64,
}

impl StreamFragment {
    pub fn new(fragment_id: LocalFragmentId) -> Self {
        Self {
            fragment_id,
            fragment_type_mask: FragmentTypeMask::empty(),
            requires_singleton: false,
            node: None,
        }
    }

    pub fn to_protobuf(&self) -> StreamFragmentProto {
        StreamFragmentProto {
            fragment_id: self.fragment_id,
            node: self.node.clone().map(|n| *n),
            fragment_type_mask: self.fragment_type_mask.into(),
            requires_singleton: self.requires_singleton,
        }
    }
}

/// [`StreamFragmentGraph`] stores a fragment graph (DAG).
#[derive(Default)]
pub struct StreamFragmentGraph {
    /// stores all the fragments in the graph.
    pub(in crate::stream_fragmenter) fragments: HashMap<LocalFragmentId, Rc<StreamFragment>>,

    /// stores edges between fragments: (upstream, downstream) => edge.
    pub(in crate::stream_fragmenter) edges:
        HashMap<(LocalFragmentId, LocalFragmentId), StreamFragmentEdgeProto>,
}

impl StreamFragmentGraph {
    /// Adds a fragment to the graph.
    pub fn add_fragment(&mut self, stream_fragment: Rc<StreamFragment>) {
        let id = stream_fragment.fragment_id;
        let ret = self.fragments.insert(id, stream_fragment);
        assert!(ret.is_none(), "fragment already exists: {:?}", id);
    }

    pub fn get_fragment(&self, fragment_id: &LocalFragmentId) -> Option<&Rc<StreamFragment>> {
        self.fragments.get(fragment_id)
    }

    /// Links upstream to downstream in the graph.
    pub fn add_edge(
        &mut self,
        upstream_id: LocalFragmentId,
        downstream_id: LocalFragmentId,
        edge: StreamFragmentEdge,
    ) {
        self.try_add_edge(upstream_id, downstream_id, edge).unwrap();
    }

    /// Try to link upstream to downstream in the graph.
    ///
    /// If the edge between upstream and downstream already exists, return an error.
    pub fn try_add_edge(
        &mut self,
        upstream_id: LocalFragmentId,
        downstream_id: LocalFragmentId,
        edge: StreamFragmentEdge,
    ) -> Result<(), String> {
        let edge = StreamFragmentEdgeProto {
            upstream_id,
            downstream_id,
            dispatch_strategy: Some(edge.dispatch_strategy),
            link_id: edge.link_id,
        };

        self.edges
            .try_insert((upstream_id, downstream_id), edge)
            .map(|_| ())
            .map_err(|e| {
                format!(
                    "edge between {upstream_id} and {downstream_id} already exists (existing: {:?})",
                    e.entry.get(),
                )
            })
    }
}
