use crate::eval::Num;
use crate::lexer::Lexer;
use crate::op::{BinOp, Op, UnOp};
use num::BigInt;
use std::str::FromStr;

enum Pending {
    Paren,
    Un(UnOp),
    Bin(BinOp),
}

struct Parser<'a> {
    r: Lexer<'a>,
    first_token: Option<&'a str>,
    // stack stores the pending operations that can't be flushed yet.
    stack: Vec<Pending>,
    result: Vec<Op>,
}

impl<'a> Parser<'a> {
    fn peek(&mut self) -> Option<&'a str> {
        if let Some(t) = self.first_token {
            return Some(t);
        }
        let t = self.r.next()?;
        self.first_token = Some(t);
        Some(t)
    }

    fn next(&mut self) -> Option<&'a str> {
        if let Some(t) = std::mem::take(&mut self.first_token) {
            return Some(t);
        }
        self.r.next()
    }

    fn symbol(&mut self, s: &str) -> bool {
        let Some(t) = self.peek() else {
            return false;
        };
        if t != s {
            return false;
        }
        let _ = self.next();
        true
    }

    fn parse_num(&mut self) -> Option<()> {
        let t = self.next()?;
        if t.contains('.') || t.contains('e') || t.contains('E') {
            let Ok(f) = f64::from_str(t) else {
                return None;
            };
            self.result.push(Op::Num(Num::Float(f)));
            return Some(());
        }
        let Ok(i) = BigInt::from_str(t) else {
            return None;
        };
        self.result.push(Op::Num(Num::Int(i)));
        Some(())
    }

    fn parse_const(&mut self) -> Option<()> {
        if self.symbol("e") {
            self.result.push(Op::Num(Num::Float(std::f64::consts::E)));
            return Some(());
        }
        if self.symbol("pi") {
            self.result.push(Op::Num(Num::Float(std::f64::consts::PI)));
            return Some(());
        }
        self.parse_num()
    }

    fn parse_fun(&mut self) -> bool {
        if self.symbol("-") {
            self.stack.push(Pending::Un(UnOp::Neg));
            return true;
        }
        if self.symbol("sin") {
            self.stack.push(Pending::Un(UnOp::Sin));
            return true;
        }
        if self.symbol("cos") {
            self.stack.push(Pending::Un(UnOp::Cos));
            return true;
        }
        if self.symbol("tan") {
            self.stack.push(Pending::Un(UnOp::Tan));
            return true;
        }
        if self.symbol("asin") || self.symbol("arcsin") {
            self.stack.push(Pending::Un(UnOp::Asin));
            return true;
        }
        if self.symbol("acos") || self.symbol("arccos") {
            self.stack.push(Pending::Un(UnOp::Acos));
            return true;
        }
        if self.symbol("atan") || self.symbol("arctan") {
            self.stack.push(Pending::Un(UnOp::Atan));
            return true;
        }
        if self.symbol("sqrt") {
            self.stack.push(Pending::Un(UnOp::Sqrt));
            return true;
        }
        if self.symbol("log") || self.symbol("ln") {
            self.stack.push(Pending::Un(UnOp::Log));
            return true;
        }
        if self.symbol("log10") {
            self.stack.push(Pending::Un(UnOp::Log10));
            return true;
        }
        if self.symbol("log2") {
            self.stack.push(Pending::Un(UnOp::Log2));
            return true;
        }
        if self.symbol("floor") {
            self.stack.push(Pending::Un(UnOp::Floor));
            return true;
        }
        if self.symbol("ceil") || self.symbol("ceiling") {
            self.stack.push(Pending::Un(UnOp::Ceil));
            return true;
        }
        if self.symbol("round") {
            self.stack.push(Pending::Un(UnOp::Round));
            return true;
        }
        if self.symbol("abs") {
            self.stack.push(Pending::Un(UnOp::Abs));
            return true;
        }
        false
    }
}

fn prec(op: BinOp) -> i8 {
    match op {
        BinOp::Pow => 2,
        BinOp::Mul => 1,
        BinOp::Div => 1,
        BinOp::IntDiv => 1,
        BinOp::Mod => 1,
        BinOp::Add => 0,
        BinOp::Sub => 0,
    }
}

