//! Shared predicates for tag analysis.

use super::walk::deep_dfs;
use polyglot_sql::expressions::Expression;

/// Returns `true` if the expression is any function-like AST node.
pub fn is_any_function(expr: &Expression) -> bool {
    matches!(
        expr,
        Expression::Function(_)
            | Expression::AggregateFunction(_)
            | Expression::WindowFunction(_)
            | Expression::Upper(_)
            | Expression::Lower(_)
            | Expression::Length(_)
            | Expression::Trim(_)
            | Expression::LTrim(_)
            | Expression::RTrim(_)
            | Expression::Replace(_)
            | Expression::Reverse(_)
            | Expression::Left(_)
            | Expression::Right(_)
            | Expression::Repeat(_)
            | Expression::Lpad(_)
            | Expression::Rpad(_)
            | Expression::Split(_)
            | Expression::RegexpLike(_)
            | Expression::RegexpReplace(_)
            | Expression::RegexpExtract(_)
            | Expression::Overlay(_)
            | Expression::ConcatWs(_)
            | Expression::Substring(_)
            | Expression::Contains(_)
            | Expression::StartsWith(_)
            | Expression::EndsWith(_)
            | Expression::Position(_)
            | Expression::Initcap(_)
            | Expression::Ascii(_)
            | Expression::Chr(_)
            | Expression::CharFunc(_)
            | Expression::Soundex(_)
            | Expression::Levenshtein(_)
            | Expression::ByteLength(_)
            | Expression::Hex(_)
            | Expression::LowerHex(_)
            | Expression::Unicode(_)
            | Expression::Abs(_)
            | Expression::Round(_)
            | Expression::Floor(_)
            | Expression::Ceil(_)
            | Expression::Power(_)
            | Expression::Sqrt(_)
            | Expression::Cbrt(_)
            | Expression::Ln(_)
            | Expression::Log(_)
            | Expression::Exp(_)
            | Expression::Sign(_)
            | Expression::Greatest(_)
            | Expression::Least(_)
            | Expression::ModFunc(_)
            | Expression::Random(_)
            | Expression::Rand(_)
            | Expression::TruncFunc(_)
            | Expression::Pi(_)
            | Expression::Radians(_)
            | Expression::Degrees(_)
            | Expression::Sin(_)
            | Expression::Cos(_)
            | Expression::Tan(_)
            | Expression::Asin(_)
            | Expression::Acos(_)
            | Expression::Atan(_)
            | Expression::Atan2(_)
            | Expression::IsNan(_)
            | Expression::IsInf(_)
            | Expression::IntDiv(_)
            | Expression::Cast(_)
            | Expression::TryCast(_)
            | Expression::SafeCast(_)
            | Expression::Coalesce(_)
            | Expression::NullIf(_)
            | Expression::IfFunc(_)
            | Expression::IfNull(_)
            | Expression::Nvl(_)
            | Expression::Nvl2(_)
            | Expression::Decode(_)
            | Expression::Case(_)
            | Expression::Count(_)
            | Expression::Sum(_)
            | Expression::Avg(_)
            | Expression::Min(_)
            | Expression::Max(_)
            | Expression::GroupConcat(_)
            | Expression::StringAgg(_)
            | Expression::ListAgg(_)
            | Expression::ArrayAgg(_)
            | Expression::CountIf(_)
            | Expression::SumIf(_)
            | Expression::Stddev(_)
            | Expression::StddevPop(_)
            | Expression::StddevSamp(_)
            | Expression::Variance(_)
            | Expression::VarPop(_)
            | Expression::VarSamp(_)
            | Expression::Median(_)
            | Expression::Mode(_)
            | Expression::First(_)
            | Expression::Last(_)
            | Expression::AnyValue(_)
            | Expression::ApproxDistinct(_)
            | Expression::ApproxCountDistinct(_)
            | Expression::ApproxPercentile(_)
            | Expression::Percentile(_)
            | Expression::LogicalAnd(_)
            | Expression::LogicalOr(_)
            | Expression::Skewness(_)
            | Expression::BitwiseCount(_)
            | Expression::ArrayConcatAgg(_)
            | Expression::ArrayUniqueAgg(_)
            | Expression::BoolXorAgg(_)
            | Expression::BitwiseAndAgg(_)
            | Expression::BitwiseOrAgg(_)
            | Expression::BitwiseXorAgg(_)
            | Expression::RowNumber(_)
            | Expression::Rank(_)
            | Expression::DenseRank(_)
            | Expression::NTile(_)
            | Expression::Lead(_)
            | Expression::Lag(_)
            | Expression::FirstValue(_)
            | Expression::LastValue(_)
            | Expression::NthValue(_)
            | Expression::PercentRank(_)
            | Expression::CumeDist(_)
            | Expression::PercentileCont(_)
            | Expression::PercentileDisc(_)
            | Expression::CurrentDate(_)
            | Expression::CurrentTime(_)
            | Expression::CurrentTimestamp(_)
            | Expression::CurrentTimestampLTZ(_)
            | Expression::AtTimeZone(_)
            | Expression::DateAdd(_)
            | Expression::DateSub(_)
            | Expression::DateDiff(_)
            | Expression::DateTrunc(_)
            | Expression::Extract(_)
            | Expression::ToDate(_)
            | Expression::ToTimestamp(_)
            | Expression::Date(_)
            | Expression::Time(_)
            | Expression::DateFromUnixDate(_)
            | Expression::UnixDate(_)
            | Expression::UnixSeconds(_)
            | Expression::UnixMillis(_)
            | Expression::UnixMicros(_)
            | Expression::UnixToTimeStr(_)
            | Expression::TimeStrToDate(_)
            | Expression::DateToDi(_)
            | Expression::DiToDate(_)
            | Expression::TsOrDiToDi(_)
            | Expression::TsOrDsToDatetime(_)
            | Expression::TsOrDsToTimestamp(_)
            | Expression::YearOfWeek(_)
            | Expression::YearOfWeekIso(_)
            | Expression::DateFormat(_)
            | Expression::FormatDate(_)
            | Expression::Year(_)
            | Expression::Month(_)
            | Expression::Day(_)
            | Expression::Hour(_)
            | Expression::Minute(_)
            | Expression::Second(_)
            | Expression::DayOfWeek(_)
            | Expression::DayOfWeekIso(_)
            | Expression::DayOfMonth(_)
            | Expression::DayOfYear(_)
            | Expression::WeekOfYear(_)
            | Expression::Quarter(_)
            | Expression::AddMonths(_)
            | Expression::MonthsBetween(_)
            | Expression::LastDay(_)
            | Expression::NextDay(_)
            | Expression::Epoch(_)
            | Expression::EpochMs(_)
            | Expression::FromUnixtime(_)
            | Expression::UnixTimestamp(_)
            | Expression::MakeDate(_)
            | Expression::MakeTimestamp(_)
            | Expression::TimestampTrunc(_)
            | Expression::TimeStrToUnix(_)
            | Expression::SessionUser(_)
            | Expression::SHA(_)
            | Expression::SHA1Digest(_)
            | Expression::TimeToUnix(_)
            | Expression::ArrayFunc(_)
            | Expression::ArrayLength(_)
            | Expression::ArraySize(_)
            | Expression::Cardinality(_)
            | Expression::ArrayContains(_)
            | Expression::ArrayPosition(_)
            | Expression::ArrayAppend(_)
            | Expression::ArrayPrepend(_)
            | Expression::ArrayConcat(_)
            | Expression::ArraySort(_)
            | Expression::ArrayReverse(_)
            | Expression::ArrayDistinct(_)
            | Expression::ArrayJoin(_)
            | Expression::ArrayToString(_)
            | Expression::Unnest(_)
            | Expression::Explode(_)
            | Expression::ExplodeOuter(_)
            | Expression::ArrayFilter(_)
            | Expression::ArrayTransform(_)
            | Expression::ArrayFlatten(_)
            | Expression::ArrayCompact(_)
            | Expression::ArrayIntersect(_)
            | Expression::ArrayUnion(_)
            | Expression::ArrayExcept(_)
            | Expression::ArrayRemove(_)
            | Expression::ArrayZip(_)
            | Expression::Sequence(_)
            | Expression::Generate(_)
            | Expression::ExplodingGenerateSeries(_)
            | Expression::ToArray(_)
            | Expression::StarMap(_)
            | Expression::StructFunc(_)
            | Expression::StructExtract(_)
            | Expression::NamedStruct(_)
            | Expression::MapFunc(_)
            | Expression::MapFromEntries(_)
            | Expression::MapFromArrays(_)
            | Expression::MapKeys(_)
            | Expression::MapValues(_)
            | Expression::MapContainsKey(_)
            | Expression::MapConcat(_)
            | Expression::ElementAt(_)
            | Expression::TransformKeys(_)
            | Expression::TransformValues(_)
            | Expression::JsonExtract(_)
            | Expression::JsonExtractScalar(_)
            | Expression::JsonExtractPath(_)
            | Expression::JsonArray(_)
            | Expression::JsonObject(_)
            | Expression::JsonQuery(_)
            | Expression::JsonValue(_)
            | Expression::JsonArrayLength(_)
            | Expression::JsonKeys(_)
            | Expression::JsonType(_)
            | Expression::ParseJson(_)
            | Expression::ToJson(_)
            | Expression::JsonSet(_)
            | Expression::JsonInsert(_)
            | Expression::JsonRemove(_)
            | Expression::JsonMergePatch(_)
            | Expression::JsonArrayAgg(_)
            | Expression::JsonObjectAgg(_)
            | Expression::Convert(_)
            | Expression::Typeof(_)
            | Expression::Filter(_)
            | Expression::Anonymous(_)
            | Expression::AnonymousAggFunc(_)
            | Expression::CombinedAggFunc(_)
            | Expression::CombinedParameterizedAgg(_)
            | Expression::ParameterizedAgg(_)
            | Expression::UserDefinedFunction(_)
            | Expression::Transform(_)
            | Expression::Translate(_)
            | Expression::Grouping(_)
            | Expression::GroupingId(_)
            | Expression::HashAgg(_)
            | Expression::Hll(_)
            | Expression::ToBoolean(_)
            | Expression::ToMap(_)
            | Expression::Pad(_)
            | Expression::ToChar(_)
            | Expression::ToNumber(_)
            | Expression::ToDouble(_)
            | Expression::Int64(_)
            | Expression::StringFunc(_)
            | Expression::ToDecfloat(_)
            | Expression::TryToDecfloat(_)
            | Expression::ToFile(_)
            | Expression::ConvertToCharset(_)
            | Expression::ConvertTimezone(_)
            | Expression::GenerateSeries(_)
            | Expression::CosineDistance(_)
            | Expression::DotProduct(_)
            | Expression::EuclideanDistance(_)
            | Expression::ManhattanDistance(_)
            | Expression::JarowinklerSimilarity(_)
            | Expression::Booland(_)
            | Expression::Boolor(_)
            | Expression::ArgMax(_)
            | Expression::ArgMin(_)
            | Expression::ApproxTopK(_)
            | Expression::ApproxTopKAccumulate(_)
            | Expression::ApproxTopKCombine(_)
            | Expression::ApproxTopKEstimate(_)
            | Expression::ApproxTopSum(_)
            | Expression::ApproxQuantiles(_)
            | Expression::Minhash(_)
            | Expression::FarmFingerprint(_)
            | Expression::Float64(_)
            | Expression::Corr(_)
            | Expression::WidthBucket(_)
            | Expression::CovarSamp(_)
            | Expression::CovarPop(_)
            | Expression::ObjectAgg(_)
            | Expression::CastToStrType(_)
            | Expression::CheckJson(_)
            | Expression::CheckXml(_)
            | Expression::TranslateCharacters(_)
            | Expression::CurrentSchemas(_)
            | Expression::CurrentDatetime(_)
            | Expression::Localtime(_)
            | Expression::Localtimestamp(_)
            | Expression::Systimestamp(_)
            | Expression::CurrentSchema(_)
            | Expression::CurrentUser(_)
            | Expression::UtcTime(_)
            | Expression::UtcTimestamp(_)
            | Expression::Timestamp(_)
            | Expression::DateBin(_)
            | Expression::Datetime(_)
            | Expression::DatetimeAdd(_)
            | Expression::DatetimeSub(_)
            | Expression::DatetimeDiff(_)
            | Expression::DatetimeTrunc(_)
            | Expression::Dayname(_)
            | Expression::MakeInterval(_)
            | Expression::PreviousDay(_)
            | Expression::Elt(_)
            | Expression::TimestampAdd(_)
            | Expression::TimestampSub(_)
            | Expression::TimestampDiff(_)
            | Expression::TimeSlice(_)
            | Expression::TimeAdd(_)
            | Expression::TimeSub(_)
            | Expression::TimeDiff(_)
            | Expression::TimeTrunc(_)
            | Expression::DateFromParts(_)
            | Expression::TimeFromParts(_)
            | Expression::DecodeCase(_)
            | Expression::Decrypt(_)
            | Expression::DecryptRaw(_)
            | Expression::Encode(_)
            | Expression::Encrypt(_)
            | Expression::EncryptRaw(_)
            | Expression::EqualNull(_)
            | Expression::ToBinary(_)
            | Expression::Base64DecodeBinary(_)
            | Expression::Base64DecodeString(_)
            | Expression::Base64Encode(_)
            | Expression::TryBase64DecodeBinary(_)
            | Expression::TryBase64DecodeString(_)
            | Expression::GapFill(_)
            | Expression::GenerateDateArray(_)
            | Expression::GenerateTimestampArray(_)
            | Expression::GetExtract(_)
            | Expression::Getbit(_)
            | Expression::HexEncode(_)
            | Expression::Compress(_)
            | Expression::DecompressBinary(_)
            | Expression::DecompressString(_)
            | Expression::Xor(_)
            | Expression::Nullif(_)
            | Expression::Format(_)
            | Expression::MD5Digest(_)
            | Expression::MD5NumberLower64(_)
            | Expression::MD5NumberUpper64(_)
            | Expression::Monthname(_)
            | Expression::Ntile(_)
            | Expression::Normalize(_)
            | Expression::Normal(_)
            | Expression::Predict(_)
            | Expression::MLTranslate(_)
            | Expression::FeaturesAtTime(_)
            | Expression::GenerateEmbedding(_)
            | Expression::MLForecast(_)
            | Expression::ModelAttribute(_)
            | Expression::VectorSearch(_)
            | Expression::Quantile(_)
            | Expression::ApproxQuantile(_)
            | Expression::ApproxPercentileEstimate(_)
            | Expression::Randn(_)
            | Expression::Randstr(_)
            | Expression::RangeN(_)
            | Expression::RangeBucket(_)
            | Expression::ReadCSV(_)
            | Expression::ReadParquet(_)
            | Expression::Reduce(_)
            | Expression::RegexpExtractAll(_)
            | Expression::RegexpILike(_)
            | Expression::RegexpFullMatch(_)
            | Expression::RegexpInstr(_)
            | Expression::RegexpSplit(_)
            | Expression::RegexpCount(_)
            | Expression::RegrValx(_)
            | Expression::RegrValy(_)
            | Expression::RegrAvgy(_)
            | Expression::RegrAvgx(_)
            | Expression::RegrCount(_)
            | Expression::RegrIntercept(_)
            | Expression::RegrR2(_)
            | Expression::RegrSxx(_)
            | Expression::RegrSxy(_)
            | Expression::RegrSyy(_)
            | Expression::RegrSlope(_)
            | Expression::SafeAdd(_)
            | Expression::SafeDivide(_)
            | Expression::SafeMultiply(_)
            | Expression::SafeSubtract(_)
            | Expression::SHA2(_)
            | Expression::SHA2Digest(_)
            | Expression::SortArray(_)
            | Expression::SplitPart(_)
            | Expression::SubstringIndex(_)
            | Expression::StandardHash(_)
            | Expression::StrPosition(_)
            | Expression::Search(_)
            | Expression::SearchIp(_)
            | Expression::StrToDate(_)
            | Expression::DateStrToDate(_)
            | Expression::DateToDateStr(_)
            | Expression::StrToTime(_)
            | Expression::StrToUnix(_)
            | Expression::StrToMap(_)
            | Expression::NumberToStr(_)
            | Expression::FromBase(_)
            | Expression::Stuff(_)
            | Expression::TimeToStr(_)
            | Expression::TimeStrToTime(_)
            | Expression::TsOrDsAdd(_)
            | Expression::TsOrDsDiff(_)
            | Expression::TsOrDsToDate(_)
            | Expression::TsOrDsToTime(_)
            | Expression::Unhex(_)
            | Expression::Uniform(_)
            | Expression::UnixToStr(_)
            | Expression::UnixToTime(_)
            | Expression::Uuid(_)
            | Expression::TimestampFromParts(_)
            | Expression::TimestampTzFromParts(_)
            | Expression::Week(_)
            | Expression::AIAgg(_)
            | Expression::AIClassify(_)
            | Expression::ArrayAll(_)
            | Expression::ArrayAny(_)
            | Expression::ArrayConstructCompact(_)
            | Expression::StPoint(_)
            | Expression::StDistance(_)
            | Expression::StringToArray(_)
            | Expression::ArraySum(_)
            | Expression::JSONBContains(_)
            | Expression::JSONBExtract(_)
            | Expression::JSONBool(_)
            | Expression::JSONExtract(_)
            | Expression::JSONExtractQuote(_)
            | Expression::JSONExtractArray(_)
            | Expression::JSONExtractScalar(_)
            | Expression::JSONBExtractScalar(_)
            | Expression::JSONFormat(_)
            | Expression::JSONArrayAppend(_)
            | Expression::JSONArrayContains(_)
            | Expression::JSONArrayInsert(_)
            | Expression::ParseJSON(_)
            | Expression::ParseUrl(_)
            | Expression::ParseIp(_)
            | Expression::ParseTime(_)
            | Expression::ParseDatetime(_)
            | Expression::MapCat(_)
            | Expression::MapDelete(_)
            | Expression::MapInsert(_)
            | Expression::MapPick(_)
            | Expression::MatchAgainst(_)
            | Expression::ObjectInsert(_)
            | Expression::JSONBExists(_)
            | Expression::Apply(_)
            | Expression::Zipf(_)
            | Expression::XMLElement(_)
            | Expression::XMLGet(_)
            | Expression::Exists(_)
    )
}

