// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! PEL Parser Library
//!
//! This module provides the function and types to parse a DataWeave expression into PEL.
//!
//! ## Primary types
//!
//! - [`InputContext`]: Input context containing inputs (e.g. payload, attributes) and DW functions (e.g. substringBefore, fromBase64) for compiling the expression.
//! - [`CompilationResult`]: Contains the compiled PEL expression if compile was successful or error messages upon failure.

use crate::dw::{
    parse_dw, unescape_string, AstNode, ComparisonType, ImportedElement, Location, LogicalType,
    Message, MessageKind,
};
use crate::json::escape_str;

/// Sets the context to parse the DataWeave expression. It receives:
/// * inputs that will be used (e.g. payload)
/// * functions that are part of the expression context (e.g. fromBase64)
#[derive(Clone, PartialEq, Debug, Default)]
pub struct InputContext<'a> {
    pub inputs: Vec<&'a str>,
    pub functions: Vec<&'a str>,
}

#[derive(Clone, PartialEq, Debug, Default)]
struct DeclaredImports {
    pub imports: Vec<ImportDefinition>,
}

#[derive(Clone, PartialEq, Debug, Default)]
struct ImportDefinition {
    module: String,
    alias: Option<String>,
    imported_elements: Vec<ImportElement>,
}

#[derive(Clone, PartialEq, Debug, Default)]
struct ImportElement {
    element: String,
    alias: Option<String>,
}

/// Contains the actual compilation result:
/// * pel: expression as a String
/// * messages: vector of [Message] containing any potential error messages.
#[derive(Clone, PartialEq, Debug)]
pub struct CompilationResult {
    pub pel: Option<String>,
    pub messages: Vec<Message>,
}

#[derive(Clone, PartialEq, Debug, Default)]
pub struct DeclaredNamespaces {
    pub namespaces: Vec<NamespaceInformation>,
}

impl DeclaredNamespaces {
    pub fn contains_prefix(&self, prefix: &String) -> bool {
        self.namespaces.iter().any(|n| n.prefix.eq(prefix))
    }

    pub fn resolve_prefix(&self, prefix: &String) -> Option<&String> {
        self.namespaces
            .iter()
            .find(|n| n.prefix.eq(prefix))
            .map(|ns| &ns.uri)
    }
}

#[derive(Clone, PartialEq, Debug, Default)]
pub struct NamespaceInformation {
    pub prefix: String,
    pub uri: String,
}

fn is_imported(var_name: &str, ctx: &InputContext, imports: &DeclaredImports) -> bool {
    //Check dw::Core as it is allways imported
    let mut fqn: String = String::new();
    fqn.push_str("dw::Core");
    fqn.push_str("::");
    fqn.push_str(var_name);
    if ctx.functions.contains(&fqn.as_str()) {
        return true;
    }

    for import in &imports.imports {
        for element in &import.imported_elements {
            if element.element == "*" {
                let mut fqn: String = String::new();
                fqn.push_str(import.module.as_str());
                fqn.push_str("::");
                fqn.push_str(var_name);
                if ctx.functions.contains(&fqn.as_str()) {
                    return true;
                }
            } else if element.element == var_name || element.alias.as_deref() == Some(var_name) {
                let mut fqn: String = String::new();
                fqn.push_str(import.module.as_str());
                fqn.push_str("::");
                fqn.push_str(element.element.as_str());
                if ctx.functions.contains(&fqn.as_str()) {
                    return true;
                }
            }
        }
    }
    false
}

