//! High-level tests for SQL generation with pre-aggregations
//!
//! These tests verify that queries correctly match and use pre-aggregations,
//! checking that the generated SQL contains references to pre-aggregation tables.

use crate::cube_bridge::member_expression::MemberExpressionExpressionDef;
use crate::cube_bridge::member_sql::MemberSql;
use crate::cube_bridge::options_member::OptionsMember;
use crate::test_fixtures::cube_bridge::{
    MockMemberExpressionDefinition, MockMemberSql, MockSchema,
};
use crate::test_fixtures::test_utils::TestContext;
use cubenativeutils::CubeError;
use indoc::indoc;
use std::rc::Rc;

#[tokio::test(flavor = "multi_thread")]
async fn test_basic_pre_agg_sql() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregations_test.yaml");
    let test_context = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - visitors.count
        dimensions:
          - visitors.source
        order:
          - id: visitors.source
    "};

    let (_sql, pre_aggrs) = test_context.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1, "Should use one pre-aggregation");
    assert_eq!(pre_aggrs[0].name(), "daily_rollup");

    if let Some(result) = test_context
        .try_execute(query_yaml, "pre_aggregation_tables.sql")
        .await
    {
        insta::assert_snapshot!("basic_pre_agg_sql_cubestore_result", result);
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_full_match_main_rollup() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["main_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
          - orders.total_amount
        dimensions:
          - orders.status
          - orders.city
        order:
          - id: orders.status
          - id: orders.city
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "main_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!("full_match_main_rollup_cubestore_result", result);
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_partial_match_main_rollup() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["main_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.status
        order:
          - id: orders.status
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "main_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!("partial_match_main_rollup_cubestore_result", result);
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_full_match_non_additive_measure() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["main_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.avg_amount
        dimensions:
          - orders.status
          - orders.city
        order:
          - id: orders.status
          - id: orders.city
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "main_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!("full_match_non_additive_measure_cubestore_result", result);
    }

    Ok(())
}

#[test]
fn test_no_match_non_additive_measure_partial() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["main_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.avg_amount
            dimensions:
              - orders.status
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_daily_rollup_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["daily_countries_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.country
        time_dimensions:
          - dimension: orders.created_at
            granularity: day
        order:
          - id: orders.country
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "daily_countries_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!("daily_rollup_full_match_cubestore_result", result);
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_daily_rollup_coarser_granularity() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["daily_countries_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.country
        time_dimensions:
          - dimension: orders.created_at
            granularity: month
        order:
          - id: orders.country
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "daily_countries_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!("daily_rollup_coarser_granularity_cubestore_result", result);
    }

    Ok(())
}

#[test]
fn test_daily_rollup_finer_granularity_no_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["daily_countries_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.count
            dimensions:
              - orders.country
            time_dimensions:
              - dimension: orders.created_at
                granularity: hour
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_daily_rollup_non_additive_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["daily_countries_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.avg_amount
        dimensions:
          - orders.country
        time_dimensions:
          - dimension: orders.created_at
            granularity: day
        order:
          - id: orders.country
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "daily_countries_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "daily_rollup_non_additive_full_match_cubestore_result",
            result
        );
    }

    Ok(())
}

#[test]
fn test_daily_rollup_non_additive_coarser_granularity_no_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["daily_countries_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.avg_amount
            dimensions:
              - orders.country
            time_dimensions:
              - dimension: orders.created_at
                granularity: month
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

// --- multi_level_measure across different pre-aggregations ---

#[tokio::test(flavor = "multi_thread")]
async fn test_multi_level_all_base_measures_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["all_base_measures_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.multi_level_measure
        dimensions:
          - orders.status
          - orders.city
        order:
          - id: orders.status
          - id: orders.city
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "all_base_measures_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "multi_level_all_base_measures_full_match_cubestore_result",
            result
        );
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_multi_level_all_base_measures_partial_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["all_base_measures_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.multi_level_measure
        dimensions:
          - orders.status
        order:
          - id: orders.status
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "all_base_measures_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "multi_level_all_base_measures_partial_match_cubestore_result",
            result
        );
    }

    Ok(())
}

#[test]
fn test_multi_level_calculated_measure_no_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["calculated_measure_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.multi_level_measure
            dimensions:
              - orders.status
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_multi_level_calculated_measure_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["calculated_measure_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.multi_level_measure
        dimensions:
          - orders.status
          - orders.city
        order:
          - id: orders.status
          - id: orders.city
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "calculated_measure_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "multi_level_calculated_measure_full_match_cubestore_result",
            result
        );
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_multi_level_mixed_measure_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["mixed_measure_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.multi_level_measure
        dimensions:
          - orders.status
          - orders.city
        order:
          - id: orders.status
          - id: orders.city
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "mixed_measure_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "multi_level_mixed_measure_full_match_cubestore_result",
            result
        );
    }

    Ok(())
}

