use anyhow::{Result, anyhow};
use pest::Parser as _;
use pest::iterators::Pair;
use pest_derive::Parser;

use super::ast::{AggFunc, BinOp, Expr, Filter, Formula};

#[derive(Parser)]
#[grammar = "formula.pest"]
struct FormulaParser;

/// Parse a formula string like "Profit = Revenue - Cost"
/// or "Tax = Revenue * 0.08 WHERE Region = \"East\""
pub fn parse_formula(raw: &str, target_category: &str) -> Result<Formula> {
    let input = raw.trim();
    let formula_pair = FormulaParser::parse(Rule::formula, input)
        .map_err(|e| anyhow!("{}", e))?
        .next()
        .ok_or_else(|| anyhow!("empty parse result"))?;
    build_formula(formula_pair, input, target_category)
}

/// Parse a bare expression (no target, no top-level WHERE clause).
/// Fails if the input contains trailing tokens after a complete expression.
pub fn parse_expr(s: &str) -> Result<Expr> {
    let input = s.trim();
    let expr_pair = FormulaParser::parse(Rule::expr_eoi, input)
        .map_err(|e| anyhow!("{}", e))?
        .next()
        .ok_or_else(|| anyhow!("empty parse result"))?
        .into_inner()
        .next()
        .ok_or_else(|| anyhow!("missing expression in expr_eoi"))?;
    build_expr(expr_pair)
}

// ---- tree walkers -------------------------------------------------------

fn build_formula(pair: Pair<Rule>, raw: &str, target_category: &str) -> Result<Formula> {
    let mut target = None;
    let mut expr = None;
    let mut filter = None;
    for inner in pair.into_inner() {
        match inner.as_rule() {
            Rule::target => target = Some(inner.as_str().trim().to_string()),
            Rule::expr => expr = Some(build_expr(inner)?),
            Rule::where_clause => filter = Some(build_filter(inner)?),
            Rule::EOI => {}
            r => return Err(anyhow!("unexpected rule in formula: {:?}", r)),
        }
    }
    let target = target.ok_or_else(|| anyhow!("missing target in formula"))?;
    let expr = expr.ok_or_else(|| anyhow!("missing expression in formula"))?;
    Ok(Formula::new(raw, target, target_category, expr, filter))
}

fn build_expr(pair: Pair<Rule>) -> Result<Expr> {
    // expr = { add_expr }
    build_add_expr(first_inner(pair, "expr")?)
}

fn build_add_expr(pair: Pair<Rule>) -> Result<Expr> {
    fold_left_binop(pair, build_mul_expr, |s| match s {
        "+" => Some(BinOp::Add),
        "-" => Some(BinOp::Sub),
        _ => None,
    })
}

fn build_mul_expr(pair: Pair<Rule>) -> Result<Expr> {
    fold_left_binop(pair, build_pow_expr, |s| match s {
        "*" => Some(BinOp::Mul),
        "/" => Some(BinOp::Div),
        _ => None,
    })
}

fn build_pow_expr(pair: Pair<Rule>) -> Result<Expr> {
    // pow_expr = { unary ~ (pow_op ~ unary)? }
    let mut pairs = pair.into_inner();
    let base_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("empty pow_expr"))?;
    let base = build_unary(base_pair)?;
    match pairs.next() {
        None => Ok(base),
        Some(op_pair) => {
            debug_assert_eq!(op_pair.as_rule(), Rule::pow_op);
            let exp_pair = pairs
                .next()
                .ok_or_else(|| anyhow!("missing exponent in pow_expr"))?;
            let exp = build_unary(exp_pair)?;
            Ok(Expr::BinOp(BinOp::Pow, Box::new(base), Box::new(exp)))
        }
    }
}

fn build_unary(pair: Pair<Rule>) -> Result<Expr> {
    // unary = { unary_minus | primary }
    let inner = first_inner(pair, "unary")?;
    match inner.as_rule() {
        Rule::unary_minus => {
            let prim = first_inner(inner, "unary_minus")?;
            Ok(Expr::UnaryMinus(Box::new(build_primary(prim)?)))
        }
        Rule::primary => build_primary(inner),
        r => Err(anyhow!("unexpected rule in unary: {:?}", r)),
    }
}