/// Check if the expression subtree contains a Column reference.
fn has_column(expr: &Expression) -> bool {
    let mut found = false;
    deep_dfs(expr, &mut |e| {
        if matches!(e, Expression::Column(_)) {
            found = true;
        }
    });
    found
}

/// Find all function nodes in `expr` whose subtree contains a Column reference.
pub fn find_column_functions(expr: &Expression) -> Vec<&Expression> {
    let mut result = Vec::new();
    collect_column_functions(expr, &mut result);
    result
}

fn collect_column_functions<'a>(expr: &'a Expression, result: &mut Vec<&'a Expression>) {
    if is_any_function(expr) && has_column(expr) {
        result.push(expr);
    }
    match expr {
        Expression::And(op)
        | Expression::Or(op)
        | Expression::Eq(op)
        | Expression::Neq(op)
        | Expression::Lt(op)
        | Expression::Lte(op)
        | Expression::Gt(op)
        | Expression::Gte(op)
        | Expression::Add(op)
        | Expression::Sub(op)
        | Expression::Mul(op)
        | Expression::Div(op)
        | Expression::Mod(op)
        | Expression::BitwiseAnd(op)
        | Expression::BitwiseOr(op)
        | Expression::BitwiseXor(op)
        | Expression::Concat(op)
        | Expression::Is(op) => {
            collect_column_functions(&op.left, result);
            collect_column_functions(&op.right, result);
        }
        Expression::Not(u) | Expression::Neg(u) | Expression::BitwiseNot(u) => {
            collect_column_functions(&u.this, result);
        }
        Expression::Like(op) | Expression::ILike(op) => {
            collect_column_functions(&op.left, result);
            collect_column_functions(&op.right, result);
        }
        Expression::Between(b) => {
            collect_column_functions(&b.this, result);
            collect_column_functions(&b.low, result);
            collect_column_functions(&b.high, result);
        }
        Expression::In(i) => {
            collect_column_functions(&i.this, result);
            for e in &i.expressions {
                collect_column_functions(e, result);
            }
        }
        Expression::IsNull(i) => {
            collect_column_functions(&i.this, result);
        }
        Expression::Paren(p) => {
            collect_column_functions(&p.this, result);
        }
        Expression::Alias(a) => {
            collect_column_functions(&a.this, result);
        }
        Expression::Case(c) => {
            if let Some(ref op) = c.operand {
                collect_column_functions(op, result);
            }
            for w in &c.whens {
                collect_column_functions(&w.0, result);
                collect_column_functions(&w.1, result);
            }
            if let Some(ref else_) = c.else_ {
                collect_column_functions(else_, result);
            }
        }
        _ if is_any_function(expr) => {
            let children = get_function_children(expr);
            for child in children {
                collect_column_functions(child, result);
            }
        }
        _ => {}
    }
}

