// Copyright 2023 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::time::Duration;

use anyhow::Result;
use risingwave_simulation::cluster::Configuration;
use risingwave_simulation::nexmark::{NexmarkCluster, THROUGHPUT};
use tokio::time::sleep;

#[tokio::test]
async fn nexmark_source() -> Result<()> {
    nexmark_source_inner(false).await
}

#[tokio::test]
async fn nexmark_source_with_watermark() -> Result<()> {
    nexmark_source_inner(true).await
}

/// Check that everything works well after scaling of source-related executor.
async fn nexmark_source_inner(watermark: bool) -> Result<()> {
    let expected_events = 20 * THROUGHPUT;
    let expected_events_range = if watermark {
        // If there's watermark, we'll possibly get fewer events due to alignment to the global
        // max watermark during scaling. The lower bound is a heuristic to reduce flakes.
        (0.95 * expected_events as f64) as usize..=expected_events
    } else {
        // If there's no watermark, we'll get exactly the expected number of events.
        expected_events..=expected_events
    };

    let configuration = Configuration::for_scale();
    let total_cores = configuration.total_streaming_cores();
    let mut cluster =
        NexmarkCluster::new(configuration, 6, Some(expected_events), watermark).await?;

    // Materialize all sources so that we can also check whether the row id generator is working
    // correctly after scaling.
    // https://github.com/risingwavelabs/risingwave/issues/7103
    let tables = ["person", "auction", "bid"];

    for table in &tables {
        cluster
            .run(&format!(
                "create materialized view materialized_{table} as select * from {table};"
            ))
            .await?;
    }

    macro_rules! reschedule {
        () => {
            for table in tables {
                let parallelism = {
                    use rand::{Rng, rng as thread_rng};
                    let rng = &mut thread_rng();
                    let parallelism = rng.random_range(1..=total_cores);
                    parallelism
                };

                cluster
                    .run(format!(
                        "alter materialized view materialized_{} set parallelism = {}",
                        table, parallelism
                    ))
                    .await?;
            }
        };
    }

    sleep(Duration::from_secs(5)).await;
    reschedule!();

    sleep(Duration::from_secs(5)).await;
    reschedule!();

    sleep(Duration::from_secs(30)).await;

    // Check the total number of events.
    let result = cluster
        .run(
            r#"
with count_p as (select count(*) count_p from materialized_person),
     count_a as (select count(*) count_a from materialized_auction),
     count_b as (select count(*) count_b from materialized_bid)
select count_p + count_a + count_b from count_p, count_a, count_b;"#,
        )
        .await?;

    let actual_events: usize = result.trim().parse()?;
    assert!(
        expected_events_range.contains(&actual_events),
        "expected event num in {:?}, got {}",
        expected_events_range,
        actual_events
    );

    Ok(())
}