fn build_primary(pair: Pair<Rule>) -> Result<Expr> {
    // primary = { number | agg_call | if_expr | paren_expr | ref_expr }
    let inner = first_inner(pair, "primary")?;
    match inner.as_rule() {
        Rule::number => Ok(Expr::Number(inner.as_str().parse()?)),
        Rule::agg_call => build_agg_call(inner),
        Rule::if_expr => build_if_expr(inner),
        Rule::paren_expr => build_expr(first_inner(inner, "paren_expr")?),
        Rule::ref_expr => {
            let id_pair = first_inner(inner, "ref_expr")?;
            Ok(Expr::Ref(identifier_to_string(id_pair)))
        }
        r => Err(anyhow!("unexpected rule in primary: {:?}", r)),
    }
}

fn build_agg_call(pair: Pair<Rule>) -> Result<Expr> {
    // agg_call = { agg_func ~ "(" ~ expr ~ inline_where? ~ ")" }
    let mut pairs = pair.into_inner();
    let func_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing agg_func"))?;
    let func = parse_agg_func(func_pair.as_str())?;
    let expr_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing aggregate argument"))?;
    let inner_expr = build_expr(expr_pair)?;
    let filter = match pairs.next() {
        Some(p) if p.as_rule() == Rule::inline_where => Some(build_filter(p)?),
        _ => None,
    };
    Ok(Expr::Agg(func, Box::new(inner_expr), filter))
}

fn parse_agg_func(s: &str) -> Result<AggFunc> {
    match s.to_ascii_uppercase().as_str() {
        "SUM" => Ok(AggFunc::Sum),
        "AVG" => Ok(AggFunc::Avg),
        "MIN" => Ok(AggFunc::Min),
        "MAX" => Ok(AggFunc::Max),
        "COUNT" => Ok(AggFunc::Count),
        f => Err(anyhow!("unknown aggregate function: {}", f)),
    }
}

fn build_if_expr(pair: Pair<Rule>) -> Result<Expr> {
    // if_expr = { ^"IF" ~ "(" ~ comparison ~ "," ~ expr ~ "," ~ expr ~ ")" }
    let mut pairs = pair.into_inner();
    let cond_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing IF condition"))?;
    let cond = build_comparison(cond_pair)?;
    let then_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing IF then-branch"))?;
    let then_e = build_expr(then_pair)?;
    let else_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing IF else-branch"))?;
    let else_e = build_expr(else_pair)?;
    Ok(Expr::If(
        Box::new(cond),
        Box::new(then_e),
        Box::new(else_e),
    ))
}

fn build_comparison(pair: Pair<Rule>) -> Result<Expr> {
    // comparison = { expr ~ cmp_op ~ expr }
    let mut pairs = pair.into_inner();
    let lhs_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing comparison lhs"))?;
    let lhs = build_expr(lhs_pair)?;
    let op_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing comparison operator"))?;
    let op = parse_cmp_op(op_pair.as_str())?;
    let rhs_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing comparison rhs"))?;
    let rhs = build_expr(rhs_pair)?;
    Ok(Expr::BinOp(op, Box::new(lhs), Box::new(rhs)))
}

fn parse_cmp_op(s: &str) -> Result<BinOp> {
    match s {
        "=" => Ok(BinOp::Eq),
        "!=" => Ok(BinOp::Ne),
        "<" => Ok(BinOp::Lt),
        ">" => Ok(BinOp::Gt),
        "<=" => Ok(BinOp::Le),
        ">=" => Ok(BinOp::Ge),
        o => Err(anyhow!("unknown comparison operator: {}", o)),
    }
}

