hamelin_lib 0.21.4

use std::collections::{HashMap, HashSet};

use crate::rewrite_table_references;
use crate::tree::ast::dataset_identifier::DatasetIdentifier;
use crate::tree::ast::identifier::SimpleIdentifier;
use crate::tree::ast::query::Query;
use crate::tree::ast::ParseWithErrors;
use crate::tree::builder::query::QueryBuilder;
use crate::tree::transform::{
    catalog_key_aliases, catalog_keys_for_dataset, RewriteTableReferencesError,
};

use super::def_name_registry::DefNameRegistry;
use super::rename_conflicting_defs::rename_conflicting_defs;
use super::CombineQueriesError;

#[derive(Debug, Clone, PartialEq)]
pub struct CteInlineEntry<Q> {
    pub catalog_key: String,
    pub query: Q,
    pub cte_name: String,
}

/// Inline upstream query as a string (not yet parsed).
pub type CteInline = CteInlineEntry<String>;

/// Inline upstream query as a parsed AST.
pub type ParsedCteInline = CteInlineEntry<Query>;

impl Eq for CteInlineEntry<String> {}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CombineQueriesOk {
    pub query: String,
    pub cte_names: HashMap<String, String>,
}

impl From<RewriteTableReferencesError> for CombineQueriesError {
    fn from(value: RewriteTableReferencesError) -> Self {
        match value {
            RewriteTableReferencesError::QueryHasParseErrors => Self::QueryHasParseErrors,
            RewriteTableReferencesError::InvalidCteName(msg) => Self::InvalidCteName(msg),
            RewriteTableReferencesError::AmbiguousReplacement(msg) => Self::Message(msg),
        }
    }
}

/// Merge already-parsed inline upstream queries as tabular `DEF`s, rewrite catalog refs, and
/// attach `main_query`.
pub fn combine_parsed_queries_for_inlining(
    main_query: Query,
    cte_inlines: &[ParsedCteInline],
    default_space: Option<&SimpleIdentifier>,
) -> Result<CombineQueriesOk, CombineQueriesError> {
    let ordered_inlines = sort_inlines_for_merge(cte_inlines, default_space)?;

    let mut def_names = DefNameRegistry::new();
    def_names.extend_from_query(&main_query);

    let mut builder = QueryBuilder::new();
    let mut allocated_by_catalog_key: HashMap<String, String> = HashMap::new();

    for inline in &ordered_inlines {
        if allocated_by_catalog_key.contains_key(&inline.catalog_key) {
            return Err(CombineQueriesError::DuplicateCatalogKey(
                inline.catalog_key.clone(),
            ));
        }

        let renamed = rename_conflicting_defs(&inline.query, def_names.used_names())?;
        def_names.extend_from_query(&renamed);

        let base_name = inline.cte_name.replace(':', "_");
        let allocated = def_names.allocate_cte_name(&base_name)?;
        allocated_by_catalog_key.insert(inline.catalog_key.clone(), allocated.as_str().to_string());

        builder = builder.merge_as_cte(renamed, allocated);
    }

    let combined = builder.merge_as_main(main_query);
    let rewritten = rewrite_table_references(&combined, &allocated_by_catalog_key, default_space)?;

    Ok(CombineQueriesOk {
        query: rewritten.to_string(),
        cte_names: allocated_by_catalog_key,
    })
}

/// Parse query strings and delegate to [`combine_parsed_queries_for_inlining`].
pub fn combine_queries_for_inlining(
    main_query: &str,
    cte_inlines: &[CteInline],
    default_space: Option<&SimpleIdentifier>,
) -> Result<CombineQueriesOk, CombineQueriesError> {
    let (main_parsed, main_errs) = Query::parse_with_errors(main_query);
    if !main_errs.is_empty() {
        return Err(CombineQueriesError::QueryHasParseErrors);
    }

    let mut parsed_inlines = Vec::with_capacity(cte_inlines.len());
    for inline in cte_inlines {
        let (parsed, errs) = Query::parse_with_errors(&inline.query);
        if !errs.is_empty() {
            return Err(CombineQueriesError::QueryHasParseErrors);
        }
        parsed_inlines.push(ParsedCteInline {
            catalog_key: inline.catalog_key.clone(),
            query: parsed,
            cte_name: inline.cte_name.clone(),
        });
    }

    combine_parsed_queries_for_inlining(main_parsed, &parsed_inlines, default_space)
}

