qail-core 3.0.1

//! Extension traits for Expr.
//!

use super::json::JsonBuilder;
use crate::ast::{Expr, JsonPathSegment};

const INVALID_JSON_SOURCE_IDENT: &str = "__qail_invalid_json_source__";

fn json_source_or_invalid(expr: Expr, _method: &str) -> String {
    match expr {
        Expr::Named(name) => name,
        Expr::Aliased { name, .. } => name,
        Expr::JsonAccess { column, .. } => column,
        _other => {
            #[cfg(debug_assertions)]
            eprintln!(
                "QAIL: {}() expects a column-like expression, got {:?}; using fallback identifier {}",
                _method, _other, INVALID_JSON_SOURCE_IDENT
            );
            INVALID_JSON_SOURCE_IDENT.to_string()
        }
    }
}

/// Extension trait to add fluent methods to Expr
pub trait ExprExt {
    /// Add an alias to this expression.
    /// # Example
    /// ```ignore
    /// col("name").with_alias("user_name")
    /// ```
    fn with_alias(self, alias: &str) -> Expr;

    /// COALESCE with a default value.
    /// # Example
    /// ```ignore
    /// col("name").or_default("Unknown")  // COALESCE(name, 'Unknown')
    /// ```
    fn or_default(self, default: impl Into<Expr>) -> Expr;

    /// JSON text extraction (column->>'key').
    /// # Example
    /// ```ignore
    /// col("contact_info").json("phone")  // contact_info->>'phone'
    /// ```
    fn json(self, key: &str) -> JsonBuilder;

    /// JSON path extraction with dot notation. A segment that is an `i64`
    /// literal is an array position; build with [`JsonBuilder::get`] for an
    /// object key such as `"0"`.
    /// # Example
    /// ```ignore
    /// col("metadata").path("vessel.0.port")  // metadata->'vessel'->0->>'port'
    /// ```
    fn path(self, dotted_path: &str) -> JsonBuilder;

    /// Cast to a type: CAST(expr AS type)
    /// # Example
    /// ```ignore
    /// col("value").cast("int")  // CAST(value AS int)
    /// ```
    fn cast(self, target_type: &str) -> Expr;

    /// UPPER(expr)
    fn upper(self) -> Expr;

    /// LOWER(expr)
    fn lower(self) -> Expr;

    /// TRIM(expr)
    fn trim(self) -> Expr;

    /// LENGTH(expr)
    fn length(self) -> Expr;

    /// ABS(expr)
    fn abs(self) -> Expr;
}

impl ExprExt for Expr {
    fn with_alias(self, alias: &str) -> Expr {
        // Star, Aliased and Literal have no alias slot and come back as-is;
        // use `Expr::set_alias` to detect that case.
        let mut expr = self;
        expr.set_alias(alias);
        expr
    }

    fn or_default(self, default: impl Into<Expr>) -> Expr {
        Expr::FunctionCall {
            name: "COALESCE".to_string(),
            args: vec![self, default.into()],
            alias: None,
        }
    }

    fn json(self, key: &str) -> JsonBuilder {
        let column = json_source_or_invalid(self, "json");
        JsonBuilder {
            column,
            path_segments: vec![(key.into(), true)], // true = text extraction (->>)
            alias: None,
        }
    }

    fn path(self, dotted_path: &str) -> JsonBuilder {
        let column = json_source_or_invalid(self, "path");

        let segments: Vec<&str> = dotted_path.split('.').collect();
        let len = segments.len();
        let path_segments: Vec<(JsonPathSegment, bool)> = segments
            .into_iter()
            .enumerate()
            .map(|(i, segment)| (JsonPathSegment::from_path_text(segment), i == len - 1)) // Last segment as text
            .collect();

        JsonBuilder {
            column,
            path_segments,
            alias: None,
        }
    }

    fn cast(self, target_type: &str) -> Expr {
        Expr::Cast {
            expr: Box::new(self),
            target_type: target_type.to_string(),
            alias: None,
        }
    }

    fn upper(self) -> Expr {
        Expr::FunctionCall {
            name: "UPPER".to_string(),
            args: vec![self],
            alias: None,
        }
    }

    fn lower(self) -> Expr {
        Expr::FunctionCall {
            name: "LOWER".to_string(),
            args: vec![self],
            alias: None,
        }
    }