impl<'a> Parser<'a> {
    fn flush(&mut self, target_prec: i8) {
        loop {
            match self.stack.last() {
                None | Some(Pending::Paren) => break,
                Some(Pending::Bin(op)) if prec(*op) < target_prec => break,
                _ => (),
            }
            match self.stack.pop().unwrap() {
                Pending::Un(op) => self.result.push(Op::Un(op)),
                Pending::Bin(op) => self.result.push(Op::Bin(op)),
                Pending::Paren => unreachable!(),
            }
        }
    }

    fn parse_op(&mut self) {
        if self.symbol("^") {
            self.flush(3);
            self.stack.push(Pending::Bin(BinOp::Pow));
            return;
        }
        if self.symbol("*") {
            self.flush(1);
            self.stack.push(Pending::Bin(BinOp::Mul));
            return;
        }
        if self.symbol("/") {
            self.flush(1);
            self.stack.push(Pending::Bin(BinOp::Div));
            return;
        }
        if self.symbol("//") {
            self.flush(1);
            self.stack.push(Pending::Bin(BinOp::IntDiv));
            return;
        }
        if self.symbol("%") {
            self.flush(1);
            self.stack.push(Pending::Bin(BinOp::Mod));
            return;
        }
        if self.symbol("+") {
            self.flush(0);
            self.stack.push(Pending::Bin(BinOp::Add));
            return;
        }
        if self.symbol("-") {
            self.flush(0);
            self.stack.push(Pending::Bin(BinOp::Sub));
            return;
        }
        self.flush(1);
        self.stack.push(Pending::Bin(BinOp::Mul));
    }

    fn parse(&mut self) -> Option<()> {
        loop {
            if self.symbol("(") {
                self.stack.push(Pending::Paren);
                continue;
            }
            if self.parse_fun() {
                continue;
            }
            self.parse_const()?;
            while self.symbol(")") {
                loop {
                    match self.stack.pop()? {
                        Pending::Un(op) => self.result.push(Op::Un(op)),
                        Pending::Bin(op) => self.result.push(Op::Bin(op)),
                        Pending::Paren => break,
                    }
                }
            }
            if self.peek().is_none() {
                break;
            }
            self.parse_op();
        }
        loop {
            match self.stack.pop() {
                Some(Pending::Paren) => return None,
                Some(Pending::Un(op)) => self.result.push(Op::Un(op)),
                Some(Pending::Bin(op)) => self.result.push(Op::Bin(op)),
                None => return Some(()),
            }
        }
    }
}

