use super::planners::multi_stage::PlanningScope;
use super::planners::QueryPlanner;
use super::state::State;
use super::QueryProperties;
use crate::logical_plan::OriginalSqlCollector;
use crate::logical_plan::PreAggregationOptimizer;
use crate::logical_plan::PreAggregationUsage;
use crate::logical_plan::RootQuery;
use crate::physical_plan_builder::PhysicalPlanBuilder;
use cubenativeutils::CubeError;
use std::collections::HashMap;
use std::rc::Rc;

pub struct TopLevelPlanner {
    query_tools: Rc<State>,
    request: Rc<QueryProperties>,
    cubestore_support_multistage: bool,
}

impl TopLevelPlanner {
    pub fn new(
        request: Rc<QueryProperties>,
        query_tools: Rc<State>,
        cubestore_support_multistage: bool,
    ) -> Self {
        Self {
            request,
            query_tools,
            cubestore_support_multistage,
        }
    }

    pub fn plan(&self) -> Result<(String, Vec<PreAggregationUsage>), CubeError> {
        let query_planner = QueryPlanner::new(self.request.clone(), self.query_tools.clone());
        let mut scope = PlanningScope::new();
        let query = query_planner.plan(&mut scope)?;
        let logical_plan = Rc::new(
            RootQuery::builder()
                .ctes(scope.into_members())
                .query(query)
                .build(),
        );

        let (optimized_plan, usages) = self.try_pre_aggregations(logical_plan.clone())?;

        // Match-only mode (refresh/metadata path): the caller only needs the matched
        // pre-aggregation(s), not the outer query SQL. Skip the physical build, which for a
        // rolling-window measure would render a time series that requires a date range the
        // refresh path doesn't provide (and which non-generated-time-series dialects can't
        // build without one). The pre-agg's own load SQL is built separately on the JS side.
        if self.request.is_pre_aggregations_match_only() {
            return Ok((String::new(), usages));
        }

        let is_external = if !usages.is_empty() {
            usages.iter().all(|usage| usage.pre_aggregation.external())
        } else {
            false
        };

        let templates = self.query_tools.plan_sql_templates(is_external)?;
        let physical_plan_builder =
            PhysicalPlanBuilder::new(self.query_tools.query_tools().clone(), templates.clone());

        // Substitute a cube's base SQL with its `originalSql` pre-aggregation table when:
        // reading (regular query), or building a rollup that opted in via
        // `useOriginalSqlPreAggregationsInPreAggregation`.
        let original_sql_pre_aggregations = if !self.request.is_pre_aggregation_query()
            || self
                .request
                .use_original_sql_pre_aggregations_in_pre_aggregation()
        {
            OriginalSqlCollector::new(self.query_tools.query_tools().clone())
                .collect(&optimized_plan)?
        } else {
            HashMap::new()
        };

        let physical_plan = physical_plan_builder.build(
            optimized_plan,
            original_sql_pre_aggregations,
            self.request.is_total_query(),
        )?;

        let sql = physical_plan.to_sql(&templates)?;

        Ok((sql, usages))
    }

    fn try_pre_aggregations(
        &self,
        plan: Rc<RootQuery>,
    ) -> Result<(Rc<RootQuery>, Vec<PreAggregationUsage>), CubeError> {
        let result = if !self.request.is_pre_aggregation_query() {
            let mut pre_aggregation_optimizer = PreAggregationOptimizer::new(
                self.query_tools.clone(),
                self.cubestore_support_multistage,
            );
            let disable_external_pre_aggregations =
                self.request.disable_external_pre_aggregations();
            let pre_aggregation_id = self.request.pre_aggregation_id();
            if let Some(result) = pre_aggregation_optimizer.try_optimize(
                plan.clone(),
                disable_external_pre_aggregations,
                pre_aggregation_id,
            )? {
                let usages = pre_aggregation_optimizer.take_usages();
                if !usages.is_empty() {
                    (result, usages)
                } else {
                    (plan.clone(), Vec::new())
                }
            } else {
                (plan.clone(), Vec::new())
            }
        } else {
            (plan.clone(), Vec::new())
        };
        Ok(result)
    }
}
