//! Recursive-descent parser for SDE expression strings.
//!
//! Converts strings like `"mu - 0.5 * h"` into an AST ([`SdeExpr`]) that the
//! compiler then lowers to an [`SdeOp`] tape.

use bt_kernel::{BtError, Result};

// ---------------------------------------------------------------------------
// AST
// ---------------------------------------------------------------------------

/// Expression AST node produced by the parser.
#[derive(Clone, Debug)]
pub enum SdeExpr {
    /// Numeric literal.
    Number(f64),
    /// Identifier — resolved later by the compiler to param / state / builtin.
    Ident(String),
    /// Unary negation.
    Neg(Box<SdeExpr>),
    /// Logical not.
    Not(Box<SdeExpr>),
    /// Binary operation.
    BinOp {
        op: BinOp,
        lhs: Box<SdeExpr>,
        rhs: Box<SdeExpr>,
    },
    /// Function call, e.g. `sqrt(x)`, `normal(mu, sigma)`, `if(c, a, b)`.
    Call {
        name: String,
        args: Vec<SdeExpr>,
    },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
    Pow,
    Gt,
    Lt,
    Gte,
    Lte,
    Eq,
    And,
    Or,
}

// ---------------------------------------------------------------------------
// Tokenizer
// ---------------------------------------------------------------------------

#[derive(Clone, Debug, PartialEq)]
enum Token {
    Number(f64),
    Ident(String),
    // Operators
    Plus,
    Minus,
    Star,
    Slash,
    StarStar, // **
    // Comparison
    Gt,
    Lt,
    Gte,
    Lte,
    EqEq,
    // Logic
    AmpAmp,
    PipePipe,
    Bang,
    // Delimiters
    LParen,
    RParen,
    Comma,
    // Ternary
    Question,
    Colon,
    // End
    Eof,
}

struct Tokenizer<'a> {
    input: &'a [u8],
    pos: usize,
}

impl<'a> Tokenizer<'a> {
    fn new(input: &'a str) -> Self {
        Self {
            input: input.as_bytes(),
            pos: 0,
        }
    }

    fn peek_byte(&self) -> Option<u8> {
        self.input.get(self.pos).copied()
    }

    fn advance(&mut self) -> u8 {
        let b = self.input[self.pos];
        self.pos += 1;
        b
    }

    fn skip_whitespace(&mut self) {
        while self.pos < self.input.len() && self.input[self.pos].is_ascii_whitespace() {
            self.pos += 1;
        }
    }

    fn next_token(&mut self) -> Result<Token> {
        self.skip_whitespace();

        let Some(b) = self.peek_byte() else {
            return Ok(Token::Eof);
        };

        match b {
            b'+' => { self.advance(); Ok(Token::Plus) }
            b'-' => { self.advance(); Ok(Token::Minus) }
            b'/' => { self.advance(); Ok(Token::Slash) }
            b'(' => { self.advance(); Ok(Token::LParen) }
            b')' => { self.advance(); Ok(Token::RParen) }
            b',' => { self.advance(); Ok(Token::Comma) }
            b'?' => { self.advance(); Ok(Token::Question) }
            b':' => { self.advance(); Ok(Token::Colon) }

            b'*' => {
                self.advance();
                if self.peek_byte() == Some(b'*') {
                    self.advance();
                    Ok(Token::StarStar)
                } else {
                    Ok(Token::Star)
                }
            }

            b'>' => {
                self.advance();
                if self.peek_byte() == Some(b'=') {
                    self.advance();
                    Ok(Token::Gte)
                } else {
                    Ok(Token::Gt)
                }
            }

            b'<' => {
                self.advance();
                if self.peek_byte() == Some(b'=') {
                    self.advance();
                    Ok(Token::Lte)
                } else {
                    Ok(Token::Lt)
                }
            }

            b'=' => {
                self.advance();
                if self.peek_byte() == Some(b'=') {
                    self.advance();
                    Ok(Token::EqEq)
                } else {
                    Err(BtError::Expression("expected '==' but got single '='".into()))
                }
            }

            b'&' => {
                self.advance();
                if self.peek_byte() == Some(b'&') {
                    self.advance();
                    Ok(Token::AmpAmp)
                } else {
                    Err(BtError::Expression("expected '&&' but got single '&'".into()))
                }
            }

            b'|' => {
                self.advance();
                if self.peek_byte() == Some(b'|') {
                    self.advance();
                    Ok(Token::PipePipe)
                } else {
                    Err(BtError::Expression("expected '||' but got single '|'".into()))
                }
            }

            b'!' => {
                self.advance();
                Ok(Token::Bang)
            }

            // Number: digits, optional decimal, optional exponent
            b'0'..=b'9' | b'.' => self.read_number(),

            // Identifier: starts with alpha or _
            b'a'..=b'z' | b'A'..=b'Z' | b'_' => self.read_ident(),

            _ => Err(BtError::Expression(format!(
                "unexpected character '{}' at position {}",
                b as char, self.pos
            ))),
        }
    }

