use super::CompiledPreAggregation;
use crate::planner::filter::{BaseFilter, FilterGroupOperator, FilterItem};
use crate::planner::query_tools::QueryTools;
use crate::planner::DimensionSymbol;
use crate::planner::GranularityHelper;
use crate::planner::MemberSymbol;
use crate::planner::TimeDimensionSymbol;
use cubenativeutils::CubeError;
use std::collections::HashMap;
use std::rc::Rc;

#[derive(Clone, Debug, PartialEq)]
pub enum MatchState {
    Partial,
    Full,
    NotMatched,
}

impl MatchState {
    pub fn combine(&self, other: &MatchState) -> MatchState {
        if matches!(self, MatchState::NotMatched) || matches!(other, MatchState::NotMatched) {
            return MatchState::NotMatched;
        }
        if matches!(self, MatchState::Partial) || matches!(other, MatchState::Partial) {
            return MatchState::Partial;
        }
        MatchState::Full
    }
}

pub struct DimensionMatcher<'a> {
    query_tools: Rc<QueryTools>,
    pre_aggregation: &'a CompiledPreAggregation,
    pre_aggregation_dimensions: HashMap<String, bool>,
    pre_aggregation_time_dimensions: HashMap<String, Vec<(Rc<TimeDimensionSymbol>, bool)>>,
    pre_aggregation_segments: HashMap<String, bool>,
    result: MatchState,
}

impl<'a> DimensionMatcher<'a> {
    pub fn new(query_tools: Rc<QueryTools>, pre_aggregation: &'a CompiledPreAggregation) -> Self {
        let pre_aggregation_dimensions = pre_aggregation
            .dimensions
            .iter()
            .map(|d| (d.full_name(), false))
            .collect();
        let mut pre_aggregation_time_dimensions =
            HashMap::<String, Vec<(Rc<TimeDimensionSymbol>, bool)>>::new();
        for dim in pre_aggregation.time_dimensions.iter() {
            if let Ok(td) = dim.as_time_dimension() {
                let key = td.base_symbol().full_name();
                pre_aggregation_time_dimensions
                    .entry(key)
                    .or_default()
                    .push((td, false));
            }
        }
        let pre_aggregation_segments = pre_aggregation
            .segments
            .iter()
            .map(|s| (s.full_name(), false))
            .collect();
        Self {
            query_tools,
            pre_aggregation,
            pre_aggregation_dimensions,
            pre_aggregation_time_dimensions,
            pre_aggregation_segments,
            result: MatchState::Full,
        }
    }

    pub fn try_match(
        &mut self,
        dimensions: &Vec<Rc<MemberSymbol>>,
        time_dimensions: &Vec<Rc<MemberSymbol>>,
        filters: &Vec<FilterItem>,
        time_dimension_filters: &Vec<FilterItem>,
        segments: &Vec<FilterItem>,
    ) -> Result<(), CubeError> {
        for dimension in dimensions.iter() {
            let dimension_match = self.try_match_symbol(dimension, true)?;
            self.result = self.result.combine(&dimension_match);
            if self.result == MatchState::NotMatched {
                return Ok(());
            }
        }
        for time_dimension in time_dimensions.iter() {
            let time_dimension_match = self.try_match_symbol(time_dimension, true)?;
            self.result = self.result.combine(&time_dimension_match);
            if self.result == MatchState::NotMatched {
                return Ok(());
            }
        }

        for filter in filters.iter() {
            let filter_match = self.try_match_filter_item(filter, true)?;
            self.result = self.result.combine(&filter_match);
            if self.result == MatchState::NotMatched {
                return Ok(());
            }
        }

        for filter in time_dimension_filters.iter() {
            let filter_match = self.try_match_filter_item(filter, true)?;
            self.result = self.result.combine(&filter_match);
            if self.result == MatchState::NotMatched {
                return Ok(());
            }
        }

        for segment in segments.iter() {
            let segment_match = self.try_match_filter_item(segment, true)?;
            self.result = self.result.combine(&segment_match);
            if self.result == MatchState::NotMatched {
                return Ok(());
            }
        }
        Ok(())
    }