#[test]
fn test_multi_level_mixed_measure_partial_no_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["mixed_measure_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.multi_level_measure
            dimensions:
              - orders.status
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_base_and_calculated_measure_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["base_and_calculated_measure_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.amount_per_count
        dimensions:
          - orders.status
          - orders.city
        order:
          - id: orders.status
          - id: orders.city
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "base_and_calculated_measure_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "base_and_calculated_measure_full_match_cubestore_result",
            result
        );
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_base_and_calculated_measure_parital_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["base_and_calculated_measure_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.amount_per_count
        dimensions:
          - orders.status
        order:
          - id: orders.status
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "base_and_calculated_measure_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "base_and_calculated_measure_parital_match_cubestore_result",
            result
        );
    }

    Ok(())
}

// --- Segment matching tests ---

// When a cube's access policy denies the queried members, RBAC
// (CompilerApi.applyRowLevelSecurity) appends a member-expression segment
// `{ expression: () => '1 = 0', cubeName, name: 'rlsAccessDenied' }`. It
// references no members (empty dependencies), so it's a constant filter on top
// of the rollup and must not disqualify pre-aggregation matching.
#[tokio::test(flavor = "multi_thread")]
async fn test_constant_member_expression_segment_keeps_pre_aggregation() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["main_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.status
    "};

    let access_denied_segment = {
        let sql: Rc<dyn MemberSql> = Rc::new(MockMemberSql::new("1 = 0")?);
        let expr = MockMemberExpressionDefinition::builder()
            .expression_name(Some("rlsAccessDenied".to_string()))
            .cube_name(Some("orders".to_string()))
            .expression(MemberExpressionExpressionDef::Sql(sql))
            .build();
        OptionsMember::MemberExpression(Rc::new(expr))
    };

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations_with_segments(
        query_yaml,
        vec![access_denied_segment],
    )?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "main_rollup");
    assert!(
        sql.contains("1 = 0"),
        "expected the constant access-denied segment in SQL, got:\n{sql}"
    );

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_segment_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["segment_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.status
        segments:
          - orders.high_priority
        time_dimensions:
          - dimension: orders.created_at
            granularity: day
        order:
          - id: orders.status
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "segment_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!("segment_full_match_cubestore_result", result);
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_segment_partial_match_unused_segment() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["segment_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.status
        time_dimensions:
          - dimension: orders.created_at
            granularity: day
        order:
          - id: orders.status
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "segment_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "segment_partial_match_unused_segment_cubestore_result",
            result
        );
    }

    Ok(())
}