    fn read_number(&mut self) -> Result<Token> {
        let start = self.pos;
        // Integer part
        while self.pos < self.input.len() && self.input[self.pos].is_ascii_digit() {
            self.pos += 1;
        }
        // Decimal part
        if self.pos < self.input.len() && self.input[self.pos] == b'.' {
            self.pos += 1;
            while self.pos < self.input.len() && self.input[self.pos].is_ascii_digit() {
                self.pos += 1;
            }
        }
        // Exponent part (e.g. 1e-4)
        if self.pos < self.input.len() && (self.input[self.pos] == b'e' || self.input[self.pos] == b'E') {
            self.pos += 1;
            if self.pos < self.input.len() && (self.input[self.pos] == b'+' || self.input[self.pos] == b'-') {
                self.pos += 1;
            }
            while self.pos < self.input.len() && self.input[self.pos].is_ascii_digit() {
                self.pos += 1;
            }
        }
        let s = std::str::from_utf8(&self.input[start..self.pos])
            .map_err(|e| BtError::Expression(format!("invalid utf8 in number: {e}")))?;
        let v: f64 = s
            .parse()
            .map_err(|e| BtError::Expression(format!("invalid number '{s}': {e}")))?;
        Ok(Token::Number(v))
    }

    fn read_ident(&mut self) -> Result<Token> {
        let start = self.pos;
        while self.pos < self.input.len()
            && (self.input[self.pos].is_ascii_alphanumeric() || self.input[self.pos] == b'_')
        {
            self.pos += 1;
        }
        let s = std::str::from_utf8(&self.input[start..self.pos])
            .map_err(|e| BtError::Expression(format!("invalid utf8 in identifier: {e}")))?;
        Ok(Token::Ident(s.to_string()))
    }

    fn tokenize_all(&mut self) -> Result<Vec<Token>> {
        let mut tokens = Vec::new();
        loop {
            let tok = self.next_token()?;
            if tok == Token::Eof {
                tokens.push(Token::Eof);
                break;
            }
            tokens.push(tok);
        }
        Ok(tokens)
    }
}

// ---------------------------------------------------------------------------
// Parser
// ---------------------------------------------------------------------------

struct Parser {
    tokens: Vec<Token>,
    pos: usize,
}

impl Parser {
    fn new(tokens: Vec<Token>) -> Self {
        Self { tokens, pos: 0 }
    }

    fn peek(&self) -> &Token {
        self.tokens.get(self.pos).unwrap_or(&Token::Eof)
    }

    fn advance(&mut self) -> Token {
        let tok = self.tokens.get(self.pos).cloned().unwrap_or(Token::Eof);
        self.pos += 1;
        tok
    }

    fn expect(&mut self, expected: &Token) -> Result<()> {
        let tok = self.advance();
        if &tok == expected {
            Ok(())
        } else {
            Err(BtError::Expression(format!(
                "expected {expected:?} but got {tok:?}"
            )))
        }
    }

    // expression = ternary
    fn parse_expr(&mut self) -> Result<SdeExpr> {
        self.parse_ternary()
    }

    // ternary = or ("?" expression ":" expression)?
    fn parse_ternary(&mut self) -> Result<SdeExpr> {
        let cond = self.parse_or()?;
        if matches!(self.peek(), Token::Question) {
            self.advance(); // consume ?
            let then_expr = self.parse_expr()?;
            self.expect(&Token::Colon)?;
            let else_expr = self.parse_expr()?;
            Ok(SdeExpr::Call {
                name: "if".to_string(),
                args: vec![cond, then_expr, else_expr],
            })
        } else {
            Ok(cond)
        }
    }

    // or = and ("||" and)*
    fn parse_or(&mut self) -> Result<SdeExpr> {
        let mut lhs = self.parse_and()?;
        while matches!(self.peek(), Token::PipePipe) {
            self.advance();
            let rhs = self.parse_and()?;
            lhs = SdeExpr::BinOp {
                op: BinOp::Or,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
            };
        }
        Ok(lhs)
    }

