//! AST visitor that extracts fields, functions, operators, and literals
//! from a parsed wirefilter expression.

use std::collections::HashSet;
use std::fmt::Display;
use std::net::IpAddr;
use std::ops::RangeInclusive;

use wirefilter::{
    ComparisonExpr, ComparisonOpExpr, ExplicitIpRange, FunctionCallArgExpr, IdentifierExpr,
    IndexExpr, IntOp, IntRange, IpRange, LogicalExpr, LogicalOp, OrderingOp, RhsValue, RhsValues,
    UnaryOp,
};

/// Maximum nesting depth before the extractor stops descending.
const MAX_NESTING_DEPTH: usize = 100;

/// Generate a dedup-add method that skips duplicates with zero allocations
/// (via `HashSet::contains(&str)` on a `HashSet<String>`), and allocates
/// once per unique entry (for the set) then clones (for the vec).
macro_rules! dedup_add {
    ($method:ident, $seen:ident, $list:ident) => {
        fn $method(&mut self, value: &str) {
            if !self.$seen.contains(value) {
                let owned = value.to_owned();
                self.$seen.insert(owned.clone());
                self.$list.push(owned);
            }
        }
    };
}

/// Extracts all components from a wirefilter AST for the ExpressionInfo contract.
pub struct ExpressionExtractor {
    pub fields: Vec<String>,
    pub functions: Vec<String>,
    pub operators: Vec<String>,
    pub string_literals: Vec<String>,
    pub regex_literals: Vec<String>,
    pub ip_literals: Vec<String>,
    pub int_literals: Vec<i64>,

    seen_fields: HashSet<String>,
    seen_functions: HashSet<String>,
    seen_operators: HashSet<String>,
    seen_strings: HashSet<String>,
    seen_regexes: HashSet<String>,
    seen_ips: HashSet<String>,
    seen_ints: HashSet<i64>,

    depth: usize,
    max_depth_exceeded: bool,
}

impl ExpressionExtractor {
    pub fn new() -> Self {
        Self {
            fields: Vec::new(),
            functions: Vec::new(),
            operators: Vec::new(),
            string_literals: Vec::new(),
            regex_literals: Vec::new(),
            ip_literals: Vec::new(),
            int_literals: Vec::new(),
            seen_fields: HashSet::new(),
            seen_functions: HashSet::new(),
            seen_operators: HashSet::new(),
            seen_strings: HashSet::new(),
            seen_regexes: HashSet::new(),
            seen_ips: HashSet::new(),
            seen_ints: HashSet::new(),
            depth: 0,
            max_depth_exceeded: false,
        }
    }

    /// Returns true if the extractor hit the maximum nesting depth.
    pub fn depth_exceeded(&self) -> bool {
        self.max_depth_exceeded
    }

    dedup_add!(add_field, seen_fields, fields);
    dedup_add!(add_function, seen_functions, functions);
    dedup_add!(add_operator, seen_operators, operators);
    dedup_add!(add_string, seen_strings, string_literals);
    dedup_add!(add_regex, seen_regexes, regex_literals);
    dedup_add!(add_ip, seen_ips, ip_literals);

    fn add_int(&mut self, val: i64) {
        if self.seen_ints.insert(val) {
            self.int_literals.push(val);
        }
    }

    /// Extract an IP address as a string.
    fn extract_ip(&mut self, addr: &IpAddr) {
        self.add_ip(&addr.to_string());
    }

    /// Extract RHS value literals.
    fn extract_rhs_value(&mut self, rhs: &RhsValue) {
        match rhs {
            RhsValue::Bytes(bytes_expr) => {
                self.add_string(&String::from_utf8_lossy(bytes_expr));
            }
            RhsValue::Int(i) => self.add_int(*i),
            RhsValue::Ip(addr) => self.extract_ip(addr),
            RhsValue::Bool(_) | RhsValue::Array(_) | RhsValue::Map(_) => {}
        }
    }

    /// Extract int from an IntRange (single value or range endpoints).
    fn extract_int_range(&mut self, range: &IntRange) {
        let r: RangeInclusive<i64> = range.into();
        let start = *r.start();
        let end = *r.end();
        self.add_int(start);
        if end != start {
            self.add_int(end);
        }
    }

    /// Extract start (and end if different) from an explicit IP range.
    fn extract_explicit_ip_range<T: Display + PartialEq>(&mut self, start: &T, end: &T) {
        self.add_ip(&start.to_string());
        if start != end {
            self.add_ip(&end.to_string());
        }
    }