#[test]
fn test_segment_no_match_missing_in_pre_agg() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["main_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.count
            dimensions:
              - orders.status
              - orders.city
            segments:
              - orders.high_priority
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

// --- Custom granularity pre-aggregation tests ---

#[tokio::test(flavor = "multi_thread")]
async fn test_custom_granularity_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml")
        .only_pre_aggregations(&["custom_half_year_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.status
        time_dimensions:
          - dimension: orders.created_at
            granularity: half_year
        order:
          - id: orders.status
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "custom_half_year_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!("custom_granularity_full_match_cubestore_result", result);
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_standard_pre_agg_coarser_custom_query() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml")
        .only_pre_aggregations(&["daily_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.status
        time_dimensions:
          - dimension: orders.created_at
            granularity: half_year
        order:
          - id: orders.status
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "daily_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "standard_pre_agg_coarser_custom_query_cubestore_result",
            result
        );
    }

    Ok(())
}

#[test]
fn test_custom_pre_agg_finer_query_no_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml")
        .only_pre_aggregations(&["custom_half_year_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.count
            dimensions:
              - orders.status
            time_dimensions:
              - dimension: orders.created_at
                granularity: day
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

#[test]
fn test_custom_pre_agg_finer_standard_query_no_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml")
        .only_pre_aggregations(&["custom_half_year_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.count
            dimensions:
              - orders.status
            time_dimensions:
              - dimension: orders.created_at
                granularity: month
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_custom_granularity_non_additive_full_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml")
        .only_pre_aggregations(&["custom_half_year_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.avg_amount
        dimensions:
          - orders.status
        time_dimensions:
          - dimension: orders.created_at
            granularity: half_year
        order:
          - id: orders.status
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "custom_half_year_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "custom_granularity_non_additive_full_match_cubestore_result",
            result
        );
    }

    Ok(())
}

#[test]
fn test_custom_granularity_non_additive_coarser_no_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml")
        .only_pre_aggregations(&["daily_rollup"]);
    let ctx = TestContext::new(schema)?;

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - orders.avg_amount
            dimensions:
              - orders.status
            time_dimensions:
              - dimension: orders.created_at
                granularity: half_year
        "})?;

    assert!(pre_aggrs.is_empty());

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_custom_granularity_non_strict_self_match() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/custom_granularity_test.yaml")
        .only_pre_aggregations(&["custom_half_year_non_strict"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        time_dimensions:
          - dimension: orders.created_at
            granularity: half_year
        order:
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "custom_half_year_non_strict");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "custom_granularity_non_strict_self_match_cubestore_result",
            result
        );
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_segment_with_coarser_granularity() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["segment_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.status
        segments:
          - orders.high_priority
        time_dimensions:
          - dimension: orders.created_at
            granularity: month
        order:
          - id: orders.status
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "segment_rollup");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!("segment_with_coarser_granularity_cubestore_result", result);
    }

    Ok(())
}

// --- Multi-stage count_distinct sum by quarter with pre-aggregation ---

#[tokio::test(flavor = "multi_thread")]
async fn test_multi_stage_count_distinct_sum_by_quarter_with_pre_aggregation(
) -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/multi_stage_sum_by_quarter_test.yaml");
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - coach.count_distinct__sum_by_quarter
        cubestoreSupportMultistage: true
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "main");

    if let Some(result) = ctx
        .try_execute(query_yaml, "multi_stage_sum_by_quarter_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "multi_stage_count_distinct_sum_by_quarter_with_pre_agg_cubestore_result",
            result
        );
    }

    Ok(())
}

// --- Multi-stage with separate pre-aggregations ---

#[tokio::test(flavor = "multi_thread")]
async fn test_multi_stage_separate_pre_aggregations() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/multi_stage_separate_pre_aggs_test.yaml");
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count_reduce_status
          - orders.revenue_reduce_status
        cubestoreSupportMultistage: true
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 2, "Expected 2 pre-aggregation usages");

    let names: Vec<&str> = pre_aggrs
        .iter()
        .map(|u| u.pre_aggregation.name().as_str())
        .collect();
    assert!(
        names.contains(&"count_rollup"),
        "Expected count_rollup, got {:?}",
        names
    );
    assert!(
        names.contains(&"revenue_rollup"),
        "Expected revenue_rollup, got {:?}",
        names
    );

    if let Some(result) = ctx
        .try_execute(query_yaml, "multi_stage_separate_pre_aggs_tables.sql")
        .await
    {
        insta::assert_snapshot!("multi_stage_separate_pre_aggs_cubestore_result", result);
    }

    Ok(())
}

// --- Multi-stage with separate pre-aggregations and time shift ---

#[tokio::test(flavor = "multi_thread")]
async fn test_multi_stage_separate_pre_aggs_with_time_shift() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/multi_stage_pre_agg_time_shift_test.yaml");
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count_prev_month
          - orders.revenue_reduce_status
        time_dimensions:
          - dimension: orders.created_at
            granularity: month
            dateRange:
              - \"2025-01-01\"
              - \"2025-03-31\"
        cubestoreSupportMultistage: true
        order:
          - id: orders.created_at
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 2, "Expected 2 pre-aggregation usages");

    // Find usages by pre-aggregation name
    let count_usage = pre_aggrs
        .iter()
        .find(|u| u.pre_aggregation.name() == "count_rollup")
        .expect("Expected count_rollup usage");
    let revenue_usage = pre_aggrs
        .iter()
        .find(|u| u.pre_aggregation.name() == "revenue_rollup")
        .expect("Expected revenue_rollup usage");

    // count_prev_month has time_shift prior 1 month, so date range should be shifted back
    assert_eq!(
        count_usage.date_range,
        Some((
            "2024-12-01T00:00:00.000".to_string(),
            "2025-02-28T23:59:59.999".to_string()
        )),
        "count_rollup should have date range shifted 1 month prior"
    );

    // revenue_reduce_status has no time shift, so original date range
    assert_eq!(
        revenue_usage.date_range,
        Some((
            "2025-01-01T00:00:00.000".to_string(),
            "2025-03-31T23:59:59.999".to_string()
        )),
        "revenue_rollup should have original date range"
    );

    if let Some(result) = ctx
        .try_execute(query_yaml, "multi_stage_pre_agg_time_shift_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "multi_stage_separate_pre_aggs_time_shift_cubestore_result",
            result
        );
    }

    Ok(())
}