fn build_filter(pair: Pair<Rule>) -> Result<Filter> {
    // where_clause / inline_where both have shape:
    //   ^"WHERE" ~ identifier ~ "=" ~ filter_value
    let mut pairs = pair.into_inner();
    let cat_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing WHERE category"))?;
    let category = identifier_to_string(cat_pair);
    let val_pair = pairs
        .next()
        .ok_or_else(|| anyhow!("missing WHERE value"))?;
    let item = filter_value_to_string(val_pair);
    Ok(Filter { category, item })
}

fn filter_value_to_string(pair: Pair<Rule>) -> String {
    // filter_value = { string | pipe_quoted | bare_ident }
    let inner = pair
        .into_inner()
        .next()
        .expect("filter_value must have an inner pair");
    let s = inner.as_str();
    match inner.as_rule() {
        Rule::string => strip_string_quotes(s),
        Rule::pipe_quoted => unquote_pipe(s),
        _ => s.to_string(),
    }
}

/// Convert an identifier pair (identifier, pipe_quoted, or bare_ident) to
/// its content string.  Pipe-quoted identifiers have their delimiters
/// stripped and backslash escapes applied; bare identifiers are returned
/// verbatim.
fn identifier_to_string(pair: Pair<Rule>) -> String {
    let s = pair.as_str();
    if is_pipe_quoted(s) {
        unquote_pipe(s)
    } else {
        s.to_string()
    }
}

fn is_pipe_quoted(s: &str) -> bool {
    s.len() >= 2 && s.starts_with('|') && s.ends_with('|')
}

fn strip_string_quotes(s: &str) -> String {
    debug_assert!(s.len() >= 2 && s.starts_with('"') && s.ends_with('"'));
    s[1..s.len() - 1].to_string()
}

/// Strip surrounding pipes and apply backslash escapes: `\|` → `|`,
/// `\\` → `\`, `\n` → newline.  Matches the escape semantics documented
/// in src/persistence/improv.pest.
fn unquote_pipe(s: &str) -> String {
    debug_assert!(is_pipe_quoted(s));
    let inner = &s[1..s.len() - 1];
    let mut out = String::with_capacity(inner.len());
    let mut chars = inner.chars();
    while let Some(c) = chars.next() {
        if c == '\\' {
            match chars.next() {
                Some('|') => out.push('|'),
                Some('\\') => out.push('\\'),
                Some('n') => out.push('\n'),
                Some(other) => {
                    out.push('\\');
                    out.push(other);
                }
                None => out.push('\\'),
            }
        } else {
            out.push(c);
        }
    }
    out
}

// ---- small helpers ------------------------------------------------------

fn first_inner<'a>(pair: Pair<'a, Rule>, ctx: &str) -> Result<Pair<'a, Rule>> {
    pair.into_inner()
        .next()
        .ok_or_else(|| anyhow!("empty rule: {}", ctx))
}

/// Fold a left-associative binary-operator rule of the shape
/// `rule = { child ~ (op ~ child)* }` into a left-leaning BinOp tree.
fn fold_left_binop<F, M>(pair: Pair<Rule>, mut build_child: F, match_op: M) -> Result<Expr>
where
    F: FnMut(Pair<Rule>) -> Result<Expr>,
    M: Fn(&str) -> Option<BinOp>,
{
    let mut pairs = pair.into_inner();
    let first = pairs
        .next()
        .ok_or_else(|| anyhow!("empty binop rule"))?;
    let mut left = build_child(first)?;
    while let Some(op_pair) = pairs.next() {
        let op = match_op(op_pair.as_str()).ok_or_else(|| {
            anyhow!("unexpected operator token: {:?}", op_pair.as_str())
        })?;
        let right_pair = pairs
            .next()
            .ok_or_else(|| anyhow!("missing rhs for operator"))?;
        let right = build_child(right_pair)?;
        left = Expr::BinOp(op, Box::new(left), Box::new(right));
    }
    Ok(left)
}

#[cfg(test)]
mod tests {
    use super::parse_formula;
    use crate::{AggFunc, BinOp, Expr};

