//! Output timestamp column metadata: when a column is defined by a `TsTrunc` (`@` truncation),
//! exposes [`TruncUnit`] and `multiplier` as inferred by the compiler. This does not encode
//! every SQL back-end's bin edges, only the Hamelin `TsTrunc` shape.

use std::sync::Arc;

use hamelin_lib::catalog::{Catalog, CatalogProvider};
use hamelin_lib::err::ContextualTranslationErrors;
use hamelin_lib::provider::EnvironmentProvider;
use hamelin_lib::tree::ast::expression::TruncUnit;
use hamelin_lib::tree::ast::identifier::{Identifier, SimpleIdentifier};
use hamelin_lib::tree::options::TypeCheckOptions;
use hamelin_lib::tree::typed_ast::clause::Projections;
use hamelin_lib::tree::typed_ast::command::TypedCommandKind;
use hamelin_lib::tree::typed_ast::command::{
    TypedAggCommand, TypedAppendCommand, TypedCommand, TypedDistinctCommand, TypedExplodeCommand,
    TypedJoinCommand, TypedLookupCommand, TypedMatchCommand, TypedNestCommand, TypedParseCommand,
    TypedSelectCommand, TypedSetCommand, TypedSuppressCommand, TypedUnnestCommand,
    TypedWindowCommand,
};
use hamelin_lib::tree::typed_ast::environment::TypeEnvironment;
use hamelin_lib::tree::typed_ast::expression::{
    TypedCast, TypedExpression, TypedExpressionKind, TypedFieldLookup, TypedTsTrunc,
};
use hamelin_lib::tree::typed_ast::query::TypedStatement;
use hamelin_lib::types::Type;
use ordermap::OrderMap;
use serde::Serialize;
use tsify::Tsify;
use wasm_bindgen::prelude::wasm_bindgen;
use wasm_bindgen::JsValue;

use crate::compile::typecheck_query;

type LineageMap = OrderMap<Identifier, Arc<TypedExpression>>;

/// Follows [`TypedExpressionKind::FieldReference`] through `map`, then strips [`TypedCast`]
/// / [`TypedFieldLookup`] to reach a [`TypedTsTrunc`] if present.
fn resolve_to_ts_trunc<'a>(
    start: &'a Arc<TypedExpression>,
    map: &'a LineageMap,
) -> Option<&'a TypedTsTrunc> {
    let mut current = start;
    let mut depth: u32 = 0;
    const MAX: u32 = 128;
    let mut seen: Vec<Identifier> = vec![];

    while depth < MAX {
        match &current.kind {
            TypedExpressionKind::FieldReference(fr) => {
                let id = fr.field_name.clone().valid().ok().map(Into::into)?;
                if seen.contains(&id) {
                    break;
                }
                seen.push(id.clone());
                if let Some(next) = map.get(&id) {
                    current = next;
                    depth += 1;
                } else {
                    break;
                }
            }
            _ => break,
        }
    }

    let mut p = current;
    let mut peel = 0u32;
    const MAX_PEEL: u32 = 32;
    while peel < MAX_PEEL {
        match &p.kind {
            TypedExpressionKind::Cast(TypedCast { value, .. }) => p = value,
            TypedExpressionKind::FieldLookup(TypedFieldLookup { value, .. }) => p = value,
            _ => break,
        }
        peel += 1;
    }

    match &p.kind {
        TypedExpressionKind::TsTrunc(t) => Some(t),
        _ => None,
    }
}

fn filter_map_to_env(input: &LineageMap, output: &TypeEnvironment) -> LineageMap {
    let mut m = LineageMap::new();
    for (name, _) in output.as_struct().iter() {
        let id: Identifier = name.clone().into();
        if let Some(e) = input.get(&id) {
            m.insert(id, e.clone());
        }
    }
    m
}

/// Mirrors [`TypeEnvironment::extend`] (SET / WINDOW): assignments shadow parent keys by top segment.
fn extend_lineage(
    parent: &LineageMap,
    assignment_exprs: &[(Identifier, Arc<TypedExpression>)],
) -> LineageMap {
    let mut touched: Vec<SimpleIdentifier> = Vec::new();
    let mut out = LineageMap::new();
    for (id, expr) in assignment_exprs {
        let top = id.first();
        if !touched.contains(&top) {
            touched.push(top.clone());
            let top_id: Identifier = top.clone().into();
            if let Some(p) = parent.get(&top_id) {
                out.insert(top_id, p.clone());
            }
        }
        out.insert(id.clone(), expr.clone());
    }
    for (k, v) in parent {
        if !touched.contains(&k.first()) {
            out.insert(k.clone(), v.clone());
        }
    }
    out
}

fn lineage_for_projections(p: &Projections) -> Vec<(Identifier, Arc<TypedExpression>)> {
    p.assignments
        .iter()
        .filter_map(|a| {
            a.identifier
                .clone()
                .valid()
                .ok()
                .map(|id| (id, a.expression.clone()))
        })
        .collect()
}

fn from_projection_pairs(pairs: Vec<(Identifier, Arc<TypedExpression>)>) -> LineageMap {
    let mut m = LineageMap::new();
    for (id, e) in pairs {
        m.insert(id, e);
    }
    m
}