    pub fn result(mut self) -> MatchState {
        let dimension_coverage_result = if self.pre_aggregation_dimensions.values().all(|v| *v) {
            MatchState::Full
        } else {
            MatchState::Partial
        };
        self.result = self.result.combine(&dimension_coverage_result);
        let time_dimension_coverage_result = if self
            .pre_aggregation_time_dimensions
            .values()
            .all(|entries| entries.iter().all(|(_, matched)| *matched))
        {
            MatchState::Full
        } else {
            MatchState::Partial
        };
        self.result = self.result.combine(&time_dimension_coverage_result);
        let segment_coverage_result = if self.pre_aggregation_segments.values().all(|v| *v) {
            MatchState::Full
        } else {
            MatchState::Partial
        };
        self.result = self.result.combine(&segment_coverage_result);
        self.result
    }

    fn try_match_symbol(
        &mut self,
        symbol: &Rc<MemberSymbol>,
        add_to_matched_dimension: bool,
    ) -> Result<MatchState, CubeError> {
        match symbol.as_ref() {
            MemberSymbol::Dimension(dimension) => {
                self.try_match_dimension(dimension, add_to_matched_dimension)
            }
            MemberSymbol::TimeDimension(time_dimension) => {
                self.try_match_time_dimension(time_dimension, add_to_matched_dimension)
            }
            MemberSymbol::MemberExpression(me) => {
                if let Some(found) = self.pre_aggregation_segments.get_mut(&me.full_name()) {
                    if add_to_matched_dimension {
                        *found = true;
                    }
                    Ok(MatchState::Full)
                } else {
                    let dependencies = me.get_dependencies();
                    if dependencies.is_empty() {
                        Ok(MatchState::NotMatched)
                    } else {
                        let mut result = MatchState::Partial;
                        for dep in dependencies {
                            let dep_match =
                                self.try_match_symbol(&dep, add_to_matched_dimension)?;
                            if dep_match == MatchState::NotMatched {
                                return Ok(MatchState::NotMatched);
                            }
                            result = result.combine(&dep_match);
                        }
                        Ok(result)
                    }
                }
            }
            _ => Ok(MatchState::NotMatched),
        }
    }

    fn try_match_dimension(
        &mut self,
        dimension: &DimensionSymbol,
        add_to_matched_dimension: bool,
    ) -> Result<MatchState, CubeError> {
        if let Some(found) = self
            .pre_aggregation_dimensions
            .get_mut(&dimension.full_name())
        {
            if add_to_matched_dimension {
                *found = true;
            }
            Ok(MatchState::Full)
        } else if dimension.is_calc_group() {
            // Calc-group dimensions are virtual enumerations: they carry no
            // stored data, so any rollup can serve them. At query time the
            // value is either pinned by a filter (rendered as a literal) or
            // re-cross-joined over the rollup scan, exactly like over a raw
            // source.
            Ok(MatchState::Full)
        } else if dimension.owned_by_cube() {
            Ok(MatchState::NotMatched)
        } else if dimension.is_multi_stage() {
            // Multi-stage dimensions carry aggregate semantics that simple
            // dependency recursion would discard; match only if pre-agg
            // exposes the multi-stage dim itself.
            Ok(MatchState::NotMatched)
        } else {
            let dependencies = dimension.get_dependencies();
            if dependencies.is_empty() {
                Ok(MatchState::NotMatched)
            } else {
                let mut result = if dimension.is_reference() {
                    MatchState::Full
                } else {
                    MatchState::Partial
                };
                for dep in dimension.get_dependencies() {
                    let dep_match = self.try_match_symbol(&dep, add_to_matched_dimension)?;
                    if dep_match == MatchState::NotMatched {
                        return Ok(MatchState::NotMatched);
                    }
                    result = result.combine(&dep_match);
                }
                Ok(result)
            }
        }
    }