    // and = comparison ("&&" comparison)*
    fn parse_and(&mut self) -> Result<SdeExpr> {
        let mut lhs = self.parse_comparison()?;
        while matches!(self.peek(), Token::AmpAmp) {
            self.advance();
            let rhs = self.parse_comparison()?;
            lhs = SdeExpr::BinOp {
                op: BinOp::And,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
            };
        }
        Ok(lhs)
    }

    // comparison = additive ((">" | "<" | ">=" | "<=" | "==") additive)?
    fn parse_comparison(&mut self) -> Result<SdeExpr> {
        let lhs = self.parse_additive()?;
        let op = match self.peek() {
            Token::Gt => Some(BinOp::Gt),
            Token::Lt => Some(BinOp::Lt),
            Token::Gte => Some(BinOp::Gte),
            Token::Lte => Some(BinOp::Lte),
            Token::EqEq => Some(BinOp::Eq),
            _ => None,
        };
        if let Some(op) = op {
            self.advance();
            let rhs = self.parse_additive()?;
            Ok(SdeExpr::BinOp {
                op,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
            })
        } else {
            Ok(lhs)
        }
    }

    // additive = multiplicative (("+" | "-") multiplicative)*
    fn parse_additive(&mut self) -> Result<SdeExpr> {
        let mut lhs = self.parse_multiplicative()?;
        loop {
            let op = match self.peek() {
                Token::Plus => Some(BinOp::Add),
                Token::Minus => Some(BinOp::Sub),
                _ => None,
            };
            if let Some(op) = op {
                self.advance();
                let rhs = self.parse_multiplicative()?;
                lhs = SdeExpr::BinOp {
                    op,
                    lhs: Box::new(lhs),
                    rhs: Box::new(rhs),
                };
            } else {
                break;
            }
        }
        Ok(lhs)
    }

    // multiplicative = power (("*" | "/") power)*
    fn parse_multiplicative(&mut self) -> Result<SdeExpr> {
        let mut lhs = self.parse_power()?;
        loop {
            let op = match self.peek() {
                Token::Star => Some(BinOp::Mul),
                Token::Slash => Some(BinOp::Div),
                _ => None,
            };
            if let Some(op) = op {
                self.advance();
                let rhs = self.parse_power()?;
                lhs = SdeExpr::BinOp {
                    op,
                    lhs: Box::new(lhs),
                    rhs: Box::new(rhs),
                };
            } else {
                break;
            }
        }
        Ok(lhs)
    }

    // power = unary ("**" unary)?
    fn parse_power(&mut self) -> Result<SdeExpr> {
        let base = self.parse_unary()?;
        if matches!(self.peek(), Token::StarStar) {
            self.advance();
            let exp = self.parse_unary()?;
            Ok(SdeExpr::BinOp {
                op: BinOp::Pow,
                lhs: Box::new(base),
                rhs: Box::new(exp),
            })
        } else {
            Ok(base)
        }
    }

    // unary = ("-" | "!") unary | call
    fn parse_unary(&mut self) -> Result<SdeExpr> {
        match self.peek() {
            Token::Minus => {
                self.advance();
                let inner = self.parse_unary()?;
                Ok(SdeExpr::Neg(Box::new(inner)))
            }
            Token::Bang => {
                self.advance();
                let inner = self.parse_unary()?;
                Ok(SdeExpr::Not(Box::new(inner)))
            }
            _ => self.parse_call(),
        }
    }

    // call = IDENT "(" args ")" | atom
    fn parse_call(&mut self) -> Result<SdeExpr> {
        if let Token::Ident(name) = self.peek().clone() {
            // Look ahead for '('
            if self.tokens.get(self.pos + 1) == Some(&Token::LParen) {
                self.advance(); // consume ident
                self.advance(); // consume '('
                let args = self.parse_args()?;
                self.expect(&Token::RParen)?;
                return Ok(SdeExpr::Call { name, args });
            }
        }
        self.parse_atom()
    }

    fn parse_args(&mut self) -> Result<Vec<SdeExpr>> {
        if matches!(self.peek(), Token::RParen) {
            return Ok(Vec::new());
        }
        let mut args = vec![self.parse_expr()?];
        while matches!(self.peek(), Token::Comma) {
            self.advance();
            args.push(self.parse_expr()?);
        }
        Ok(args)
    }