fn apply_lineage_step(map: &LineageMap, cmd: &TypedCommand) -> LineageMap {
    match &cmd.kind {
        TypedCommandKind::Where(_)
        | TypedCommandKind::Within(_)
        | TypedCommandKind::Limit(_)
        | TypedCommandKind::Trimstrings(_)
        | TypedCommandKind::Sort(_)
        | TypedCommandKind::From(_)
        | TypedCommandKind::Union(_)
        | TypedCommandKind::Error(_) => map.clone(),

        TypedCommandKind::Suppress(TypedSuppressCommand { group_by, .. }) => {
            extend_lineage(map, &lineage_for_projections(group_by))
        }

        TypedCommandKind::Set(TypedSetCommand { projections, .. }) => {
            extend_lineage(map, &lineage_for_projections(projections))
        }

        TypedCommandKind::Select(TypedSelectCommand { projections, .. }) => {
            // Overlay projections on the incoming map so `SELECT t = a` can still resolve
            // field-reference chains to columns dropped from the new output schema.
            let mut m = map.clone();
            for a in &projections.assignments {
                if let Ok(id) = a.identifier.clone().valid() {
                    m.insert(id, a.expression.clone());
                }
            }
            m
        }

        TypedCommandKind::Drop(d) => {
            let mut m = map.clone();
            for id in &d.dropped_fields {
                m.remove(id);
            }
            m
        }

        TypedCommandKind::Parse(TypedParseCommand { .. }) => {
            filter_map_to_env(map, &cmd.output_schema)
        }

        TypedCommandKind::Join(TypedJoinCommand { .. })
        | TypedCommandKind::Lookup(TypedLookupCommand { .. }) => {
            filter_map_to_env(map, &cmd.output_schema)
        }

        TypedCommandKind::Distinct(TypedDistinctCommand { keys, .. }) => {
            from_projection_pairs(lineage_for_projections(keys))
        }

        TypedCommandKind::Agg(TypedAggCommand {
            group_by,
            aggregates,
            ..
        }) => {
            let mut pairs = lineage_for_projections(group_by);
            pairs.extend(lineage_for_projections(aggregates));
            from_projection_pairs(pairs)
        }

        TypedCommandKind::Window(TypedWindowCommand {
            projections,
            group_by,
            ..
        }) => {
            let mut v = lineage_for_projections(projections);
            v.extend(lineage_for_projections(group_by));
            extend_lineage(map, &v)
        }

        TypedCommandKind::Match(TypedMatchCommand { agg, group_by, .. }) => {
            let mut pairs = lineage_for_projections(agg);
            pairs.extend(lineage_for_projections(group_by));
            let mut n = from_projection_pairs(pairs);
            for (name, _) in cmd.output_schema.as_struct().iter() {
                let id: Identifier = name.clone().into();
                if n.contains_key(&id) {
                    continue;
                }
                if let Some(e) = map.get(&id) {
                    n.insert(id, e.clone());
                }
            }
            n
        }

        TypedCommandKind::Append(TypedAppendCommand { .. }) => {
            filter_map_to_env(map, &cmd.output_schema)
        }

        TypedCommandKind::Explode(TypedExplodeCommand { items, .. }) => {
            let mut m = map.clone();
            for it in items {
                if let Ok(id) = it.assignment.identifier.clone().valid() {
                    m.insert(id, it.assignment.expression.clone());
                }
            }
            m
        }

        TypedCommandKind::Unnest(TypedUnnestCommand { .. }) => {
            filter_map_to_env(map, &cmd.output_schema)
        }

        TypedCommandKind::Nest(TypedNestCommand { .. }) => LineageMap::new(),
    }
}

fn output_ts_columns(typed: &TypedStatement) -> Vec<OutputTimestampColumn> {
    let Some(valid) = typed.pipeline.valid_ref().ok() else {
        return vec![];
    };
    let commands = &valid.commands;
    let mut line = LineageMap::new();
    for c in commands {
        line = apply_lineage_step(&line, c);
    }
    let out_env = &valid.final_schema;

    let mut res = vec![];
    for (name, typ) in out_env.as_struct().iter() {
        if *typ != Type::Timestamp {
            continue;
        }
        let id: Identifier = name.clone().into();
        let (unit, multiplier) = line
            .get(&id)
            .and_then(|def| resolve_to_ts_trunc(def, &line))
            .and_then(|t| {
                t.unit
                    .as_fixed()
                    .map(|(unit, multiplier)| (Some(unit), Some(multiplier)))
            })
            .unwrap_or((None, None));
        res.push(OutputTimestampColumn {
            name: name.to_string(),
            unit,
            multiplier,
        });
    }
    res
}

/// One output column of type `timestamp` with optional `TruncUnit` and `multiplier` from `@` truncation.
#[derive(Debug, Clone, Serialize, Tsify)]
#[tsify(into_wasm_abi)]
pub struct OutputTimestampColumn {
    pub name: String,
    #[tsify(optional)]
    pub unit: Option<TruncUnit>,
    #[tsify(optional)]
    pub multiplier: Option<u32>,
}

