// Copyright 2025 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 risingwave_common::catalog::FragmentTypeFlag;
use risingwave_common::id::{FragmentId, JobId, TableId};
use risingwave_common::types::Fields;
use risingwave_common::util::stream_graph_visitor::{
    visit_stream_node_body, visit_stream_node_source_backfill, visit_stream_node_stream_scan,
};
use risingwave_frontend_macro::system_catalog;
use risingwave_pb::id::RelationId;
use risingwave_pb::meta::FragmentDistribution;
use risingwave_pb::stream_plan::stream_node::NodeBody;

use crate::catalog::system_catalog::SysCatalogReaderImpl;
use crate::catalog::system_catalog::rw_catalog::common::CatalogBackfillType;
use crate::error::Result;

#[derive(Fields)]
struct RwBackfillInfo {
    job_id: JobId,
    #[primary_key]
    fragment_id: FragmentId,
    backfill_state_table_id: TableId,
    backfill_target_relation_id: RelationId,
    backfill_type: String,
    backfill_epoch: i64,
}

fn extract_stream_scan(fragment_distribution: &FragmentDistribution) -> Option<RwBackfillInfo> {
    let fragment_type_mask = fragment_distribution.fragment_type_mask;
    let is_source_backfill = fragment_type_mask & (FragmentTypeFlag::SourceScan as u32) != 0;
    let is_snapshot_backfill = fragment_type_mask
        & (FragmentTypeFlag::SnapshotBackfillStreamScan as u32
            | FragmentTypeFlag::CrossDbSnapshotBackfillStreamScan as u32)
        != 0;
    let is_arrangement_or_no_shuffle =
        fragment_type_mask & (FragmentTypeFlag::StreamScan as u32) != 0;
    let is_locality_backfill =
        fragment_type_mask & (FragmentTypeFlag::LocalityProvider as u32) != 0;

    let backfill_type = if is_source_backfill {
        CatalogBackfillType::Source
    } else if is_snapshot_backfill {
        CatalogBackfillType::SnapshotBackfill
    } else if is_arrangement_or_no_shuffle || is_locality_backfill {
        CatalogBackfillType::ArrangementOrNoShuffle
    } else {
        return None;
    };

    let stream_node = fragment_distribution.node.as_ref()?;

    let mut scan = None;
    match backfill_type {
        CatalogBackfillType::Source => {
            visit_stream_node_source_backfill(stream_node, |node| {
                scan = Some(RwBackfillInfo {
                    job_id: fragment_distribution.table_id,
                    fragment_id: fragment_distribution.fragment_id,
                    backfill_state_table_id: node
                        .state_table
                        .as_ref()
                        .map(|table| table.id)
                        .unwrap_or(TableId::placeholder()),
                    backfill_target_relation_id: node.upstream_source_id.as_relation_id(),
                    backfill_type: backfill_type.to_string(),
                    backfill_epoch: 0,
                });
            });
        }
        CatalogBackfillType::SnapshotBackfill | CatalogBackfillType::ArrangementOrNoShuffle => {
            if is_locality_backfill {
                let mut backfill_state_table_id = None;
                visit_stream_node_body(stream_node, |body| {
                    if let NodeBody::LocalityProvider(node) = body
                        && backfill_state_table_id.is_none()
                    {
                        backfill_state_table_id =
                            node.progress_table.as_ref().map(|table| table.id);
                    }
                });

                let mut backfill_target_relation_id = None;
                visit_stream_node_stream_scan(stream_node, |node| {
                    if backfill_target_relation_id.is_none() {
                        backfill_target_relation_id = Some(node.table_id.as_relation_id());
                    }
                });

                scan = Some(RwBackfillInfo {
                    job_id: fragment_distribution.table_id,
                    fragment_id: fragment_distribution.fragment_id,
                    backfill_state_table_id: backfill_state_table_id
                        .unwrap_or(TableId::placeholder()),
                    backfill_target_relation_id: backfill_target_relation_id?,
                    backfill_type: backfill_type.to_string(),
                    backfill_epoch: 0,
                });
            } else {
                visit_stream_node_stream_scan(stream_node, |node| {
                    scan = Some(RwBackfillInfo {
                        job_id: fragment_distribution.table_id,
                        fragment_id: fragment_distribution.fragment_id,
                        backfill_state_table_id: node
                            .state_table
                            .as_ref()
                            .map(|table| table.id)
                            .unwrap_or(TableId::placeholder()),
                        backfill_target_relation_id: node.table_id.as_relation_id(),
                        backfill_type: backfill_type.to_string(),
                        backfill_epoch: node.snapshot_backfill_epoch() as _,
                    });
                });
            }
        }
    }

    scan
}

#[system_catalog(table, "rw_catalog.rw_backfill_info")]
async fn read_rw_backfill_info(reader: &SysCatalogReaderImpl) -> Result<Vec<RwBackfillInfo>> {
    let distributions = reader
        .meta_client
        .list_creating_fragment_distribution()
        .await?;

    Ok(distributions
        .into_iter()
        .filter_map(|distribution| extract_stream_scan(&distribution))
        .collect())
}