    #[test]
    fn parse_simple_subtraction() {
        let f = parse_formula("Profit = Revenue - Cost", "Foo").unwrap();
        assert_eq!(f.target, "Profit");
        assert_eq!(f.target_category, "Foo");
        assert!(matches!(f.expr, Expr::BinOp(BinOp::Sub, _, _)));
    }

    #[test]
    fn parse_where_clause() {
        let f = parse_formula("EastRev = Revenue WHERE Region = \"East\"", "Foo").unwrap();
        assert_eq!(f.target, "EastRev");
        let filter = f.filter.as_ref().unwrap();
        assert_eq!(filter.category, "Region");
        assert_eq!(filter.item, "East");
    }

    #[test]
    fn parse_sum_aggregation() {
        let f = parse_formula("Total = SUM(Revenue)", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::Agg(AggFunc::Sum, _, _)));
    }

    #[test]
    fn parse_avg_aggregation() {
        let f = parse_formula("Avg = AVG(Revenue)", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::Agg(AggFunc::Avg, _, _)));
    }

    #[test]
    fn parse_if_expression() {
        let f = parse_formula("Capped = IF(Revenue > 1000, 1000, Revenue)", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::If(_, _, _)));
    }

    #[test]
    fn parse_numeric_literal() {
        let f = parse_formula("Fixed = 42", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::Number(n) if (n - 42.0).abs() < 1e-10));
    }

    #[test]
    fn parse_chained_arithmetic() {
        parse_formula("X = (A + B) * (C - D)", "Cat").unwrap();
    }

    #[test]
    fn parse_missing_equals_returns_error() {
        assert!(parse_formula("BadFormula Revenue Cost", "Cat").is_err());
    }

    // ── Aggregate functions ─────────────────────────────────────────────

    #[test]
    fn parse_min_aggregation() {
        let f = parse_formula("Lo = MIN(Revenue)", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::Agg(AggFunc::Min, _, _)));
    }

    #[test]
    fn parse_max_aggregation() {
        let f = parse_formula("Hi = MAX(Revenue)", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::Agg(AggFunc::Max, _, _)));
    }

    #[test]
    fn parse_count_aggregation() {
        let f = parse_formula("N = COUNT(Revenue)", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::Agg(AggFunc::Count, _, _)));
    }

    // ── Aggregate with WHERE filter ─────────────────────────────────────

    #[test]
    fn parse_sum_with_top_level_where_works() {
        let f = parse_formula("EastTotal = SUM(Revenue) WHERE Region = \"East\"", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::Agg(AggFunc::Sum, _, _)));
        let filter = f.filter.as_ref().unwrap();
        assert_eq!(filter.category, "Region");
        assert_eq!(filter.item, "East");
    }

    /// Regression: WHERE inside aggregate parens must parse as the
    /// aggregate's inline filter, not as a top-level WHERE clause.
    #[test]
    fn parse_sum_with_inline_where_filter() {
        let f = parse_formula("EastTotal = SUM(Revenue WHERE Region = \"East\")", "Foo").unwrap();
        if let Expr::Agg(AggFunc::Sum, inner, Some(filter)) = &f.expr {
            assert!(matches!(**inner, Expr::Ref(_)));
            assert_eq!(filter.category, "Region");
            assert_eq!(filter.item, "East");
        } else {
            panic!("Expected SUM with inline WHERE filter, got: {:?}", f.expr);
        }
    }

    // ── Comparison operators ────────────────────────────────────────────

    #[test]
    fn parse_if_with_comparison_operators() {
        let f = parse_formula("X = IF(A != 0, A, 1)", "Cat").unwrap();
        assert!(matches!(f.expr, Expr::If(_, _, _)));

        let f = parse_formula("X = IF(A < 10, A, 10)", "Cat").unwrap();
        assert!(matches!(f.expr, Expr::If(_, _, _)));

        let f = parse_formula("X = IF(A <= 10, A, 10)", "Cat").unwrap();
        assert!(matches!(f.expr, Expr::If(_, _, _)));

        let f = parse_formula("X = IF(A >= 10, 10, A)", "Cat").unwrap();
        assert!(matches!(f.expr, Expr::If(_, _, _)));

        let f = parse_formula("X = IF(A = B, 1, 0)", "Cat").unwrap();
        assert!(matches!(f.expr, Expr::If(_, _, _)));
    }

    // ── Quoted strings in WHERE ─────────────────────────────────────────

    #[test]
    fn parse_where_with_quoted_string_inside_expression() {
        let f = parse_formula("X = Revenue WHERE Region = \"West Coast\"", "Foo").unwrap();
        let filter = f.filter.as_ref().unwrap();
        assert_eq!(filter.item, "West Coast");
    }

    // ── Power operator ──────────────────────────────────────────────────

    #[test]
    fn parse_power_operator() {
        let f = parse_formula("Sq = X ^ 2", "Cat").unwrap();
        assert!(matches!(f.expr, Expr::BinOp(BinOp::Pow, _, _)));
    }

    // ── Unary minus ─────────────────────────────────────────────────────

    #[test]
    fn parse_unary_minus() {
        let f = parse_formula("Neg = -Revenue", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::UnaryMinus(_)));
    }

    // ── Division and multiplication ─────────────────────────────────────

    #[test]
    fn parse_multiplication() {
        let f = parse_formula("Double = Revenue * 2", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::BinOp(BinOp::Mul, _, _)));
    }

    #[test]
    fn parse_division() {
        let f = parse_formula("Half = Revenue / 2", "Foo").unwrap();
        assert!(matches!(f.expr, Expr::BinOp(BinOp::Div, _, _)));
    }

    // ── Parenthesized expression ────────────────────────────────────────

    #[test]
    fn parse_nested_parens() {
        let f = parse_formula("X = ((A + B))", "Cat").unwrap();
        assert!(matches!(f.expr, Expr::BinOp(BinOp::Add, _, _)));
    }

    // ── Aggregate function name used as ref (no parens) ─────────────────

    #[test]
    fn parse_aggregate_name_without_parens_is_ref() {
        // "SUM" without parens should be treated as a reference, not a function
        let f = parse_formula("X = SUM + 1", "Cat").unwrap();
        assert!(matches!(f.expr, Expr::BinOp(BinOp::Add, _, _)));
        if let Expr::BinOp(_, lhs, _) = &f.expr {
            assert!(matches!(**lhs, Expr::Ref(_)));
        }
    }

    #[test]
    fn parse_if_without_parens_is_ref() {
        // "IF" without parens should be treated as a reference
        let f = parse_formula("X = IF + 1", "Cat").unwrap();
        if let Expr::BinOp(BinOp::Add, lhs, _) = &f.expr {
            assert!(matches!(**lhs, Expr::Ref(_)));
        } else {
            panic!("Expected BinOp(Add), got: {:?}", f.expr);
        }
    }

    // ── Quoted string in WHERE ──────────────────────────────────────────

    #[test]
    fn parse_quoted_string_in_where() {
        let f = parse_formula("X = Revenue WHERE Region = \"East\"", "Cat").unwrap();
        let filter = f.filter.as_ref().unwrap();
        assert_eq!(filter.item, "East");
    }

    // ── Error paths ─────────────────────────────────────────────────────

    #[test]
    fn parse_unexpected_token_error() {
        use super::parse_expr;
        // Extra tokens after a valid expression
        assert!(parse_expr("1 + 2 3").is_err());
    }

    #[test]
    fn parse_unexpected_character_error() {
        use super::parse_expr;
        assert!(parse_expr("@invalid").is_err());
    }

    #[test]
    fn parse_empty_expression_error() {
        use super::parse_expr;
        assert!(parse_expr("").is_err());
    }

    // ── Multi-word identifiers must be pipe-quoted ──────────────────────

    #[test]
    fn multi_word_bare_identifier_is_rejected() {
        // Multi-word identifiers must be pipe-quoted; bare multi-word fails
        // the `bare_name`-compatible grammar rule.
        assert!(parse_formula("Total Revenue = Base Revenue + Bonus", "Foo").is_err());
    }

    // ── WHERE inside quotes in the expression ───────────────────────────

    #[test]
    fn where_inside_quotes_is_not_a_keyword() {
        // A filter value containing the literal text "WHERE" is parsed as
        // a string, not as a nested WHERE keyword.
        let f = parse_formula("X = Revenue WHERE Region = \"WHERE\"", "Foo").unwrap();
        let filter = f.filter.as_ref().unwrap();
        assert_eq!(filter.item, "WHERE");
    }

    // ── Pipe-quoted identifiers ─────────────────────────────────────────

    #[test]
    fn pipe_quoted_identifier_in_expression() {
        let f = parse_formula("|Total Revenue| = |Base Revenue| + Bonus", "Foo").unwrap();
        assert_eq!(f.target, "|Total Revenue|");
        if let Expr::BinOp(BinOp::Add, lhs, rhs) = &f.expr {
            assert!(matches!(**lhs, Expr::Ref(ref s) if s == "Base Revenue"));
            assert!(matches!(**rhs, Expr::Ref(ref s) if s == "Bonus"));
        } else {
            panic!("Expected Add, got: {:?}", f.expr);
        }
    }

    #[test]
    fn pipe_quoted_keyword_as_identifier() {
        // A category named "WHERE" can be referenced with pipes
        let f = parse_formula("X = |WHERE| + |SUM|", "Cat").unwrap();
        if let Expr::BinOp(BinOp::Add, lhs, rhs) = &f.expr {
            assert!(matches!(**lhs, Expr::Ref(ref s) if s == "WHERE"));
            assert!(matches!(**rhs, Expr::Ref(ref s) if s == "SUM"));
        } else {
            panic!("Expected Add, got: {:?}", f.expr);
        }
    }

    #[test]
    fn pipe_quoted_identifier_with_special_chars() {
        // Pipes allow characters that would normally break tokenization
        let f = parse_formula("X = |Revenue (USD)| + |Cost + Tax|", "Cat").unwrap();
        if let Expr::BinOp(BinOp::Add, lhs, rhs) = &f.expr {
            assert!(matches!(**lhs, Expr::Ref(ref s) if s == "Revenue (USD)"));
            assert!(matches!(**rhs, Expr::Ref(ref s) if s == "Cost + Tax"));
        } else {
            panic!("Expected Add, got: {:?}", f.expr);
        }
    }

    #[test]
    fn pipe_quoted_in_aggregate() {
        let f = parse_formula("X = SUM(|Net Revenue|)", "Cat").unwrap();
        if let Expr::Agg(AggFunc::Sum, inner, None) = &f.expr {
            assert!(matches!(**inner, Expr::Ref(ref s) if s == "Net Revenue"));
        } else {
            panic!("Expected SUM aggregate, got: {:?}", f.expr);
        }
    }

    #[test]
    fn pipe_quoted_in_where_filter_value() {
        let f = parse_formula("X = Revenue WHERE Region = |East Coast|", "Foo").unwrap();
        let filter = f.filter.as_ref().unwrap();
        assert_eq!(filter.item, "East Coast");
    }

    #[test]
    fn pipe_quoted_in_inline_where() {
        let f =
            parse_formula("X = SUM(Revenue WHERE |Region Name| = |East Coast|)", "Foo").unwrap();
        if let Expr::Agg(AggFunc::Sum, _, Some(filter)) = &f.expr {
            assert_eq!(filter.category, "Region Name");
            assert_eq!(filter.item, "East Coast");
        } else {
            panic!("Expected SUM with WHERE filter, got: {:?}", f.expr);
        }
    }

    // ── Pipe-quoted escape semantics ────────────────────────────────────

    #[test]
    fn pipe_quoted_escape_literal_pipe() {
        // \| inside a pipe-quoted identifier is a literal pipe
        let f = parse_formula("X = |A\\|B|", "Cat").unwrap();
        if let Expr::Ref(ref s) = f.expr {
            assert_eq!(s, "A|B");
        } else {
            panic!("Expected Ref, got: {:?}", f.expr);
        }
    }
}