    // atom = NUMBER | IDENT | "(" expression ")"
    fn parse_atom(&mut self) -> Result<SdeExpr> {
        match self.peek().clone() {
            Token::Number(v) => {
                self.advance();
                Ok(SdeExpr::Number(v))
            }
            Token::Ident(name) => {
                self.advance();
                Ok(SdeExpr::Ident(name))
            }
            Token::LParen => {
                self.advance();
                let expr = self.parse_expr()?;
                self.expect(&Token::RParen)?;
                Ok(expr)
            }
            other => Err(BtError::Expression(format!(
                "unexpected token {other:?}, expected number, identifier, or '('"
            ))),
        }
    }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Parse an SDE expression string into an AST.
pub fn parse_sde_expr(input: &str) -> Result<SdeExpr> {
    if input.trim().is_empty() {
        return Err(BtError::Expression("empty expression".into()));
    }
    let mut tokenizer = Tokenizer::new(input);
    let tokens = tokenizer.tokenize_all()?;
    let mut parser = Parser::new(tokens);
    let expr = parser.parse_expr()?;

    // Ensure we consumed everything
    if !matches!(parser.peek(), Token::Eof) {
        return Err(BtError::Expression(format!(
            "unexpected trailing token {:?}",
            parser.peek()
        )));
    }

    Ok(expr)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_simple_arithmetic() {
        let expr = parse_sde_expr("mu - 0.5 * h").unwrap();
        // Should be Sub(Ident("mu"), Mul(Literal(0.5), Ident("h")))
        match expr {
            SdeExpr::BinOp { op: BinOp::Sub, .. } => {}
            other => panic!("expected Sub, got {other:?}"),
        }
    }

    #[test]
    fn test_function_call() {
        let expr = parse_sde_expr("sqrt(h)").unwrap();
        match expr {
            SdeExpr::Call { name, args } => {
                assert_eq!(name, "sqrt");
                assert_eq!(args.len(), 1);
            }
            other => panic!("expected Call, got {other:?}"),
        }
    }

    #[test]
    fn test_nested_functions() {
        let expr = parse_sde_expr("normal(mu_j, sigma_j)").unwrap();
        match expr {
            SdeExpr::Call { name, args } => {
                assert_eq!(name, "normal");
                assert_eq!(args.len(), 2);
            }
            other => panic!("expected Call, got {other:?}"),
        }
    }

    #[test]
    fn test_power() {
        let expr = parse_sde_expr("(ret - mu) ** 2").unwrap();
        match expr {
            SdeExpr::BinOp { op: BinOp::Pow, .. } => {}
            other => panic!("expected Pow, got {other:?}"),
        }
    }

    #[test]
    fn test_ternary() {
        let expr = parse_sde_expr("h > 0.01 ? 0.01 : h").unwrap();
        match expr {
            SdeExpr::Call { name, args } => {
                assert_eq!(name, "if");
                assert_eq!(args.len(), 3);
            }
            other => panic!("expected if Call, got {other:?}"),
        }
    }

    #[test]
    fn test_negation() {
        let expr = parse_sde_expr("-mu").unwrap();
        match expr {
            SdeExpr::Neg(_) => {}
            other => panic!("expected Neg, got {other:?}"),
        }
    }

    #[test]
    fn test_scientific_notation() {
        let expr = parse_sde_expr("1e-4 + 2.5E3").unwrap();
        match expr {
            SdeExpr::BinOp { op: BinOp::Add, lhs, rhs } => {
                match *lhs {
                    SdeExpr::Number(v) => assert!((v - 1e-4).abs() < 1e-20),
                    other => panic!("expected Number, got {other:?}"),
                }
                match *rhs {
                    SdeExpr::Number(v) => assert!((v - 2500.0).abs() < 1e-10),
                    other => panic!("expected Number, got {other:?}"),
                }
            }
            other => panic!("expected Add, got {other:?}"),
        }
    }

    #[test]
    fn test_complex_garch_expr() {
        let expr = parse_sde_expr("omega + alpha * (ret - mu) ** 2 + beta * h").unwrap();
        // omega + (alpha * ((ret - mu) ** 2)) + (beta * h)
        match expr {
            SdeExpr::BinOp { op: BinOp::Add, .. } => {}
            other => panic!("expected Add, got {other:?}"),
        }
    }

    #[test]
    fn test_empty_error() {
        assert!(parse_sde_expr("").is_err());
        assert!(parse_sde_expr("  ").is_err());
    }
}