// --- rollupJoin with calculated measures through view ---

#[test]
fn test_rollup_join_calculated_measures_through_view() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/rollup_join_calculated_measures.yaml");
    let ctx = TestContext::new(schema)?;

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(indoc! {"
            measures:
              - my_view.facts_avg_cost
            time_dimensions:
              - dimension: my_view.facts_day
                granularity: day
        "})?;

    let pre_agg_names: Vec<_> = pre_aggrs
        .iter()
        .map(|pa| format!("{}.{}", pa.cube_name(), pa.name()))
        .collect();
    assert!(
        pre_agg_names
            .iter()
            .any(|n| n == "line_items.combined_rollup_join"),
        "Should use combined_rollup_join, got: {:?}\nSQL:\n{}",
        pre_agg_names,
        sql
    );
    assert!(
        sql.contains("campaigns_rollup"),
        "SQL should reference campaigns_rollup, got:\n{}",
        sql
    );
    assert!(
        sql.contains("facts_rollup"),
        "SQL should reference facts_rollup, got:\n{}",
        sql
    );
    assert!(
        sql.contains("li_rollup"),
        "SQL should reference li_rollup, got:\n{}",
        sql
    );

    Ok(())
}

// A rolling-window count_distinct_approx whose pre-aggregation stores the
// rolling measure itself: the leaf reads the rollup's HLL state and must
// keep it MERGED (state only), so the outer rolling-window stage can merge
// across the window and finalize to a cardinality. This pins the state
// branch — the read must NOT collapse the state to a cardinality too early.
#[test]
fn test_count_distinct_approx_rolling_pre_agg_keeps_state() -> Result<(), CubeError> {
    let ctx = TestContext::new(MockSchema::from_yaml_file(
        "common/integration_rolling_window.yaml",
    ))?;

    let query = indoc! {r#"
        measures:
          - orders.rolling_unique_customers_approx_7d
        dimensions:
          - orders.category
        time_dimensions:
          - dimension: orders.created_at
            granularity: day
            dateRange:
              - "2024-01-10"
              - "2024-01-25"
        cubestoreSupportMultistage: true
    "#};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "approx_rolling");

    // Leaf reads the rollup column as a bare merged state (hll_merge).
    assert!(
        sql.contains("merge(\"orders__rolling_unique_customers_approx_7d\")"),
        "Rolling leaf should keep the merged HLL state, got:\n{}",
        sql
    );
    // The rolling-window stage finalizes the merged states to a cardinality.
    assert!(
        sql.contains("cardinality(merge("),
        "Rolling window should finalize merged states to a cardinality, got:\n{}",
        sql
    );

    Ok(())
}

// --- HLL count_distinct_approx through a pre-aggregation ---
//
// A count_distinct_approx measure materialized in a rollup keeps an HLL
// state per group. The build (load) SQL must serialize that state
// (hll_init), and a query reading the rollup must merge the states and
// take their cardinality (hll_cardinality_merge) rather than recompute
// an approximate distinct over the state column (count_distinct_approx).
//
// The mock driver renders these as distinct, CubeStore-style forms:
//   hll_init               -> hll_add_agg(hll_hash_any(x))
//   hll_merge              -> merge(x)                       (state only)
//   hll_cardinality_merge  -> cardinality(merge(x))
//   count_distinct_approx  -> round(hll_cardinality(hll_add_agg(hll_hash_any(x))))

#[test]
fn test_count_distinct_approx_pre_agg_read_merges_state() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["approx_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.approx_unique_count
        dimensions:
          - orders.status
          - orders.city
    "};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "approx_rollup");

    // A top-level read merges the stored states and finalizes them to a
    // cardinality. It must NOT re-hash the state column as a fresh
    // approximate distinct would (hll_add_agg).
    assert!(
        sql.contains("cardinality(merge("),
        "Read query should merge HLL states and take cardinality, got:\n{}",
        sql
    );
    assert!(
        !sql.contains("hll_add_agg"),
        "Read query should not re-init HLL from the state column, got:\n{}",
        sql
    );

    Ok(())
}

#[test]
fn test_count_distinct_approx_pre_agg_build_emits_state() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["approx_rollup"]);
    let ctx = TestContext::new(schema)?;

    // pre_aggregation_query: true renders the rollup build (load) SQL.
    let query_yaml = indoc! {"
        measures:
          - orders.approx_unique_count
        dimensions:
          - orders.status
          - orders.city
        pre_aggregation_query: true
    "};

    let sql = ctx.build_sql(query_yaml)?;

    // The build must serialize the HLL state (hll_init) without merging
    // or taking its cardinality — that happens only on read.
    assert!(
        sql.contains("hll_add_agg(hll_hash_any("),
        "Build SQL should emit HLL state, got:\n{}",
        sql
    );
    assert!(
        !sql.contains("hll_cardinality"),
        "Build SQL should not compute cardinality, got:\n{}",
        sql
    );
    assert!(
        !sql.contains("merge("),
        "Build SQL should not merge states, got:\n{}",
        sql
    );

    Ok(())
}

// A multi-stage measure that sums a count_distinct_approx must read the
// rollup as a finalized cardinality, exactly as it would without the
// pre-aggregation — the outer `sum` aggregates counts, not raw HLL states.
// This pins that the pre-agg read does not leak a bare merged state here.
#[test]
fn test_count_distinct_approx_multistage_pre_agg_reads_cardinality() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/multi_stage_sum_by_quarter_test.yaml");
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - coach.approx__sum_by_quarter
        cubestoreSupportMultistage: true
    "};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "main_approx");

    // Leaf finalizes the merged state to a cardinality so the outer sum
    // aggregates numbers (same shape as the non-pre-agg rendering).
    assert!(
        sql.contains("cardinality(merge("),
        "Multi-stage leaf should finalize HLL to cardinality, got:\n{}",
        sql
    );

    Ok(())
}