    fn trim(self) -> Expr {
        Expr::FunctionCall {
            name: "TRIM".to_string(),
            args: vec![self],
            alias: None,
        }
    }

    fn length(self) -> Expr {
        Expr::FunctionCall {
            name: "LENGTH".to_string(),
            args: vec![self],
            alias: None,
        }
    }

    fn abs(self) -> Expr {
        Expr::FunctionCall {
            name: "ABS".to_string(),
            args: vec![self],
            alias: None,
        }
    }
}

// Implement ExprExt for &str to enable: "col_name".or_default("X")
impl ExprExt for &str {
    fn with_alias(self, alias: &str) -> Expr {
        Expr::Aliased {
            name: self.to_string(),
            alias: alias.to_string(),
        }
    }

    fn or_default(self, default: impl Into<Expr>) -> Expr {
        Expr::FunctionCall {
            name: "COALESCE".to_string(),
            args: vec![Expr::Named(self.to_string()), default.into()],
            alias: None,
        }
    }

    fn json(self, key: &str) -> JsonBuilder {
        JsonBuilder {
            column: self.to_string(),
            path_segments: vec![(key.into(), true)],
            alias: None,
        }
    }

    fn path(self, dotted_path: &str) -> JsonBuilder {
        let segments: Vec<&str> = dotted_path.split('.').collect();
        let len = segments.len();
        let path_segments: Vec<(JsonPathSegment, bool)> = segments
            .into_iter()
            .enumerate()
            .map(|(i, segment)| (JsonPathSegment::from_path_text(segment), i == len - 1))
            .collect();

        JsonBuilder {
            column: self.to_string(),
            path_segments,
            alias: None,
        }
    }

    fn cast(self, target_type: &str) -> Expr {
        Expr::Cast {
            expr: Box::new(Expr::Named(self.to_string())),
            target_type: target_type.to_string(),
            alias: None,
        }
    }

    fn upper(self) -> Expr {
        Expr::FunctionCall {
            name: "UPPER".to_string(),
            args: vec![Expr::Named(self.to_string())],
            alias: None,
        }
    }

    fn lower(self) -> Expr {
        Expr::FunctionCall {
            name: "LOWER".to_string(),
            args: vec![Expr::Named(self.to_string())],
            alias: None,
        }
    }

    fn trim(self) -> Expr {
        Expr::FunctionCall {
            name: "TRIM".to_string(),
            args: vec![Expr::Named(self.to_string())],
            alias: None,
        }
    }

    fn length(self) -> Expr {
        Expr::FunctionCall {
            name: "LENGTH".to_string(),
            args: vec![Expr::Named(self.to_string())],
            alias: None,
        }
    }

    fn abs(self) -> Expr {
        Expr::FunctionCall {
            name: "ABS".to_string(),
            args: vec![Expr::Named(self.to_string())],
            alias: None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::builders::col;

    #[test]
    fn test_or_default() {
        let expr = col("name").or_default("Unknown");
        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "COALESCE"));
    }

    #[test]
    fn test_json_fluent() {
        let expr: Expr = col("info").json("phone").into();
        assert!(matches!(expr, Expr::JsonAccess { column, .. } if column == "info"));
    }

    #[test]
    fn test_path_fluent() {
        let expr: Expr = col("metadata").path("vessel.0.port").into();
        if let Expr::JsonAccess { path_segments, .. } = expr {
            assert_eq!(path_segments.len(), 3);
            assert_eq!(path_segments[0], ("vessel".into(), false)); // JSON
            assert_eq!(path_segments[1], (JsonPathSegment::Index(0), false)); // JSON
            assert_eq!(path_segments[2], ("port".into(), true)); // Text
        } else {
            panic!("Expected JsonAccess");
        }
    }

    #[test]
    fn test_str_or_default() {
        let expr = "name".or_default("N/A");
        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "COALESCE"));
    }

    #[test]
    fn test_cast_fluent() {
        let expr = col("value").cast("int4");
        assert!(matches!(expr, Expr::Cast { target_type, .. } if target_type == "int4"));
    }

    #[test]
    fn test_upper_fluent() {
        let expr = col("name").upper();
        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "UPPER"));
    }

    #[test]
    fn test_lower_fluent() {
        let expr = "email".lower();
        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "LOWER"));
    }

    #[test]
    fn test_trim_fluent() {
        let expr = col("text").trim();
        assert!(matches!(expr, Expr::FunctionCall { name, .. } if name == "TRIM"));
    }
}