    /// Extract IP from an IpRange.
    fn extract_ip_range(&mut self, range: &IpRange) {
        match range {
            IpRange::Cidr(cidr) => {
                self.add_ip(&cidr.to_string());
            }
            IpRange::Explicit(explicit) => match explicit {
                ExplicitIpRange::V4(r) => self.extract_explicit_ip_range(r.start(), r.end()),
                ExplicitIpRange::V6(r) => self.extract_explicit_ip_range(r.start(), r.end()),
            },
        }
    }

    /// Extract literals from an RhsValues set (used in `in {...}` expressions).
    fn extract_rhs_values(&mut self, values: &RhsValues) {
        match values {
            RhsValues::Bytes(items) => {
                for item in items {
                    self.add_string(&String::from_utf8_lossy(item));
                }
            }
            RhsValues::Int(ranges) => {
                for range in ranges {
                    self.extract_int_range(range);
                }
            }
            RhsValues::Ip(ranges) => {
                for range in ranges {
                    self.extract_ip_range(range);
                }
            }
            RhsValues::Bool(_) | RhsValues::Array(_) | RhsValues::Map(_) => {}
        }
    }

    /// Walk a comparison operator and extract the operator name + RHS literals.
    fn extract_comparison_op(&mut self, op: &ComparisonOpExpr) {
        match op {
            ComparisonOpExpr::IsTrue => {
                // Bare field truthiness — no explicit operator.
            }
            ComparisonOpExpr::Ordering { op: ord_op, rhs } => {
                self.add_operator(match ord_op {
                    OrderingOp::Equal => "eq",
                    OrderingOp::NotEqual => "ne",
                    OrderingOp::GreaterThan => "gt",
                    OrderingOp::GreaterThanEqual => "ge",
                    OrderingOp::LessThan => "lt",
                    OrderingOp::LessThanEqual => "le",
                });
                self.extract_rhs_value(rhs);
            }
            ComparisonOpExpr::Int { op, rhs } => {
                self.add_operator(match op {
                    IntOp::BitwiseAnd => "bitwise_and",
                });
                self.add_int(*rhs);
            }
            ComparisonOpExpr::Contains(bytes_expr) => {
                self.add_operator("contains");
                self.add_string(&String::from_utf8_lossy(bytes_expr));
            }
            ComparisonOpExpr::Matches(regex) => {
                self.add_operator("matches");
                self.add_regex(&regex.to_string());
            }
            ComparisonOpExpr::Wildcard(wc) => {
                self.add_operator("wildcard");
                self.add_string(&String::from_utf8_lossy(wc.pattern()));
            }
            ComparisonOpExpr::StrictWildcard(wc) => {
                self.add_operator("strict_wildcard");
                self.add_string(&String::from_utf8_lossy(wc.pattern()));
            }
            ComparisonOpExpr::OneOf(values) => {
                self.add_operator("in");
                self.extract_rhs_values(values);
            }
            ComparisonOpExpr::ContainsOneOf(items) => {
                self.add_operator("contains");
                self.add_operator("in");
                for item in items {
                    self.add_string(&String::from_utf8_lossy(item));
                }
            }
            ComparisonOpExpr::InList { .. } => {
                self.add_operator("in");
            }
        }
    }

    /// Recursively walk an IndexExpr to extract fields and functions.
    fn walk_index_expr(&mut self, index_expr: &IndexExpr) {
        match index_expr.identifier() {
            IdentifierExpr::Field(field) => {
                self.add_field(field.name());
            }
            IdentifierExpr::FunctionCallExpr(func_call) => {
                self.add_function(func_call.function().name());
                for arg in func_call.args() {
                    self.walk_function_arg(arg);
                }
            }
        }
    }

    /// Walk a function call argument.
    fn walk_function_arg(&mut self, arg: &FunctionCallArgExpr) {
        if self.depth >= MAX_NESTING_DEPTH {
            self.max_depth_exceeded = true;
            return;
        }
        self.depth += 1;
        match arg {
            FunctionCallArgExpr::IndexExpr(index_expr) => {
                self.walk_index_expr(index_expr);
            }
            FunctionCallArgExpr::Literal(rhs) => {
                self.extract_rhs_value(rhs);
            }
            FunctionCallArgExpr::Logical(logical) => {
                self.walk_logical(logical);
            }
        }
        self.depth -= 1;
    }