fn validate_pel(
    expression: &AstNode,
    ctx: &InputContext,
    imports: &DeclaredImports,
    namespace: &DeclaredNamespaces,
) -> Vec<Message> {
    match expression {
        AstNode::Name {
            prefix, location, ..
        } => {
            let maybe_prefix: &Option<AstNode> = prefix;
            match maybe_prefix {
                None => {
                    vec![]
                }
                Some(prefix_node) => {
                    let prefix_string = prefix_node.get_value().unwrap();
                    if !namespace.contains_prefix(&prefix_string) {
                        vec![Message {
                            message: format!(
                                "Unable to resolve namespace prefix of : `{}`",
                                &prefix_string
                            ),
                            kind: MessageKind::Error,
                            location: Location {
                                start: location.start,
                                end: location.end,
                            },
                        }]
                    } else {
                        vec![]
                    }
                }
            }
        }
        AstNode::VariableReference { reference, .. } => match reference.as_ref() {
            AstNode::NameIdentifier { name, location } => {
                if name.len() == 1 {
                    let var_name = name.first().unwrap();
                    if !ctx.inputs.contains(&var_name.as_str()) {
                        if !is_imported(var_name, ctx, imports) {
                            vec![Message {
                                message: format!("Unable to resolve reference of: `{var_name}`."),
                                kind: MessageKind::Error,
                                location: Location {
                                    start: location.start,
                                    end: location.end,
                                },
                            }]
                        } else {
                            vec![]
                        }
                    } else {
                        vec![]
                    }
                } else {
                    let fqn = name.join("::");
                    if !ctx.functions.contains(&fqn.as_str()) {
                        vec![Message {
                            message: format!("Unable to resolve module with identifier: `{fqn}`."),
                            kind: MessageKind::Error,
                            location: Location {
                                start: location.start,
                                end: location.end,
                            },
                        }]
                    } else {
                        vec![]
                    }
                }
            }
            _ => {
                vec![]
            }
        },

        AstNode::Array { items, .. } => items
            .iter()
            .flat_map(|node| validate_pel(node, ctx, imports, namespace))
            .collect(),
        AstNode::Comparison { lhs, rhs, .. } => [lhs, rhs]
            .iter()
            .flat_map(|node| validate_pel(node.as_ref(), ctx, imports, namespace))
            .collect(),
        AstNode::Logical { lhs, rhs, .. } => [lhs, rhs]
            .iter()
            .flat_map(|node| validate_pel(node.as_ref(), ctx, imports, namespace))
            .collect(),
        AstNode::Not { rhs, .. } => validate_pel(rhs, ctx, imports, namespace),
        AstNode::ValueSelector { lhs, rhs, .. } => [lhs, rhs]
            .iter()
            .flat_map(|node| validate_pel(node.as_ref(), ctx, imports, namespace))
            .collect(),
        AstNode::InfixFunctionCall { lhs, func, rhs, .. } => [lhs, func, rhs]
            .iter()
            .flat_map(|node| validate_pel(node.as_ref(), ctx, imports, namespace))
            .collect(),
        AstNode::FunctionCall { lhs, args, .. } => {
            let mut vec1 = args.clone();
            vec1.push(lhs.as_ref().to_owned());
            vec1.iter()
                .flat_map(|node| validate_pel(node, ctx, imports, namespace))
                .collect()
        }
        AstNode::IfElse {
            condition,
            if_branch,
            else_branch,
            ..
        } => [condition, if_branch, else_branch]
            .iter()
            .flat_map(|node| validate_pel(node.as_ref(), ctx, imports, namespace))
            .collect(),
        AstNode::Default { lhs, rhs, .. } => [lhs, rhs]
            .iter()
            .flat_map(|node| validate_pel(node.as_ref(), ctx, imports, namespace))
            .collect(),
        AstNode::Document {
            header: _header,
            body,
            ..
        } => [body]
            .iter()
            .flat_map(|node| validate_pel(node.as_ref(), ctx, imports, namespace))
            .collect(),
        AstNode::Number {
            value, location, ..
        } => match value.parse::<f64>() {
            Ok(_) => {
                vec![]
            }
            Err(_) => {
                vec![Message {
                    message: format!("Number `{value}` should be a valid Double."),
                    kind: MessageKind::Error,
                    location: Location {
                        start: location.start,
                        end: location.end,
                    },
                }]
            }
        },
        AstNode::Object { location, .. } => {
            vec![Message {
                message: "Object expressions are not yet supported.".to_string(),
                kind: MessageKind::Error,
                location: Location {
                    start: location.start,
                    end: location.end,
                },
            }]
        }
        AstNode::NullSafeSelector { lhs, .. } => {
            validate_pel(lhs.as_ref(), ctx, imports, namespace)
        }
        n => n
            .children()
            .iter()
            .flat_map(|node| validate_pel(node, ctx, imports, namespace))
            .collect(),
    }
}