fn inline_catalog_keys_referenced_by_query(
    query: &Query,
    inline_catalog_keys: &HashSet<String>,
    default_space: Option<&SimpleIdentifier>,
) -> Result<HashSet<String>, CombineQueriesError> {
    let datasets = query
        .datasets(default_space)
        .map_err(|e| CombineQueriesError::Message(e.to_string()))?;

    let mut refs = HashSet::new();
    for ds in datasets {
        let ds_id: DatasetIdentifier = ds.into();
        for key in inline_catalog_keys {
            let aliases = catalog_key_aliases(key, default_space);
            let ref_keys = catalog_keys_for_dataset(&ds_id, default_space);
            if ref_keys
                .iter()
                .any(|ref_key| aliases.iter().any(|alias| alias == ref_key))
            {
                refs.insert(key.clone());
            }
        }
    }
    Ok(refs)
}

/// Order inlines so each hoisted CTE is defined before any other inline that references its
/// catalog key.
fn sort_inlines_for_merge(
    inlines: &[ParsedCteInline],
    default_space: Option<&SimpleIdentifier>,
) -> Result<Vec<ParsedCteInline>, CombineQueriesError> {
    if inlines.len() <= 1 {
        return Ok(inlines.to_vec());
    }

    let catalog_keys: HashSet<String> = inlines.iter().map(|i| i.catalog_key.clone()).collect();
    let n = inlines.len();
    let mut depends_on: Vec<HashSet<usize>> = vec![HashSet::new(); n];

    for (i, inline) in inlines.iter().enumerate() {
        let refs =
            inline_catalog_keys_referenced_by_query(&inline.query, &catalog_keys, default_space)?;
        for (j, other) in inlines.iter().enumerate() {
            if i != j && refs.contains(&other.catalog_key) {
                depends_on[i].insert(j);
            }
        }
    }

    let mut in_degree: Vec<usize> = depends_on.iter().map(|deps| deps.len()).collect();
    let mut ready: Vec<usize> = (0..n).filter(|&i| in_degree[i] == 0).collect();
    let mut order = Vec::with_capacity(n);

    while !ready.is_empty() {
        ready.sort_unstable();
        let i = ready.remove(0);
        order.push(i);
        for j in 0..n {
            if depends_on[j].contains(&i) {
                in_degree[j] -= 1;
                if in_degree[j] == 0 {
                    ready.push(j);
                }
            }
        }
    }

    if order.len() != n {
        return Err(CombineQueriesError::Message(
            "cyclic catalog dependency among cte_inlines".into(),
        ));
    }

    Ok(order.into_iter().map(|i| inlines[i].clone()).collect())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn inline(catalog_key: &str, query: &str, cte_name: &str) -> CteInline {
        CteInline {
            catalog_key: catalog_key.to_string(),
            query: query.to_string(),
            cte_name: cte_name.to_string(),
        }
    }

    #[test]
    fn hoists_upstream_defs_and_rewrites_main_ref() {
        let upstream = "DEF network_events = FROM events:sysmon_events\n| WHERE event.code == \"3\";\n\nFROM network_events\n| SET host = host.name";
        let main_query = "FROM events:my_view\n| LIMIT 10";
        let inlines = vec![inline("events:my_view", upstream, "events_my_view")];

        let result = combine_queries_for_inlining(main_query, &inlines, None).unwrap();

        assert!(!result
            .query
            .contains("DEF events_my_view = DEF network_events"));
        assert!(result
            .query
            .contains("DEF network_events = FROM events:sysmon_events"));
        assert!(result
            .query
            .contains("DEF events_my_view = FROM network_events"));
        assert!(result.query.contains("FROM events_my_view\n| LIMIT 10"));
        assert_eq!(
            result.cte_names.get("events:my_view"),
            Some(&"events_my_view".to_string())
        );
    }

    #[test]
    fn auto_renames_cte_on_collision_with_existing_def() {
        let upstream = "FROM events:upstream | SET x = 1 AS int";
        let main_query = "DEF events_my_view = FROM events:base | LIMIT 1;\nFROM events:upstream | SET y = 2 AS int";
        let inlines = vec![inline("events:upstream", upstream, "events_my_view")];

        let result = combine_queries_for_inlining(main_query, &inlines, None).unwrap();

        assert_eq!(
            result.cte_names.get("events:upstream"),
            Some(&"events_my_view_2".to_string())
        );
        assert!(result
            .query
            .contains("DEF events_my_view_2 = FROM events_my_view_2"));
        assert!(result
            .query
            .contains("FROM events_my_view_2\n| SET y = 2 AS int"));
    }

    #[test]
    fn chains_multiple_inlines() {
        let inner = "FROM events:connector | SET a = 1 AS int";
        let middle = "FROM events:inner_table | SET b = 2 AS int";
        let main_query = "FROM events:outer_table | LIMIT 5";
        let inlines = vec![
            inline("events:inner_table", inner, "events_inner_table"),
            inline("events:outer_table", middle, "events_outer_table"),
        ];

        let result = combine_queries_for_inlining(main_query, &inlines, None).unwrap();

        assert!(result
            .query
            .contains("DEF events_inner_table = FROM events:connector"));
        assert!(result
            .query
            .contains("DEF events_outer_table = FROM events_inner_table"));
        assert!(result.query.contains("FROM events_outer_table\n| LIMIT 5"));
    }

    #[test]
    fn uses_allocated_name_when_two_inlines_share_base_cte_name() {
        let upstream_a = "FROM events:dataset_a | SET x = 1 AS int";
        let upstream_b = "FROM events:dataset_b | SET y = 2 AS int";
        let main_query = "FROM events:dataset_a, events:dataset_b | LIMIT 1";
        let inlines = vec![
            inline("events:dataset_a", upstream_a, "shared_cte"),
            inline("events:dataset_b", upstream_b, "shared_cte"),
        ];

        let result = combine_queries_for_inlining(main_query, &inlines, None).unwrap();

        assert_eq!(
            result.cte_names.get("events:dataset_a"),
            Some(&"shared_cte".to_string())
        );
        assert_eq!(
            result.cte_names.get("events:dataset_b"),
            Some(&"shared_cte_2".to_string())
        );
        assert!(result.query.contains("DEF shared_cte = FROM shared_cte"));
        assert!(result
            .query
            .contains("DEF shared_cte_2 = FROM shared_cte_2"));
        assert!(result.query.contains("FROM shared_cte, shared_cte_2"));
    }

    #[test]
    fn renames_duplicate_internal_def_names_across_inlines() {
        let upstream_a = "DEF network_events = FROM events:sysmon_a | LIMIT 1;\nFROM network_events | SET a = 1 AS int";
        let upstream_b = "DEF network_events = FROM events:sysmon_b | LIMIT 1;\nFROM network_events | SET b = 2 AS int";
        let main_query = "FROM events:view_a, events:view_b | LIMIT 1";
        let inlines = vec![
            inline("events:view_a", upstream_a, "events_view_a"),
            inline("events:view_b", upstream_b, "events_view_b"),
        ];

        let result = combine_queries_for_inlining(main_query, &inlines, None).unwrap();

        assert!(result
            .query
            .contains("DEF network_events = FROM events:sysmon_a"));
        assert!(result
            .query
            .contains("DEF network_events_2 = FROM events:sysmon_b"));
        assert!(result
            .query
            .contains("DEF events_view_a = FROM network_events"));
        assert!(result
            .query
            .contains("DEF events_view_b = FROM network_events_2"));
    }

    #[test]
    fn rewrites_own_catalog_ref_in_inline_body() {
        let upstream = "FROM events:my_view | LIMIT 10";
        let main_query = "FROM events:my_view | SET x = 1 AS int";
        let inlines = vec![inline("events:my_view", upstream, "events_my_view")];

        let result = combine_queries_for_inlining(main_query, &inlines, None).unwrap();

        assert!(!result.query.contains("FROM events:my_view"));
        assert!(result
            .query
            .contains("DEF events_my_view = FROM events_my_view"));
    }

    #[test]
    fn avoids_hoisted_cte_name_collision_with_scalar_def() {
        let upstream = "DEF my_limit = 10 AS int;\nFROM events:source | LIMIT my_limit";
        let main_query = "FROM events:dataset | LIMIT 1";
        let inlines = vec![inline("events:dataset", upstream, "my_limit")];

        let result = combine_queries_for_inlining(main_query, &inlines, None).unwrap();

        assert_eq!(
            result.cte_names.get("events:dataset"),
            Some(&"my_limit_2".to_string())
        );
        assert!(result.query.contains("DEF my_limit = 10 AS int"));
        assert!(result.query.contains("DEF my_limit_2 = FROM events:source"));
    }

    #[test]
    fn rejects_duplicate_catalog_key_inlines() {
        let upstream_a = "FROM events:dataset_a | SET x = 1 AS int";
        let upstream_b = "FROM events:dataset_b | SET y = 2 AS int";
        let main_query = "FROM events:dup | LIMIT 1";
        let inlines = vec![
            inline("events:dup", upstream_a, "events_dup"),
            inline("events:dup", upstream_b, "events_dup_alt"),
        ];

        let err = combine_queries_for_inlining(main_query, &inlines, None).unwrap_err();
        assert!(matches!(
            err,
            CombineQueriesError::DuplicateCatalogKey(key) if key == "events:dup"
        ));
    }

    #[test]
    fn rejects_conflicting_scalar_def_across_merge() {
        let upstream = "DEF row_limit = 20 AS int;\nFROM events:source | LIMIT row_limit";
        let main_query = "DEF row_limit = 10 AS int;\nFROM events:dataset | LIMIT 1";
        let inlines = vec![inline("events:dataset", upstream, "events_dataset")];

        let err = combine_queries_for_inlining(main_query, &inlines, None).unwrap_err();
        assert!(
            matches!(err, CombineQueriesError::Message(msg) if msg.contains("scalar DEF name collision"))
        );
    }

    #[test]
    fn rewrites_dependent_inline_when_dependency_comes_later() {
        let inner = "FROM events:connector | SET a = 1 AS int";
        let outer = "FROM events:inner_table | SET b = 2 AS int";
        let main_query = "FROM events:outer_table | LIMIT 5";
        let inlines = vec![
            inline("events:outer_table", outer, "events_outer_table"),
            inline("events:inner_table", inner, "events_inner_table"),
        ];

        let result = combine_queries_for_inlining(main_query, &inlines, None).unwrap();

        assert!(!result.query.contains("FROM events:inner_table"));
        assert!(result
            .query
            .contains("DEF events_outer_table = FROM events_inner_table"));
        assert!(result
            .query
            .contains("DEF events_inner_table = FROM events:connector"));
        assert!(result.query.contains("FROM events_outer_table\n| LIMIT 5"));

        let inner_pos = result.query.find("DEF events_inner_table").unwrap();
        let outer_pos = result.query.find("DEF events_outer_table").unwrap();
        assert!(
            inner_pos < outer_pos,
            "inner CTE must be defined before outer when dependency comes later in input"
        );
    }

    #[test]
    fn rejects_cyclic_inline_catalog_dependency() {
        let a = "FROM events:b | SET x = 1 AS int";
        let b = "FROM events:a | SET y = 2 AS int";
        let main_query = "FROM events:a | LIMIT 1";
        let inlines = vec![
            inline("events:a", a, "events_a"),
            inline("events:b", b, "events_b"),
        ];

        let err = combine_queries_for_inlining(main_query, &inlines, None).unwrap_err();
        assert!(matches!(err, CombineQueriesError::Message(_)));
    }

    #[test]
    fn reports_dataset_error_when_sorting_inlines_without_default_space() {
        let inner = "FROM my_table | SET a = 1 AS int";
        let outer = "FROM events:inner_table | SET b = 2 AS int";
        let main_query = "FROM events:outer_table | LIMIT 5";
        let inlines = vec![
            inline("events:outer_table", outer, "events_outer_table"),
            inline("events:inner_table", inner, "events_inner_table"),
        ];

        let err = combine_queries_for_inlining(main_query, &inlines, None).unwrap_err();
        assert!(matches!(err, CombineQueriesError::Message(msg) if !msg.is_empty()));
    }

    #[test]
    fn rejects_invalid_requested_cte_name() {
        let main_query = "FROM events:dataset | LIMIT 1";
        let inlines = vec![inline(
            "events:dataset",
            "FROM events:source | SET x = 1 AS int",
            "123invalid",
        )];

        let err = combine_queries_for_inlining(main_query, &inlines, None).unwrap_err();
        assert!(matches!(err, CombineQueriesError::InvalidCteName(_)));
    }
}