    /// Recursively walk a logical expression tree.
    fn walk_logical(&mut self, expr: &LogicalExpr) {
        if self.depth >= MAX_NESTING_DEPTH {
            self.max_depth_exceeded = true;
            return;
        }
        self.depth += 1;
        match expr {
            LogicalExpr::Combining { op, items } => {
                self.add_operator(match op {
                    LogicalOp::And => "and",
                    LogicalOp::Or => "or",
                    LogicalOp::Xor => "xor",
                });
                for item in items {
                    self.walk_logical(item);
                }
            }
            LogicalExpr::Comparison(cmp) => {
                self.walk_comparison(cmp);
            }
            LogicalExpr::Parenthesized(paren) => {
                self.walk_logical(&paren.expr);
            }
            LogicalExpr::Unary { op, arg } => {
                match op {
                    UnaryOp::Not => self.add_operator("not"),
                }
                self.walk_logical(arg);
            }
        }
        self.depth -= 1;
    }

    /// Walk a comparison expression: LHS + operator + RHS literals.
    fn walk_comparison(&mut self, cmp: &ComparisonExpr) {
        self.walk_index_expr(cmp.lhs_expr());
        self.extract_comparison_op(cmp.operator());
    }

    /// Entry point: walk the root LogicalExpr of a FilterAst.
    pub fn extract(&mut self, root: &LogicalExpr) {
        self.walk_logical(root);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::scheme::SCHEME;
    use std::sync::LazyLock;
    use wirefilter::{AlwaysList, Scheme, SchemeBuilder, Type};

    /// Helper: parse an expression against the production SCHEME and extract.
    fn extract_expr(expr: &str) -> ExpressionExtractor {
        let ast = SCHEME.parse(expr).expect("parse failed");
        let mut ex = ExpressionExtractor::new();
        ex.extract(ast.expression());
        ex
    }

    /// A minimal scheme with an Int list registered so `in $list_name`
    /// syntax can be parsed.
    static LIST_SCHEME: LazyLock<Scheme> = LazyLock::new(|| {
        let mut b = SchemeBuilder::new();
        b.add_field("ip.src", Type::Ip).unwrap();
        b.add_field("http.host", Type::Bytes).unwrap();
        b.add_field("cf.bot_management.score", Type::Int).unwrap();
        b.add_list(Type::Int, AlwaysList {}).unwrap();
        b.add_list(Type::Ip, AlwaysList {}).unwrap();
        b.add_list(Type::Bytes, AlwaysList {}).unwrap();
        b.build()
    });

    /// Helper: parse against the LIST_SCHEME and extract.
    fn extract_list_expr(expr: &str) -> ExpressionExtractor {
        let ast = LIST_SCHEME.parse(expr).expect("parse failed");
        let mut ex = ExpressionExtractor::new();
        ex.extract(ast.expression());
        ex
    }

    // ── Deduplication ────────────────────────────────────────

    #[test]
    fn dedup_field_referenced_twice() {
        let ex = extract_expr(r#"http.host eq "a.com" or http.host eq "b.com""#);
        assert_eq!(ex.fields, vec!["http.host"]);
    }

    #[test]
    fn dedup_operator_used_twice() {
        let ex = extract_expr(r#"http.host eq "a.com" or http.host eq "b.com""#);
        // "eq" should appear once, "or" should appear once
        assert_eq!(ex.operators.iter().filter(|o| *o == "eq").count(), 1);
        assert_eq!(ex.operators.iter().filter(|o| *o == "or").count(), 1);
    }

    #[test]
    fn dedup_string_literal_used_twice() {
        let ex = extract_expr(r#"http.host eq "same" or http.request.uri.path eq "same""#);
        assert_eq!(
            ex.string_literals.iter().filter(|s| *s == "same").count(),
            1
        );
    }

    #[test]
    fn dedup_int_literal_used_twice() {
        let ex = extract_expr("cf.bot_management.score eq 30 or cf.threat_score eq 30");
        assert_eq!(ex.int_literals.iter().filter(|&&v| v == 30).count(), 1);
    }

    #[test]
    fn dedup_ip_literal_used_twice() {
        let ex = extract_expr(r#"ip.src eq 1.2.3.4 or ip.src ne 1.2.3.4"#);
        assert_eq!(ex.ip_literals.iter().filter(|s| *s == "1.2.3.4").count(), 1);
    }

    // ── Bool / Array / Map no-ops ────────────────────────────

    #[test]
    fn bool_field_truthiness_no_crash() {
        // Bare bool field — IsTrue operator, no crash
        let ex = extract_expr("ssl");
        assert_eq!(ex.fields, vec!["ssl"]);
        assert!(ex.operators.is_empty());
        assert!(ex.string_literals.is_empty());
    }

    #[test]
    fn bool_field_negated() {
        let ex = extract_expr("not ssl");
        assert_eq!(ex.fields, vec!["ssl"]);
        assert!(ex.operators.contains(&"not".to_string()));
    }

    // ── InList ($list_name) ──────────────────────────────────

    #[test]
    fn in_list_extracts_in_operator() {
        let ex = extract_list_expr("cf.bot_management.score in $my_list");
        assert_eq!(ex.fields, vec!["cf.bot_management.score"]);
        assert!(ex.operators.contains(&"in".to_string()));
        // No literals should be extracted from a named list
        assert!(ex.int_literals.is_empty());
        assert!(ex.string_literals.is_empty());
    }

    #[test]
    fn in_list_ip_extracts_in_operator() {
        let ex = extract_list_expr("ip.src in $blocked_ips");
        assert_eq!(ex.fields, vec!["ip.src"]);
        assert!(ex.operators.contains(&"in".to_string()));
        assert!(ex.ip_literals.is_empty());
    }

    // ── Explicit IP range extraction ─────────────────────────

    #[test]
    fn explicit_ip_v4_range_extracts_both_endpoints() {
        let ex = extract_expr("ip.src in {10.0.0.1..10.0.0.5}");
        assert!(ex.ip_literals.contains(&"10.0.0.1".to_string()));
        assert!(ex.ip_literals.contains(&"10.0.0.5".to_string()));
    }

    #[test]
    fn explicit_ip_v4_single_value_range() {
        // A range where start == end should only produce one entry
        let ex = extract_expr("ip.src in {10.0.0.1..10.0.0.1}");
        assert_eq!(
            ex.ip_literals.iter().filter(|s| *s == "10.0.0.1").count(),
            1
        );
    }

    #[test]
    fn cidr_ip_range_extracted() {
        let ex = extract_expr("ip.src in {192.168.0.0/16}");
        assert!(ex.ip_literals.contains(&"192.168.0.0/16".to_string()));
    }

    #[test]
    fn single_ip_comparison() {
        let ex = extract_expr("ip.src eq 10.0.0.1");
        assert_eq!(ex.ip_literals, vec!["10.0.0.1"]);
        assert_eq!(ex.fields, vec!["ip.src"]);
    }

    // ── ContainsOneOf (via `contains` on a set) ──────────────

    // Note: ContainsOneOf cannot be constructed via the parser
    // ("Node should not be constructed as there is no syntax to do so").
    // We test the `contains` operator with a single value instead.
    #[test]
    fn contains_operator_extraction() {
        let ex = extract_expr(r#"http.request.uri.path contains "admin""#);
        assert!(ex.operators.contains(&"contains".to_string()));
        assert!(ex.string_literals.contains(&"admin".to_string()));
    }

    // ── Ordering operators ───────────────────────────────────

    #[test]
    fn all_ordering_ops_extracted() {
        let ex = extract_expr(
            "cf.bot_management.score gt 10 and \
             cf.bot_management.score ge 20 and \
             cf.bot_management.score lt 30 and \
             cf.bot_management.score le 40 and \
             cf.bot_management.score eq 50 and \
             cf.bot_management.score ne 60",
        );
        for op in &["gt", "ge", "lt", "le", "eq", "ne"] {
            assert!(
                ex.operators.contains(&op.to_string()),
                "missing operator: {op}"
            );
        }
        for val in &[10, 20, 30, 40, 50, 60] {
            assert!(ex.int_literals.contains(val), "missing int literal: {val}");
        }
    }

    // ── Regex extraction ─────────────────────────────────────

    #[test]
    fn regex_literal_extracted() {
        let ex = extract_expr(r#"http.request.uri.path matches "^/api/v[0-9]+""#);
        assert!(ex.operators.contains(&"matches".to_string()));
        assert_eq!(ex.regex_literals.len(), 1);
        assert!(ex.regex_literals[0].contains("/api/v[0-9]+"));
    }

    // ── Wildcard extraction ──────────────────────────────────

    #[test]
    fn wildcard_operator_extracted() {
        let ex = extract_expr(r#"http.request.uri.path wildcard "/api/*""#);
        assert!(ex.operators.contains(&"wildcard".to_string()));
        assert!(ex.string_literals.contains(&"/api/*".to_string()));
    }

    // ── OneOf (in set) extraction ────────────────────────────

    #[test]
    fn one_of_bytes_extraction() {
        let ex = extract_expr(r#"http.host in {"a.com" "b.com" "c.com"}"#);
        assert!(ex.operators.contains(&"in".to_string()));
        assert_eq!(ex.string_literals.len(), 3);
    }

    #[test]
    fn one_of_ip_extraction() {
        let ex = extract_expr("ip.src in {1.2.3.4 5.6.7.8}");
        assert!(ex.operators.contains(&"in".to_string()));
        assert_eq!(ex.ip_literals.len(), 2);
    }

    #[test]
    fn one_of_int_extraction() {
        let ex = extract_expr("cf.bot_management.score in {10 20 30}");
        assert!(ex.operators.contains(&"in".to_string()));
        assert_eq!(ex.int_literals, vec![10, 20, 30]);
    }

    // ── Function extraction ──────────────────────────────────

    #[test]
    fn function_call_extracted() {
        let ex = extract_expr(r#"lower(http.host) eq "example.com""#);
        assert!(ex.functions.contains(&"lower".to_string()));
        assert!(ex.fields.contains(&"http.host".to_string()));
        assert!(ex.string_literals.contains(&"example.com".to_string()));
    }

    // ── Logical operators ────────────────────────────────────

    #[test]
    fn and_or_not_operators() {
        let ex = extract_expr(r#"(http.host eq "a.com" and not ssl) or http.host eq "b.com""#);
        assert!(ex.operators.contains(&"and".to_string()));
        assert!(ex.operators.contains(&"or".to_string()));
        assert!(ex.operators.contains(&"not".to_string()));
    }

    #[test]
    fn xor_operator() {
        let ex = extract_expr(r#"http.host eq "a.com" xor http.host eq "b.com""#);
        assert!(ex.operators.contains(&"xor".to_string()));
    }

    // ── Depth guard ──────────────────────────────────────────

    #[test]
    fn depth_not_exceeded_for_simple_expr() {
        let ex = extract_expr(r#"http.host eq "example.com""#);
        assert!(!ex.depth_exceeded());
    }

    #[test]
    fn depth_not_exceeded_for_moderate_nesting() {
        // 10 levels of nesting via parentheses — well within limit
        let mut expr = String::from(r#"http.host eq "example.com""#);
        for _ in 0..10 {
            expr = format!("({expr})");
        }
        let ex = extract_expr(&expr);
        assert!(!ex.depth_exceeded());
        assert_eq!(ex.fields, vec!["http.host"]);
    }

    // ── Int range extraction ─────────────────────────────────

    #[test]
    fn int_range_both_endpoints() {
        let ex = extract_expr("cf.bot_management.score in {10..20}");
        assert!(ex.int_literals.contains(&10));
        assert!(ex.int_literals.contains(&20));
    }

    #[test]
    fn int_range_single_value() {
        let ex = extract_expr("cf.bot_management.score in {42..42}");
        // Should only have one entry for equal endpoints
        assert_eq!(ex.int_literals.iter().filter(|&&v| v == 42).count(), 1);
    }

    // ── Bitwise AND operator ─────────────────────────────────

    #[test]
    fn bitwise_and_operator() {
        // bitwise_and is a comparison-level op: `field & int` evaluates to bool
        let ex = extract_expr("cf.bot_management.score & 255");
        assert!(ex.operators.contains(&"bitwise_and".to_string()));
        assert!(ex.int_literals.contains(&255));
    }

    // ── Empty / new extractor ────────────────────────────────

    #[test]
    fn new_extractor_is_empty() {
        let ex = ExpressionExtractor::new();
        assert!(ex.fields.is_empty());
        assert!(ex.functions.is_empty());
        assert!(ex.operators.is_empty());
        assert!(ex.string_literals.is_empty());
        assert!(ex.regex_literals.is_empty());
        assert!(ex.ip_literals.is_empty());
        assert!(ex.int_literals.is_empty());
        assert!(!ex.depth_exceeded());
    }
}