fn collect_declared_ns(expression: &AstNode) -> DeclaredNamespaces {
    match expression {
        AstNode::Document { header: None, .. } => DeclaredNamespaces::default(),
        AstNode::Document {
            header: Some(directives),
            ..
        } => directives
            .iter()
            .map(collect_declared_ns)
            .reduce(|acc, e| {
                let x: Vec<NamespaceInformation> =
                    [acc.namespaces.as_slice(), e.namespaces.as_slice()].concat();
                DeclaredNamespaces { namespaces: x }
            })
            .unwrap_or_default(),
        AstNode::NamespaceDirective { prefix, uri, .. } => DeclaredNamespaces {
            namespaces: vec![NamespaceInformation {
                prefix: prefix.get_value().unwrap(),
                uri: uri.get_value().unwrap(),
            }],
        },
        _ => DeclaredNamespaces::default(),
    }
}
fn collect_declared_imports(expression: &AstNode) -> DeclaredImports {
    match expression {
        AstNode::Document { header: None, .. } => DeclaredImports::default(),
        AstNode::Document {
            header: Some(directives),
            ..
        } => directives
            .iter()
            .map(collect_declared_imports)
            .reduce(|acc, e| {
                let x: Vec<ImportDefinition> =
                    [acc.imports.as_slice(), e.imports.as_slice()].concat();
                DeclaredImports { imports: x }
            })
            .unwrap_or_default(),
        AstNode::ImportDirective {
            module, elements, ..
        } => DeclaredImports {
            imports: vec![ImportDefinition {
                module: name_as_string(module.as_ref()),
                alias: alias_as_string(&module.alias),
                imported_elements: elements
                    .iter()
                    .map(|e| ImportElement {
                        element: name_as_string(e),
                        alias: alias_as_string(&e.alias),
                    })
                    .collect(),
            }],
        },
        _ => DeclaredImports::default(),
    }
}

fn name_as_string(module: &ImportedElement) -> String {
    match module.name.as_ref() {
        AstNode::NameIdentifier { name, .. } => name.join("::"),
        _ => "".to_string(),
    }
}

fn alias_as_string(option: &Option<AstNode>) -> Option<String> {
    option.as_ref().map(|a| match a {
        AstNode::NameIdentifier { name, .. } => name.join("::"),
        _ => "".to_string(),
    })
}

/// Compiles the provided DataWeave expression and returns a [CompilationResult].
pub fn compile_to_pel_expr(
    script_name: &str,
    expression_str: &str,
    ctx: &InputContext,
) -> CompilationResult {
    let result = parse_dw(script_name, expression_str);
    match result {
        Ok(parse_result) => {
            let expression = parse_result.ast;
            let declared_imports: DeclaredImports = collect_declared_imports(&expression);
            let declared_ns: DeclaredNamespaces = collect_declared_ns(&expression);
            let messages = validate_pel(&expression, ctx, &declared_imports, &declared_ns);
            let errors = messages.iter().any(|m| m.kind == MessageKind::Error);
            CompilationResult {
                pel: (!errors).then(|| to_pel(&expression, &declared_ns)),
                messages,
            }
        }
        Err(e) => CompilationResult {
            pel: None,
            messages: e.messages,
        },
    }
}