#[derive(Debug, Clone, Serialize, Tsify)]
#[tsify(into_wasm_abi)]
pub struct GetOutputTimestampColumnsResultOk {
    pub columns: Vec<OutputTimestampColumn>,
}

#[derive(Serialize, Tsify)]
#[tsify(into_wasm_abi)]
pub enum GetOutputTimestampColumnsResult {
    Ok(GetOutputTimestampColumnsResultOk),
    Err(ContextualTranslationErrors),
}

/// Lists top-level `timestamp` output columns and, when the defining expression
/// reduces to a `@` truncation, `unit` and `multiplier` from [`TypedTsTrunc`].
#[wasm_bindgen]
pub fn get_output_timestamp_columns(
    query: String,
    catalog: Catalog,
) -> Result<GetOutputTimestampColumnsResult, JsValue> {
    let catalog_provider = Arc::new(
        CatalogProvider::try_from(catalog)
            .map_err(|e| JsValue::from_str(e.to_string().as_str()))?,
    );
    Ok(get_output_timestamp_columns_with_provider(
        query,
        catalog_provider,
    ))
}

pub fn get_output_timestamp_columns_with_provider(
    query: String,
    provider: Arc<dyn EnvironmentProvider>,
) -> GetOutputTimestampColumnsResult {
    let opts = TypeCheckOptions::builder().provider(provider).build();
    match typecheck_query(query, opts) {
        Ok(typed_statement) => {
            GetOutputTimestampColumnsResult::Ok(GetOutputTimestampColumnsResultOk {
                columns: output_ts_columns(&typed_statement),
            })
        }
        Err(errors) => GetOutputTimestampColumnsResult::Err(errors),
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use hamelin_lib::catalog::{Catalog, CatalogBuilder, DataSetBuilder};
    use hamelin_lib::tree::ast::identifier::SimpleIdentifier;
    use hamelin_lib::types::TIMESTAMP;

    use super::*;

    fn test_provider() -> Arc<CatalogProvider> {
        let catalog: Catalog = CatalogBuilder::default()
            .with_data_set(
                DataSetBuilder::new(SimpleIdentifier::new("logs").into())
                    .with(SimpleIdentifier::new("timestamp"), TIMESTAMP)
                    .with(SimpleIdentifier::new("x"), hamelin_lib::types::INT),
            )
            .build();
        Arc::new(CatalogProvider::try_from(catalog).expect("test catalog ok"))
    }

    #[test]
    fn get_output_by_trunc_hamelin_2h() {
        let q = "FROM logs | AGG count() BY ts = timestamp@2h".to_string();
        let r = get_output_timestamp_columns_with_provider(q, test_provider());
        let GetOutputTimestampColumnsResult::Ok(GetOutputTimestampColumnsResultOk { columns }) = r
        else {
            panic!("expected ok");
        };
        let t = columns.iter().find(|c| c.name == "ts").expect("ts col");
        assert_eq!(t.unit, Some(TruncUnit::Hour));
        assert_eq!(t.multiplier, Some(2));
    }

    /// SELECT must overlay assignments on the incoming lineage so `t = a` can resolve
    /// after the output schema drops `a` as a top-level field.
    #[test]
    fn get_output_select_aliasing_field_with_trunc() {
        let q = "FROM logs | SET a = timestamp@1d | SELECT t = a".to_string();
        let r = get_output_timestamp_columns_with_provider(q, test_provider());
        let GetOutputTimestampColumnsResult::Ok(GetOutputTimestampColumnsResultOk { columns }) = r
        else {
            panic!("expected ok");
        };
        let t = columns.iter().find(|c| c.name == "t").expect("t col");
        assert_eq!(t.unit, Some(TruncUnit::Day));
        assert_eq!(t.multiplier, Some(1));
    }

    #[test]
    fn get_output_agg_bare_by_timestamp_2h() {
        let q = "FROM logs | AGG count() BY timestamp@2h".to_string();
        let r = get_output_timestamp_columns_with_provider(q, test_provider());
        let GetOutputTimestampColumnsResult::Ok(GetOutputTimestampColumnsResultOk { columns }) = r
        else {
            panic!("expected ok");
        };
        let t = columns
            .iter()
            .find(|c| c.name == "timestamp")
            .expect("timestamp col");
        assert_eq!(t.unit, Some(TruncUnit::Hour));
        assert_eq!(t.multiplier, Some(2));
    }

    #[test]
    fn get_output_agg_bare_by_timestamp_m_is_minute() {
        let q = "FROM logs | AGG count() BY timestamp@m".to_string();
        let r = get_output_timestamp_columns_with_provider(q, test_provider());
        let GetOutputTimestampColumnsResult::Ok(GetOutputTimestampColumnsResultOk { columns }) = r
        else {
            panic!("expected ok");
        };
        let t = columns
            .iter()
            .find(|c| c.name == "timestamp")
            .expect("timestamp col");
        assert_eq!(t.unit, Some(TruncUnit::Minute));
        assert_eq!(t.multiplier, Some(1));
    }
}