// A cube `foo` whose `originalSql` pre-aggregation (`main`) materializes its base
// SQL, plus a `second` rollup. The mock renders the originalSql pre-agg table as
// `foo__main` and the raw cube SQL references `foo_table`.
fn use_original_sql_pre_aggregations_in_pre_aggregation_schema() -> Result<MockSchema, CubeError> {
    MockSchema::from_yaml(indoc! {"
        cubes:
          - name: foo
            sql: SELECT * FROM foo_table
            dimensions:
              - name: time
                type: time
                sql: timestamp
            measures:
              - name: total
                type: sum
                sql: amount
            pre_aggregations:
              - name: main
                type: originalSql
              - name: second
                type: rollup
                measures:
                  - total
                time_dimension: time
                granularity: day
    "})
}

// Building a rollup with `useOriginalSqlPreAggregations` must source it from the cube's
// `originalSql` pre-aggregation table instead of the raw cube SQL
#[test]
fn test_rollup_build_with_use_original_sql_pre_aggregations_in_pre_aggregation_reads_original_sql_pre_agg(
) -> Result<(), CubeError> {
    let ctx = TestContext::new(use_original_sql_pre_aggregations_in_pre_aggregation_schema()?)?;

    let query_yaml = indoc! {"
        measures:
          - foo.total
        time_dimensions:
          - dimension: foo.time
            granularity: day
        pre_aggregation_query: true
        use_original_sql_pre_aggregations_in_pre_aggregation: true
    "};

    let sql = ctx.build_sql(query_yaml)?;

    assert!(
        sql.contains("foo__main"),
        "Build SQL should source from the originalSql pre-agg table, got:\n{}",
        sql
    );
    assert!(
        !sql.contains("foo_table"),
        "Build SQL should not read the raw cube table when useOriginalSqlPreAggregations is set, got:\n{}",
        sql
    );
    Ok(())
}

// Without the flag, a rollup build reads the raw cube SQL — the originalSql
// pre-agg table must not be substituted in.
#[test]
fn test_rollup_build_without_use_original_sql_pre_aggregations_in_pre_aggregation_reads_raw_table(
) -> Result<(), CubeError> {
    let ctx = TestContext::new(use_original_sql_pre_aggregations_in_pre_aggregation_schema()?)?;

    let query_yaml = indoc! {"
        measures:
          - foo.total
        time_dimensions:
          - dimension: foo.time
            granularity: day
        pre_aggregation_query: true
    "};

    let sql = ctx.build_sql(query_yaml)?;

    assert!(
        sql.contains("foo_table"),
        "Build SQL should read the raw cube table without useOriginalSqlPreAggregations, got:\n{}",
        sql
    );
    assert!(
        !sql.contains("foo__main"),
        "Build SQL should not source from the originalSql pre-agg table without the flag, got:\n{}",
        sql
    );
    Ok(())
}

// A measure referenced only in ORDER BY (not in the selected measures) is dropped
// from the ORDER BY when reading a pre-aggregation. CubeStore cannot ORDER BY an
// aggregate of a rollup column that isn't projected, and the legacy planner
// likewise ignores such keys, so the remaining keys (here the time dimension and
// the selected dimension) drive the order. Mirrors the driver-test "partitioned
// pre-agg" queries that order by `created_at asc, <unselected measure> desc, <dim> asc`.
#[tokio::test(flavor = "multi_thread")]
async fn test_order_by_only_measure_dropped_from_pre_agg() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["daily_countries_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.count
        dimensions:
          - orders.country
        time_dimensions:
          - dimension: orders.created_at
            granularity: day
        order:
          - id: orders.created_at
            desc: false
          - id: orders.total_amount
            desc: true
          - id: orders.country
            desc: false
    "};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "daily_countries_rollup");
    // total_amount is neither selected nor projected by the rollup read, so it must
    // not appear in ORDER BY — neither as the base column nor the rollup column.
    assert!(
        !sql.contains("\"orders\".amount"),
        "ORDER BY must not reference the base-table column, got:\n{sql}"
    );
    assert!(
        !sql.contains("orders__total_amount"),
        "order-by-only measure must be dropped, not reference the rollup column, got:\n{sql}"
    );

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "order_by_only_measure_dropped_from_pre_agg_cubestore_result",
            result
        );
    }

    Ok(())
}

// An ungrouped query reading from a rollup must reference the stored measure
// column as-is (no `sum()`), not re-render the measure's base-table SQL.
#[tokio::test(flavor = "multi_thread")]
async fn test_ungrouped_pre_agg_measure_reads_rollup_column() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
        .only_pre_aggregations(&["primary_key_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        measures:
          - orders.total_amount
        dimensions:
          - orders.id
          - orders.status
          - orders.city
        ungrouped: true
        order:
          - id: orders.id
    "};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "primary_key_rollup");
    assert!(
        sql.contains("\"orders__total_amount\" \"orders__total_amount\""),
        "ungrouped measure should project the rollup column, got:\n{sql}"
    );
    assert!(
        !sql.contains("\"orders\".amount"),
        "ungrouped measure must not reference the base-table column, got:\n{sql}"
    );

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_matching_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "ungrouped_pre_agg_measure_reads_rollup_column_cubestore_result",
            result
        );
    }

    Ok(())
}

// `lambda_union` UNIONs `visitor_checkins.for_lambda` with
// `visitor_checkins2.for_lambda`. The lambda exposes the first member
// rollup's symbols, so the dimension's unified output alias is
// `visitor_checkins__visitor_id`, but the second branch stores it as
// `visitor_checkins2__visitor_id`.
#[tokio::test(flavor = "multi_thread")]
async fn test_rollup_lambda_cross_cube_union_aliases() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregations_test.yaml");
    let ctx = TestContext::new(schema)?;

    // Force the lambda; otherwise the matcher would pick the plain `for_lambda`
    // rollup that the lambda is built from (it is declared first). The total
    // `order:` pins row order so the CubeStore result snapshot is stable.
    let query_yaml = indoc! {"
        measures:
          - visitor_checkins.count
        dimensions:
          - visitor_checkins.visitor_id
        time_dimensions:
          - dimension: visitor_checkins.created_at
            granularity: day
        order:
          - id: visitor_checkins.visitor_id
          - id: visitor_checkins.created_at
        pre_aggregation_id: visitor_checkins.lambda_union
    "};

    let (_sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(pre_aggrs.len(), 1);
    assert_eq!(pre_aggrs[0].name(), "lambda_union");

    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_tables.sql")
        .await
    {
        insta::assert_snapshot!("rollup_lambda_cross_cube_union_cubestore_result", result);
    }

    Ok(())
}

// `joined_rollup` is a rollup on the fact cube `visitor_checkins` that includes
// a cross-cube dimension from the many_to_one joined `visitors` cube, plus a
// `time_dimension`. An ungrouped query with no measures and no time dimension
// requests raw rows, so it may only use a pre-aggregation whose dimensions
// cover the primary keys of every cube in the query. Here they don't, so
// reading flat rollup columns would silently drop the join between the two
// cubes: the query must fall back to base SQL.
#[tokio::test(flavor = "multi_thread")]
async fn test_ungrouped_no_measures_cross_cube_query_should_not_match_rollup(
) -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregations_test.yaml")
        .only_pre_aggregations(&["joined_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        dimensions:
          - visitor_checkins.visitor_id
          - visitors.source
        filters:
          - dimension: visitors.source
            operator: equals
            values:
              - google
        ungrouped: true
    "};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert!(
        pre_aggrs.is_empty(),
        "expected no pre-aggregation, got `{}`. Generated SQL:\n{sql}",
        pre_aggrs
            .iter()
            .map(|p| p.name().clone())
            .collect::<Vec<_>>()
            .join(", ")
    );

    assert!(
        sql.to_lowercase().contains("join"),
        "expected base SQL with a JOIN between visitor_checkins and visitors, got:\n{sql}"
    );

    Ok(())
}

// Same rejection as above, reached through a view. The view renames every
// member, so the ungrouped gate must compare the cubes and primary keys the
// view members resolve to rather than the names the query spells them with.
#[tokio::test(flavor = "multi_thread")]
async fn test_ungrouped_cross_cube_view_query_should_not_match_rollup() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregations_test.yaml")
        .only_pre_aggregations(&["joined_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        dimensions:
          - visitors_view.visitor_checkins_visitor_id
          - visitors_view.visitors_source
        filters:
          - dimension: visitors_view.visitors_source
            operator: equals
            values:
              - google
        ungrouped: true
    "};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert!(
        pre_aggrs.is_empty(),
        "expected no pre-aggregation, got `{}`. Generated SQL:\n{sql}",
        pre_aggrs
            .iter()
            .map(|p| p.name().clone())
            .collect::<Vec<_>>()
            .join(", ")
    );

    assert!(
        sql.to_lowercase().contains("join"),
        "expected base SQL with a JOIN between visitor_checkins and visitors, got:\n{sql}"
    );

    Ok(())
}

// The mirror case: `for_join` covers the primary key of the only cube the query
// resolves to, so an ungrouped view query must still match it. Guards the gate
// against reading the view itself as a separate cube, which would make the cube
// sets differ and reject every ungrouped query issued through a view.
#[tokio::test(flavor = "multi_thread")]
async fn test_ungrouped_view_query_matches_rollup_covering_primary_key() -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregations_test.yaml")
        .only_pre_aggregations(&["for_join"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        dimensions:
          - visitors_view.visitors_id
          - visitors_view.visitors_source
        ungrouped: true
    "};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(
        pre_aggrs.len(),
        1,
        "expected `visitors.for_join` to be used, got SQL:\n{sql}"
    );
    assert_eq!(pre_aggrs[0].name(), "for_join");
    assert_eq!(pre_aggrs[0].cube_name(), "visitors");

    // The rows are the claim: one per raw visitor, not per (id, source) group.
    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_tables.sql")
        .await
    {
        insta::assert_snapshot!("ungrouped_view_query_for_join_cubestore_result", result);
    }

    Ok(())
}

// `joined_keys_rollup` has the same cross-cube shape as `joined_rollup` but is
// grouped by the primary keys of both cubes, so it does hold one row per raw
// joined row and an ungrouped view query spanning both cubes may use it. The
// contrast with the rejected `joined_rollup` is what the gate is discriminating.
#[tokio::test(flavor = "multi_thread")]
async fn test_ungrouped_cross_cube_view_query_matches_rollup_covering_both_primary_keys(
) -> Result<(), CubeError> {
    let schema = MockSchema::from_yaml_file("common/pre_aggregations_test.yaml")
        .only_pre_aggregations(&["joined_keys_rollup"]);
    let ctx = TestContext::new(schema)?;

    let query_yaml = indoc! {"
        dimensions:
          - visitors_view.visitors_source
          - visitors_view.visitor_checkins_visitor_id
        ungrouped: true
    "};

    let (sql, pre_aggrs) = ctx.build_sql_with_used_pre_aggregations(query_yaml)?;

    assert_eq!(
        pre_aggrs.len(),
        1,
        "expected `joined_keys_rollup` to be used, got SQL:\n{sql}"
    );
    assert_eq!(pre_aggrs[0].name(), "joined_keys_rollup");
    // Tokenized so the check is immune to line breaks and to `join` appearing
    // inside the rollup's own table name.
    assert!(
        !sql.to_lowercase().split_whitespace().any(|t| t == "join"),
        "reading the rollup should not re-join the cubes, got:\n{sql}"
    );

    // The rows are the claim the gate rests on: one per raw joined row.
    if let Some(result) = ctx
        .try_execute(query_yaml, "pre_aggregation_tables.sql")
        .await
    {
        insta::assert_snapshot!(
            "ungrouped_cross_cube_view_joined_keys_rollup_cubestore_result",
            result
        );
    }

    Ok(())
}

// --- References written as a cube-name interpolation ---
//
// `interpolated_refs_rollup` names the same members as `segment_rollup`, but
// through `` `${CUBE}.count` ``-style references. The two must be picked for the
// same queries and read back the same rows. The rows themselves are pinned by
// the snapshot, which is only compared when Postgres execution is enabled.

async fn interpolated_and_symbol_refs_agree(
    query_yaml: &str,
    snapshot_name: &str,
) -> Result<(), CubeError> {
    let interpolated_ctx = TestContext::new(
        MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
            .only_pre_aggregations(&["interpolated_refs_rollup"]),
    )?;
    let symbol_ctx = TestContext::new(
        MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml")
            .only_pre_aggregations(&["segment_rollup"]),
    )?;

    let (_sql, interpolated_pre_aggrs) =
        interpolated_ctx.build_sql_with_used_pre_aggregations(query_yaml)?;
    assert_eq!(interpolated_pre_aggrs.len(), 1);
    assert_eq!(interpolated_pre_aggrs[0].name(), "interpolated_refs_rollup");

    let (_sql, symbol_pre_aggrs) = symbol_ctx.build_sql_with_used_pre_aggregations(query_yaml)?;
    assert_eq!(symbol_pre_aggrs.len(), 1);
    assert_eq!(symbol_pre_aggrs[0].name(), "segment_rollup");

    let interpolated_result = interpolated_ctx
        .try_execute_pg(query_yaml, "pre_aggregation_matching_tables.sql")
        .await;
    let symbol_result = symbol_ctx
        .try_execute_pg(query_yaml, "pre_aggregation_matching_tables.sql")
        .await;

    assert_eq!(interpolated_result, symbol_result);

    // Without Postgres execution there are no rows to compare, so make it
    // explicit that the row assertions below do run when it is enabled.
    #[cfg(feature = "integration-postgres")]
    assert!(
        interpolated_result.is_some(),
        "Postgres execution is enabled but returned no result"
    );

    if let Some(result) = interpolated_result {
        insta::assert_snapshot!(snapshot_name, result);
    }

    Ok(())
}

#[tokio::test(flavor = "multi_thread")]
async fn test_interpolated_refs_full_match() -> Result<(), CubeError> {
    interpolated_and_symbol_refs_agree(
        indoc! {"
            measures:
              - orders.count
              - orders.total_amount
            dimensions:
              - orders.status
            segments:
              - orders.high_priority
            time_dimensions:
              - dimension: orders.created_at
                granularity: day
            order:
              - id: orders.status
              - id: orders.created_at
        "},
        "interpolated_refs_full_match_result",
    )
    .await
}

#[tokio::test(flavor = "multi_thread")]
async fn test_interpolated_refs_with_coarser_granularity() -> Result<(), CubeError> {
    interpolated_and_symbol_refs_agree(
        indoc! {"
            measures:
              - orders.count
            dimensions:
              - orders.status
            segments:
              - orders.high_priority
            time_dimensions:
              - dimension: orders.created_at
                granularity: month
            order:
              - id: orders.status
              - id: orders.created_at
        "},
        "interpolated_refs_with_coarser_granularity_result",
    )
    .await
}