/// Get the direct child expressions of a function node for recursion.
fn get_function_children(expr: &Expression) -> Vec<&Expression> {
    match expr {
        Expression::Function(f) => f.args.iter().collect(),
        Expression::AggregateFunction(f) => f.args.iter().collect(),
        Expression::WindowFunction(wf) => vec![&wf.this],
        Expression::Upper(u)
        | Expression::Lower(u)
        | Expression::Length(u)
        | Expression::LTrim(u)
        | Expression::RTrim(u)
        | Expression::Reverse(u)
        | Expression::Abs(u)
        | Expression::Sqrt(u)
        | Expression::Cbrt(u)
        | Expression::Ln(u)
        | Expression::Exp(u)
        | Expression::Sign(u) => vec![&u.this],
        Expression::Cast(c) | Expression::TryCast(c) | Expression::SafeCast(c) => vec![&c.this],
        Expression::Coalesce(v) | Expression::Greatest(v) | Expression::Least(v) => {
            v.expressions.iter().collect()
        }
        Expression::NullIf(b)
        | Expression::IfNull(b)
        | Expression::Nvl(b)
        | Expression::Power(b)
        | Expression::Contains(b)
        | Expression::StartsWith(b)
        | Expression::EndsWith(b) => vec![&b.this, &b.expression],
        Expression::Count(c) => match &c.this {
            Some(this) => vec![this],
            None => vec![],
        },
        Expression::Sum(a)
        | Expression::Avg(a)
        | Expression::Min(a)
        | Expression::Max(a)
        | Expression::ArrayAgg(a)
        | Expression::CountIf(a)
        | Expression::Stddev(a)
        | Expression::StddevPop(a)
        | Expression::StddevSamp(a)
        | Expression::Variance(a) => {
            vec![&a.this]
        }
        Expression::IfFunc(i) => {
            let mut v = vec![&i.condition, &i.true_value];
            if let Some(ref fv) = i.false_value {
                v.push(fv);
            }
            v
        }
        Expression::ConcatWs(c) => {
            let mut v = vec![&c.separator];
            v.extend(c.expressions.iter());
            v
        }
        Expression::Trim(t) => vec![&t.this],
        Expression::Replace(r) => vec![&r.this, &r.old, &r.new],
        Expression::Substring(s) => vec![&s.this],
        Expression::Extract(e) => vec![&e.this],
        Expression::Nvl2(n) => vec![&n.this, &n.true_value, &n.false_value],
        _ => vec![],
    }
}