    fn try_match_time_dimension(
        &mut self,
        time_dimension: &TimeDimensionSymbol,
        add_to_matched_dimension: bool,
    ) -> Result<MatchState, CubeError> {
        let granularity = if self.pre_aggregation.allow_non_strict_date_range_match {
            time_dimension.granularity().clone()
        } else {
            time_dimension.rollup_granularity(self.query_tools.clone())?
        };
        let base_symbol_name = time_dimension.base_symbol().full_name();

        if let Some(entries) = self
            .pre_aggregation_time_dimensions
            .get_mut(&base_symbol_name)
        {
            // First, look for exact granularity match
            let exact_match = entries
                .iter_mut()
                .find(|(td, _)| granularity.is_none() || td.granularity() == &granularity);
            if let Some((_, matched)) = exact_match {
                if add_to_matched_dimension {
                    *matched = true;
                }
                return Ok(MatchState::Full);
            }

            // No exact match — find the finest pre-agg granularity that covers the query
            let mut best_match = MatchState::NotMatched;
            for (pre_agg_td, matched) in entries.iter_mut() {
                let pre_aggr_granularity = pre_agg_td.granularity();
                if pre_aggr_granularity.is_none() {
                    continue;
                }
                let min_granularity = GranularityHelper::min_granularity_for_time_dimensions(
                    (&granularity, time_dimension),
                    (pre_aggr_granularity, pre_agg_td),
                )?;
                if min_granularity == *pre_aggr_granularity {
                    if add_to_matched_dimension {
                        *matched = true;
                    }
                    best_match = MatchState::Partial;
                    break;
                }
            }
            Ok(best_match)
        } else {
            if time_dimension.owned_by_cube() {
                Ok(MatchState::NotMatched)
            } else if time_dimension.is_multi_stage() {
                Ok(MatchState::NotMatched)
            } else {
                let mut result = if time_dimension.is_reference() {
                    MatchState::Full
                } else {
                    MatchState::Partial
                };
                for dep in time_dimension.get_dependencies_as_time_dimensions() {
                    let dep_match = self.try_match_symbol(&dep, add_to_matched_dimension)?;
                    if dep_match == MatchState::NotMatched {
                        return Ok(MatchState::NotMatched);
                    }
                    result = result.combine(&dep_match);
                }
                Ok(result)
            }
        }
    }

    fn try_match_filter_item(
        &mut self,
        filter_item: &FilterItem,
        add_to_matched_dimension: bool,
    ) -> Result<MatchState, CubeError> {
        match filter_item {
            FilterItem::Item(filter) => self.try_match_filter(filter, add_to_matched_dimension),
            FilterItem::Group(group) => {
                let add_to_matched_dimension =
                    add_to_matched_dimension && group.operator == FilterGroupOperator::And;
                let mut result = MatchState::Full;
                for item in group.items.iter() {
                    result = result
                        .combine(&self.try_match_filter_item(item, add_to_matched_dimension)?);
                }
                Ok(result)
            }
            FilterItem::Segment(segment) => {
                let symbol = segment.member_evaluator();
                // An ad-hoc member-expression segment that references no members
                // (a constant like the `1 = 0` rlsAccessDenied segment RBAC injects
                // on access denial) is a filter pushable on top of any rollup, so it
                // doesn't need pre-aggregation coverage. Named segments still must be
                // covered and fall through to the regular matcher.
                if segment.is_member_expression() && symbol.get_dependencies().is_empty() {
                    return Ok(MatchState::Full);
                }
                self.try_match_symbol(&symbol, add_to_matched_dimension)
            }
        }
    }

