// Copyright 2024 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.

//! Provides E2E Test runner functionality.

use risingwave_sqlparser::ast::Statement;
use tokio_postgres::Client;

use crate::config::Configuration;
use crate::test_runners::utils::{
    create_base_tables, create_mviews, drop_mview_table, drop_tables, format_drop_mview,
    generate_rng, populate_tables, run_query, set_variable, test_batch_queries,
    test_session_variable, test_sqlsmith, test_stream_queries, update_base_tables,
};
use crate::utils::read_file_contents;
use crate::{mview_sql_gen, parse_sql, sql_gen};

/// e2e test runner for pre-generated queries from sqlsmith
pub async fn run_pre_generated(client: &Client, outdir: &str) {
    let timeout_duration = 12; // allow for some variance.
    let queries_path = format!("{}/queries.sql", outdir);
    let queries = read_file_contents(queries_path).unwrap();
    for statement in parse_sql(&queries) {
        let sql = statement.to_string();
        tracing::info!("[EXECUTING STATEMENT]: {}", sql);
        run_query(timeout_duration, client, &sql).await.unwrap();
    }
    tracing::info!("[EXECUTION SUCCESS]");
}

/// Query Generator
/// If we encounter an expected error, just skip.
/// If we encounter an unexpected error,
/// Sqlsmith should stop execution, but writeout ddl and queries so far.
/// If query takes too long -> cancel it, **mark it as error**.
/// NOTE(noel): It will still fail if DDL creation fails.
pub async fn generate(
    client: &Client,
    testdata: &str,
    count: usize,
    _outdir: &str,
    config: &Configuration,
    seed: Option<u64>,
) {
    let timeout_duration = 5;

    set_variable(client, "RW_IMPLICIT_FLUSH", "TRUE").await;
    set_variable(client, "QUERY_MODE", "DISTRIBUTED").await;
    tracing::info!("Set session variables");

    let mut rng = generate_rng(seed);
    let base_tables = create_base_tables(testdata, client).await.unwrap();

    let rows_per_table = 50;
    let max_rows_inserted = rows_per_table * base_tables.len();
    let inserts = populate_tables(
        client,
        &mut rng,
        base_tables.clone(),
        rows_per_table,
        config,
    )
    .await;
    tracing::info!("Populated base tables");

    let (tables, _mviews) = create_mviews(&mut rng, base_tables.clone(), client, config)
        .await
        .unwrap();

    // Generate an update for some inserts, on the corresponding table.
    update_base_tables(client, &mut rng, &base_tables, &inserts, config).await;

    test_sqlsmith(
        client,
        &mut rng,
        tables.clone(),
        base_tables.clone(),
        max_rows_inserted,
        config,
    )
    .await;
    tracing::info!("Passed sqlsmith tests");

    tracing::info!("Ran updates");

    let mut generated_queries = 0;
    for _ in 0..count {
        test_session_variable(client, &mut rng).await;
        let sql = sql_gen(&mut rng, tables.clone(), config);
        tracing::info!("[EXECUTING TEST_BATCH]: {}", sql);
        let result = run_query(timeout_duration, client, sql.as_str()).await;
        match result {
            Err(_e) => {
                generated_queries += 1;
                tracing::info!("Generated {} batch queries", generated_queries);
                tracing::error!("Unrecoverable error encountered.");
                return;
            }
            Ok(0) => {
                generated_queries += 1;
            }
            _ => {}
        }
    }
    tracing::info!("Generated {} batch queries", generated_queries);

    let mut generated_queries = 0;
    for _ in 0..count {
        test_session_variable(client, &mut rng).await;
        let (sql, table) = mview_sql_gen(&mut rng, tables.clone(), "stream_query", config);
        tracing::info!("[EXECUTING TEST_STREAM]: {}", sql);
        let result = run_query(timeout_duration, client, sql.as_str()).await;
        match result {
            Err(_e) => {
                generated_queries += 1;
                tracing::info!("Generated {} stream queries", generated_queries);
                tracing::error!("Unrecoverable error encountered.");
                return;
            }
            Ok(0) => {
                generated_queries += 1;
            }
            _ => {}
        }
        tracing::info!("[EXECUTING DROP MVIEW]: {}", &format_drop_mview(&table));
        drop_mview_table(&table, client).await;
    }
    tracing::info!("Generated {} stream queries", generated_queries);

    drop_tables(testdata, client).await;
}

/// e2e test runner for sqlsmith
pub async fn run(
    client: &Client,
    testdata: &str,
    count: usize,
    config: &Configuration,
    seed: Option<u64>,
) {
    let mut rng = generate_rng(seed);

    set_variable(client, "RW_IMPLICIT_FLUSH", "TRUE").await;
    set_variable(client, "QUERY_MODE", "DISTRIBUTED").await;
    tracing::info!("Set session variables");

    let base_tables = create_base_tables(testdata, client).await.unwrap();

    let rows_per_table = 50;
    let inserts = populate_tables(
        client,
        &mut rng,
        base_tables.clone(),
        rows_per_table,
        config,
    )
    .await;
    tracing::info!("Populated base tables");

    let (tables, _mviews) = create_mviews(&mut rng, base_tables.clone(), client, config)
        .await
        .unwrap();
    tracing::info!("Created tables");

    // Filter out not-append-only tables.
    let updatable_base_tables: Vec<_> = base_tables
        .iter()
        .filter(|table| !table.is_append_only)
        .cloned()
        .collect();

    // Filter out inserts on not-append-only tables.
    let updatable_inserts: Vec<_> = inserts
        .iter()
        .filter(|stmt| {
            if let Statement::Insert { table_name, .. } = stmt {
                updatable_base_tables
                    .iter()
                    .any(|table| table.name == table_name.base_name())
            } else {
                false
            }
        })
        .cloned()
        .collect();

    // Generate an update for some inserts, on the corresponding table.
    update_base_tables(
        client,
        &mut rng,
        &updatable_base_tables,
        &updatable_inserts,
        config,
    )
    .await;
    tracing::info!("Ran updates");

    let max_rows_inserted = rows_per_table * base_tables.len();
    test_sqlsmith(
        client,
        &mut rng,
        tables.clone(),
        base_tables.clone(),
        max_rows_inserted,
        config,
    )
    .await;
    tracing::info!("Passed sqlsmith tests");

    test_batch_queries(client, &mut rng, tables.clone(), count, config)
        .await
        .unwrap();
    tracing::info!("Passed batch queries");
    test_stream_queries(client, &mut rng, tables.clone(), count, config)
        .await
        .unwrap();
    tracing::info!("Passed stream queries");

    drop_tables(testdata, client).await;
    tracing::info!("[EXECUTION SUCCESS]");
}