// parse turns an expression into bytecode.
pub fn parse(expr: &str) -> Option<Vec<Op>> {
    let mut p = Parser {
        r: Lexer { buf: expr },
        first_token: None,
        stack: Vec::new(),
        result: Vec::new(),
    };
    p.parse()?;
    Some(p.result)
}

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

    fn int(i: i64) -> Op {
        Op::Num(Num::Int(BigInt::from(i)))
    }

    fn float(f: f64) -> Op {
        Op::Num(Num::Float(f))
    }

    #[test]
    fn parse_int() {
        assert_eq!(parse("500"), Some(vec![int(500)]));
    }

    #[test]
    fn parse_float() {
        assert_eq!(parse("1e2"), Some(vec![float(100.)]));
    }

    #[test]
    fn parse_fun() {
        assert_eq!(
            parse("sin pi"),
            Some(vec![float(std::f64::consts::PI), Op::Un(UnOp::Sin)])
        );
    }

    #[test]
    fn parse_nested_fun() {
        assert_eq!(
            parse("log log 100"),
            Some(vec![int(100), Op::Un(UnOp::Log), Op::Un(UnOp::Log)])
        );
    }

    #[test]
    fn parse_power() {
        assert_eq!(
            parse("2^1^2"),
            Some(vec![
                int(2),
                int(1),
                int(2),
                Op::Bin(BinOp::Pow),
                Op::Bin(BinOp::Pow)
            ])
        );
    }

    #[test]
    fn parse_fun_power() {
        assert_eq!(
            parse("log 2^2"),
            Some(vec![int(2), Op::Un(UnOp::Log), int(2), Op::Bin(BinOp::Pow)])
        );
    }

    #[test]
    fn parse_power_fun() {
        assert_eq!(
            parse("2^log 2"),
            Some(vec![int(2), int(2), Op::Un(UnOp::Log), Op::Bin(BinOp::Pow)])
        );
    }

    #[test]
    fn parse_pow_mul() {
        assert_eq!(
            parse("2^2*2"),
            Some(vec![
                int(2),
                int(2),
                Op::Bin(BinOp::Pow),
                int(2),
                Op::Bin(BinOp::Mul)
            ])
        );
    }

    #[test]
    fn parse_mul_pow() {
        assert_eq!(
            parse("2*2^2"),
            Some(vec![
                int(2),
                int(2),
                int(2),
                Op::Bin(BinOp::Pow),
                Op::Bin(BinOp::Mul)
            ])
        );
    }

    #[test]
    fn parse_mul_div() {
        assert_eq!(
            parse("2*2/2"),
            Some(vec![
                int(2),
                int(2),
                Op::Bin(BinOp::Mul),
                int(2),
                Op::Bin(BinOp::Div)
            ])
        );
    }

    #[test]
    fn parse_add_mul() {
        assert_eq!(
            parse("2+2*2"),
            Some(vec![
                int(2),
                int(2),
                int(2),
                Op::Bin(BinOp::Mul),
                Op::Bin(BinOp::Add)
            ])
        );
    }

    #[test]
    fn parse_mul_add() {
        assert_eq!(
            parse("2*2+2"),
            Some(vec![
                int(2),
                int(2),
                Op::Bin(BinOp::Mul),
                int(2),
                Op::Bin(BinOp::Add)
            ])
        );
    }

    #[test]
    fn parse_fun_add() {
        assert_eq!(
            parse("log 2+2"),
            Some(vec![int(2), Op::Un(UnOp::Log), int(2), Op::Bin(BinOp::Add)])
        );
    }

    #[test]
    fn parse_parens() {
        assert_eq!(
            parse("(1+2)*3"),
            Some(vec![
                int(1),
                int(2),
                Op::Bin(BinOp::Add),
                int(3),
                Op::Bin(BinOp::Mul)
            ])
        );
    }

    #[test]
    fn parse_implicit_multiplication() {
        assert_eq!(
            parse("2pi"),
            Some(vec![
                int(2),
                float(std::f64::consts::PI),
                Op::Bin(BinOp::Mul)
            ])
        );
    }

    #[test]
    fn parse_implicit_multiplication_neg() {
        assert_eq!(
            parse("-2 2 -2"),
            Some(vec![
                int(2),
                Op::Un(UnOp::Neg),
                int(2),
                Op::Bin(BinOp::Mul),
                int(2),
                Op::Bin(BinOp::Sub)
            ])
        );
    }

    #[test]
    fn parse_fun_implicit_multiplication() {
        assert_eq!(
            parse("sin 2pi"),
            Some(vec![
                int(2),
                Op::Un(UnOp::Sin),
                float(std::f64::consts::PI),
                Op::Bin(BinOp::Mul)
            ])
        );
    }

    #[test]
    fn parse_unary_negate() {
        assert_eq!(
            parse("1---2"),
            Some(vec![
                int(1),
                int(2),
                Op::Un(UnOp::Neg),
                Op::Un(UnOp::Neg),
                Op::Bin(BinOp::Sub)
            ])
        );
    }

    #[test]
    fn parse_negate_fun() {
        assert_eq!(
            parse("log-log 2"),
            Some(vec![
                int(2),
                Op::Un(UnOp::Log),
                Op::Un(UnOp::Neg),
                Op::Un(UnOp::Log)
            ])
        );
    }

    #[test]
    fn parse_unmatched_parens() {
        assert_eq!(parse("0))))"), None);
    }

    #[test]
    fn parse_unmatched_op() {
        assert_eq!(parse("2^2+"), None);
    }

    #[test]
    fn parse_unicode() {
        assert_eq!(parse("你好"), None);
    }
}