    fn try_match_filter(
        &mut self,
        filter: &Rc<BaseFilter>,
        add_to_matched_dimension: bool,
    ) -> Result<MatchState, CubeError> {
        let symbol = if let Some(time_dimension) = filter.time_dimension_symbol() {
            time_dimension
        } else {
            filter.member_evaluator().clone()
        };
        let add_to_matched_dimension = add_to_matched_dimension && filter.is_single_value_equal();
        let res = self.try_match_symbol(&symbol, add_to_matched_dimension)?;
        Ok(res)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cube_bridge::member_expression::MemberExpressionExpressionDef;
    use crate::cube_bridge::member_sql::MemberSql;
    use crate::cube_bridge::options_member::OptionsMember;
    use crate::logical_plan::optimizers::pre_aggregation::{
        PreAggregationFullName, PreAggregationsCompiler,
    };
    use crate::test_fixtures::cube_bridge::{
        MockMemberExpressionDefinition, MockMemberSql, MockSchema,
    };
    use crate::test_fixtures::test_utils::TestContext;
    use indoc::indoc;
    use std::rc::Rc;

    fn create_test_context() -> TestContext {
        let schema = MockSchema::from_yaml_file("common/pre_aggregation_matching_test.yaml");
        TestContext::new(schema).unwrap()
    }

    /// Builds an ad-hoc member-expression segment (no registered `segments:`
    /// path) from constant SQL — mirrors the `rlsAccessDenied` `1 = 0` segment.
    fn const_expr_segment(name: &str, cube: &str, sql: &str) -> OptionsMember {
        let member_sql: Rc<dyn MemberSql> = Rc::new(MockMemberSql::new(sql).unwrap());
        let expr = MockMemberExpressionDefinition::builder()
            .expression_name(Some(name.to_string()))
            .cube_name(Some(cube.to_string()))
            .expression(MemberExpressionExpressionDef::Sql(member_sql))
            .build();
        OptionsMember::MemberExpression(Rc::new(expr))
    }

    fn match_pre_agg(ctx: &TestContext, pre_agg_name: &str, query_yaml: &str) -> MatchState {
        match_pre_agg_with_segments(ctx, pre_agg_name, query_yaml, vec![])
    }

    fn match_pre_agg_with_segments(
        ctx: &TestContext,
        pre_agg_name: &str,
        query_yaml: &str,
        extra_segments: Vec<OptionsMember>,
    ) -> MatchState {
        let cube_names = vec!["orders".to_string()];
        let mut compiler =
            PreAggregationsCompiler::try_new(ctx.query_tools().clone(), &cube_names).unwrap();
        let name = PreAggregationFullName::new("orders".to_string(), pre_agg_name.to_string());
        let pre_agg = compiler.compile_pre_aggregation(&name).unwrap();

        let qp = ctx
            .create_query_properties_with_segments(query_yaml, extra_segments)
            .unwrap();
        let mut matcher = DimensionMatcher::new(ctx.query_tools().query_tools().clone(), &pre_agg);
        matcher
            .try_match(
                qp.dimensions(),
                qp.time_dimensions(),
                qp.dimensions_filters(),
                qp.time_dimensions_filters(),
                qp.segments(),
            )
            .unwrap();
        matcher.result()
    }

    #[test]
    fn test_full_match_dimensions() {
        let ctx = create_test_context();
        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                      - orders.city
                "},
            ),
            MatchState::Full,
        );
    }

    #[test]
    fn test_partial_match_unused_dimension() {
        let ctx = create_test_context();
        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                "},
            ),
            MatchState::Partial,
        );
    }

    #[test]
    fn test_not_matched_missing_dimension() {
        let ctx = create_test_context();
        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.id
                "},
            ),
            MatchState::NotMatched,
        );
    }

    #[test]
    fn test_time_dimension_matching() {
        let ctx = create_test_context();

        assert_eq!(
            match_pre_agg(
                &ctx,
                "daily_countries_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.country
                    time_dimensions:
                      - dimension: orders.created_at
                        granularity: day
                "},
            ),
            MatchState::Full,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "daily_countries_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.country
                    time_dimensions:
                      - dimension: orders.created_at
                        granularity: month
                "},
            ),
            MatchState::Partial,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "daily_countries_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.country
                    time_dimensions:
                      - dimension: orders.created_at
                        granularity: hour
                "},
            ),
            MatchState::NotMatched,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "daily_countries_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.country
                "},
            ),
            MatchState::Partial,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "daily_countries_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    time_dimensions:
                      - dimension: orders.created_at
                        granularity: day
                "},
            ),
            MatchState::Partial,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "daily_countries_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                    time_dimensions:
                      - dimension: orders.updated_at
                        granularity: day
                "},
            ),
            MatchState::NotMatched,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "daily_countries_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.city
                    time_dimensions:
                      - dimension: orders.created_at
                        granularity: day
                "},
            ),
            MatchState::NotMatched,
        );
    }

    #[test]
    fn test_reference_dimension_full_match() {
        let ctx = create_test_context();
        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status_ref
                      - orders.city
                "},
            ),
            MatchState::Full,
        );
    }

    #[test]
    fn test_compound_dimension_matching() {
        let ctx = create_test_context();
        let query = indoc! {"
            measures:
              - orders.count
            dimensions:
              - orders.location_and_status
        "};

        assert_eq!(
            match_pre_agg(&ctx, "compound_dimension_rollup", query),
            MatchState::Full,
        );

        assert_eq!(
            match_pre_agg(&ctx, "base_dimensions_rollup", query),
            MatchState::Partial,
        );

        assert_eq!(
            match_pre_agg(&ctx, "mixed_dimensions_rollup", query),
            MatchState::Partial,
        );
    }

    #[test]
    fn test_filter_matching() {
        let ctx = create_test_context();

        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.city
                    filters:
                      - dimension: orders.status
                        operator: equals
                        values:
                          - shipped
                "},
            ),
            MatchState::Full,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.city
                    filters:
                      - dimension: orders.status
                        operator: contains
                        values:
                          - ship
                "},
            ),
            MatchState::Partial,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                      - orders.city
                    filters:
                      - dimension: orders.id
                        operator: gt
                        values:
                          - \"5\"
                "},
            ),
            MatchState::NotMatched,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    filters:
                      - or:
                          - dimension: orders.status
                            operator: equals
                            values:
                              - shipped
                          - dimension: orders.status
                            operator: equals
                            values:
                              - processing
                "},
            ),
            MatchState::Partial,
        );

        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    filters:
                      - and:
                          - dimension: orders.status
                            operator: equals
                            values:
                              - shipped
                          - dimension: orders.city
                            operator: equals
                            values:
                              - New York
                "},
            ),
            MatchState::Full,
        );
    }

    #[test]
    fn test_segment_full_match() {
        let ctx = create_test_context();
        assert_eq!(
            match_pre_agg(
                &ctx,
                "segment_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                    segments:
                      - orders.high_priority
                    time_dimensions:
                      - dimension: orders.created_at
                        granularity: day
                "},
            ),
            MatchState::Full,
        );
    }

    #[test]
    fn test_segment_partial_match_unused_segment() {
        let ctx = create_test_context();
        // No segment in query, but pre-agg has segment => partial (unused segment coverage)
        assert_eq!(
            match_pre_agg(
                &ctx,
                "segment_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                    time_dimensions:
                      - dimension: orders.created_at
                        granularity: day
                "},
            ),
            MatchState::Partial,
        );
    }

    #[test]
    fn test_constant_member_expression_segment_full_match() {
        let ctx = create_test_context();
        // The ad-hoc `1 = 0` rlsAccessDenied segment references no members, so it's a
        // constant filter pushable on top of any rollup and imposes no coverage
        // requirement: with both rollup dimensions selected the match stays Full
        // (without the fix the bare segment forces NotMatched).
        assert_eq!(
            match_pre_agg_with_segments(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                      - orders.city
                "},
                vec![const_expr_segment("rlsAccessDenied", "orders", "1 = 0")],
            ),
            MatchState::Full,
        );
    }

    #[test]
    fn test_segment_not_matched_missing_in_pre_agg() {
        let ctx = create_test_context();
        // Segment in query, but pre-agg without segments => NotMatched
        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                      - orders.city
                    segments:
                      - orders.high_priority
                "},
            ),
            MatchState::NotMatched,
        );
    }

    #[test]
    fn test_expression_segment_matched_via_dimensions() {
        let ctx = create_test_context();
        // status_filter depends on {CUBE.status}, status is in main_rollup dimensions
        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                      - orders.city
                    segments:
                      - orders.status_filter
                "},
            ),
            MatchState::Partial,
        );
    }

    #[test]
    fn test_expression_segment_not_matched_missing_dependency() {
        let ctx = create_test_context();
        // status_filter depends on {CUBE.status}, but daily_countries_rollup has no status dimension
        assert_eq!(
            match_pre_agg(
                &ctx,
                "daily_countries_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.country
                    segments:
                      - orders.status_filter
                    time_dimensions:
                      - dimension: orders.created_at
                        granularity: day
                "},
            ),
            MatchState::NotMatched,
        );
    }

    #[test]
    fn test_plain_segment_not_matched_via_dimensions() {
        let ctx = create_test_context();
        // high_priority has no {CUBE.x} deps, main_rollup has no segments => NotMatched
        assert_eq!(
            match_pre_agg(
                &ctx,
                "main_rollup",
                indoc! {"
                    measures:
                      - orders.count
                    dimensions:
                      - orders.status
                      - orders.city
                    segments:
                      - orders.high_priority
                "},
            ),
            MatchState::NotMatched,
        );
    }
}