fn to_pel(expression: &AstNode, declared_namespaces: &DeclaredNamespaces) -> String {
    match expression {
        AstNode::Str {
            value: content,
            quote,
            location,
        } => {
            format!(
                "[{}, {}, {}]",
                str(":str"),
                loc_str(location),
                str(unescape_string(content.to_string(), quote).as_str())
            )
        }
        AstNode::Boolean { value, location } => {
            format!(
                "[{}, {}, {}]",
                str(":bool"),
                loc_str(location),
                str(value.to_string().as_str())
            )
        }
        AstNode::Number { value, location } => {
            format!(
                "[{}, {}, {}]",
                str(":nbr"),
                loc_str(location),
                str(value.to_string().as_str())
            )
        }
        AstNode::Null { location } => {
            format!("[{}, {}]", str(":null"), loc_str(location))
        }
        AstNode::DateTime { value, location } => {
            format!(
                "[{}, {}, {}]",
                str(":datetime"),
                loc_str(location),
                str(value.as_str())
            )
        }
        AstNode::LocalDateTime { value, location } => {
            format!(
                "[{}, {}, {}]",
                str(":ldatetime"),
                loc_str(location),
                str(value.as_str())
            )
        }
        AstNode::Array { items, location } => {
            if items.is_empty() {
                format!("[{}, {}]", str(":array"), loc_str(location))
            } else {
                let elements: Vec<String> = items
                    .iter()
                    .map(|a| to_pel(a, declared_namespaces))
                    .collect();
                let array_elements = elements.join(", ");
                format!(
                    "[{}, {}, {}]",
                    str(":array"),
                    loc_str(location),
                    array_elements
                )
            }
        }
        AstNode::Comparison {
            lhs,
            rhs,
            comp,
            location,
        } => {
            let c = match comp {
                ComparisonType::Greater => ">",
                ComparisonType::Less => "<",
                ComparisonType::Equal => "==",
                ComparisonType::NotEqual => "!=",
                ComparisonType::Similar => "~=",
                ComparisonType::GreaterEqual => ">=",
                _ => "<=",
            };
            format!(
                "[{}, {}, {}, {}]",
                str(c),
                loc_str(location),
                to_pel(lhs.as_ref(), declared_namespaces),
                to_pel(rhs.as_ref(), declared_namespaces)
            )
        }
        AstNode::Logical {
            lhs,
            rhs,
            comp,
            location,
        } => {
            let c = match comp {
                LogicalType::And => "&&",
                LogicalType::Or => "||",
            };
            format!(
                "[{}, {}, {}, {}]",
                str(c),
                loc_str(location),
                to_pel(lhs.as_ref(), declared_namespaces),
                to_pel(rhs.as_ref(), declared_namespaces)
            )
        }
        AstNode::Not { rhs: exp, location } => {
            format!(
                "[{}, {}, {}]",
                str("!"),
                loc_str(location),
                to_pel(exp.as_ref(), declared_namespaces)
            )
        }
        AstNode::ValueSelector { lhs, rhs, location } => {
            format!(
                "[{}, {}, {}, {}]",
                str("."),
                loc_str(location),
                to_pel(lhs.as_ref(), declared_namespaces),
                to_pel(rhs.as_ref(), declared_namespaces)
            )
        }
        AstNode::AttributeSelector { lhs, rhs, location } => {
            format!(
                "[{}, {}, {}, {}]",
                str("@"),
                loc_str(location),
                to_pel(lhs.as_ref(), declared_namespaces),
                to_pel(rhs.as_ref(), declared_namespaces)
            )
        }
        AstNode::MultiValueSelector { lhs, rhs, location } => {
            format!(
                "[{}, {}, {}, {}]",
                str("*"),
                loc_str(location),
                to_pel(lhs.as_ref(), declared_namespaces),
                to_pel(rhs.as_ref(), declared_namespaces)
            )
        }
        AstNode::Name {
            prefix,
            value,
            location,
        } => {
            let maybe_prefix: &Option<AstNode> = prefix;
            match maybe_prefix {
                None => {
                    //For backwards compatibility
                    to_pel(value, declared_namespaces)
                }
                Some(prefix) => {
                    let prefix_str = &prefix.get_value().unwrap();
                    let namespace = &declared_namespaces.resolve_prefix(prefix_str).unwrap();
                    let ns_node = format!(
                        "[{}, {}, {}]",
                        str(":ns"),
                        loc_str(&prefix.location()),
                        str(namespace)
                    );
                    format!(
                        "[{}, {}, {}, {}]",
                        str(":name"),
                        loc_str(location),
                        ns_node,
                        to_pel(value.as_ref(), declared_namespaces)
                    )
                }
            }
        }
        AstNode::NameIdentifier {
            name,
            location: _location,
        } => {
            // TODO We need to execute the mapping between
            str(name.last().unwrap().as_str())
        }
        AstNode::VariableReference {
            reference,
            location,
        } => {
            format!(
                "[{}, {}, {}]",
                str(":ref"),
                loc_str(location),
                to_pel(reference.as_ref(), declared_namespaces)
            )
        }
        AstNode::InfixFunctionCall {
            lhs,
            func,
            rhs,
            location,
        } => {
            format!(
                "[{}, {}, {}, {}, {}]",
                str(":apply"),
                loc_str(location),
                to_pel(func.as_ref(), declared_namespaces),
                to_pel(lhs.as_ref(), declared_namespaces),
                to_pel(rhs.as_ref(), declared_namespaces)
            )
        }
        AstNode::FunctionCall {
            lhs,
            args,
            location,
        } => {
            if args.is_empty() {
                format!(
                    "[{}, {}, {}]",
                    str(":apply"),
                    loc_str(location),
                    to_pel(lhs, declared_namespaces)
                )
            } else {
                let elements: Vec<String> = args
                    .iter()
                    .map(|n| to_pel(n, declared_namespaces))
                    .collect();
                let array_elements = elements.join(", ");
                format!(
                    "[{}, {}, {}, {}]",
                    str(":apply"),
                    loc_str(location),
                    to_pel(lhs, declared_namespaces),
                    array_elements
                )
            }
        }
        AstNode::IfElse {
            condition,
            if_branch,
            else_branch,
            location,
        } => {
            format!(
                "[{}, {}, {}, {}, {}]",
                str(":if"),
                loc_str(location),
                to_pel(condition, declared_namespaces),
                to_pel(if_branch, declared_namespaces),
                to_pel(else_branch, declared_namespaces)
            )
        }
        AstNode::Default { lhs, rhs, location } => {
            format!(
                "[{}, {}, {}, {}]",
                str(":default"),
                loc_str(location),
                to_pel(lhs, declared_namespaces),
                to_pel(rhs, declared_namespaces)
            )
        }
        AstNode::Period { value, location } => {
            format!(
                "[{}, {}, {}]",
                str(":period"),
                loc_str(location),
                str(value)
            )
        }
        AstNode::Time { value, location } => {
            format!("[{}, {}, {}]", str(":time"), loc_str(location), str(value))
        }
        AstNode::LocalTime { value, location } => {
            format!("[{}, {}, {}]", str(":ltime"), loc_str(location), str(value))
        }
        AstNode::Date { value, location } => {
            format!("[{}, {}, {}]", str(":ldate"), loc_str(location), str(value))
        }
        AstNode::NullSafeSelector { lhs, .. } => to_pel(lhs.as_ref(), declared_namespaces),
        AstNode::Document { body, .. } => to_pel(body.as_ref(), declared_namespaces),
        AstNode::StrInterpolation {
            segments, location, ..
        } => segments
            .iter()
            .map(|n| to_pel(n, declared_namespaces))
            .reduce(|acc, value| {
                format!(
                    "[{}, {}, {}, {}]",
                    str(":++"),
                    loc_str(location),
                    acc,
                    value
                )
            })
            .unwrap(),
        _ => "".to_string(),
    }
}

fn loc_str(location: &Location) -> String {
    str(format!("{}-{}", location.start, location.end).as_str())
}

fn str(text: &str) -> String {
    escape_str(text)
}
