use super::block_expression::{block_expr, BlockExpr};
use super::literals;
use super::tokens::{token, Token, TokenKind as tk, TokenSpan};
use super::{IResult, Span};
use nom::{
    branch::alt,
    combinator::{map, opt},
    multi::separated_list0,
    sequence::{delimited, separated_pair, tuple},
};

#[derive(Debug, Clone)]
pub struct Literal<'a> {
    pub tok: Token<'a>,
    pub kind: LiteralKind,
}

#[derive(Debug, Clone)]
pub enum LiteralKind {
    Num(u32),
    Const(u32, u64),
}

#[derive(Debug, Clone)]
pub struct Call<'a> {
    pub name: Span<'a>,
    pub args: Vec<Expression<'a>>,
}

#[derive(Debug, Clone)]
pub struct StructInit<'a> {
    pub name: Span<'a>,
    pub args: Vec<(Token<'a>, Expression<'a>)>,
}

#[derive(Debug, Clone)]
pub enum BinOpKind {
    And,
    Or,
    Xor,
}

#[derive(Debug, Clone)]
pub enum UnOpKind {
    Not,
    BitNot,
}

#[derive(Debug, Clone)]
pub struct BinOp<'a> {
    pub a: Expression<'a>,
    pub b: Expression<'a>,
    pub kind: BinOpKind,
}

#[derive(Debug, Clone)]
pub struct UnOp<'a> {
    pub a: Expression<'a>,
    pub kind: UnOpKind,
}

#[derive(Debug, Clone)]
pub enum Expression<'a> {
    Path(&'a str),
    Literal(Literal<'a>),
    UnOp(Box<UnOp<'a>>),
    BinOp(Box<BinOp<'a>>),
    Call(Box<Call<'a>>),
    StructInit(Box<StructInit<'a>>),
    BlockExpr(Box<BlockExpr<'a>>),
}

/// expressions that can't be subdivided further
fn atom(input: TokenSpan) -> IResult<TokenSpan, Expression> {
    alt((
        map(call_item, |call| Expression::Call(Box::new(call))),
        map(struct_init_item, |sinit| {
            Expression::StructInit(Box::new(sinit))
        }),
        map(token(tk::Ident), |it| {
            Expression::Path(it.span().fragment())
        }),
        map(token(tk::Number), |tok| {
            Expression::Literal(Literal {
                tok,
                kind: LiteralKind::Num(tok.span().parse().expect("invalid literal")),
            })
        }),
        map(token(tk::Constant), |tok| {
            let (_, (width, val)) = literals::const_bits(tok.span()).expect("invalid literal");
            Expression::Literal(Literal {
                tok,
                kind: LiteralKind::Const(width, val),
            })
        }),
        delimited(token(tk::LParen), expression, token(tk::RParen)),
    ))(input)
}

fn unop_kind(input: TokenSpan) -> IResult<TokenSpan, UnOpKind> {
    alt((
        map(token(tk::Not), |_| UnOpKind::Not),
        map(token(tk::BitNot), |_| UnOpKind::BitNot),
    ))(input)
}

/// unary expressions e.g `~expr`, `!expr`
fn unary(input: TokenSpan) -> IResult<TokenSpan, Expression> {
    alt((
        map(tuple((unop_kind, unary)), |(kind, expr)| {
            Expression::UnOp(Box::new(UnOp { a: expr, kind }))
        }),
        atom,
    ))(input)
}

fn bitop_kind(input: TokenSpan) -> IResult<TokenSpan, BinOpKind> {
    alt((
        map(token(tk::BitXor), |_| BinOpKind::Xor),
        map(token(tk::BitOr), |_| BinOpKind::Or),
        map(token(tk::BitAnd), |_| BinOpKind::And),
    ))(input)
}

/// bit and, or, xor e.g. a ^ b & c
/// TODO: make precedence rules for bit ops
fn bitop(input: TokenSpan) -> IResult<TokenSpan, Expression> {
    // special care is given to avoid parsing `unary` twice, as that would
    // make this parser quadratic
    map(
        tuple((unary, opt(tuple((bitop_kind, bitop))))),
        |(a, rest)| {
            if let Some((kind, b)) = rest {
                Expression::BinOp(Box::new(BinOp { a, b, kind }))
            } else {
                a
            }
        },
    )(input)
}

pub fn call_item(input: TokenSpan) -> IResult<TokenSpan, Call> {
    map(
        tuple((
            token(tk::Ident),
            delimited(
                token(tk::LParen),
                separated_list0(token(tk::Comma), expression),
                token(tk::RParen),
            ),
        )),
        |(name, args)| Call {
            name: name.span(),
            args,
        },
    )(input)
}

pub fn struct_init_item(input: TokenSpan) -> IResult<TokenSpan, StructInit> {
    map(
        tuple((
            token(tk::Ident),
            delimited(
                token(tk::LBrace),
                separated_list0(
                    token(tk::Comma),
                    separated_pair(token(tk::Ident), token(tk::Colon), expression),
                ),
                token(tk::RBrace),
            ),
        )),
        |(name, args)| StructInit {
            name: name.span(),
            args,
        },
    )(input)
}

pub fn expression(input: TokenSpan) -> IResult<TokenSpan, Expression> {
    alt((
        bitop,
        map(block_expr, |blk| Expression::BlockExpr(Box::new(blk))),
    ))(input)
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::parser::tokens::{lex, Token, TokenSpan};
    use crate::parser::Span;
    use nom::combinator::all_consuming;

    fn tok(input: &'static str) -> Vec<Token> {
        let input = Span::from(input);
        lex(input).unwrap().1
    }

    fn fullexpr(input: TokenSpan) -> IResult<TokenSpan, Expression> {
        all_consuming(expression)(input)
    }

    #[test]
    fn test_atoms() {
        fullexpr(TokenSpan::new(&tok("a"))).unwrap();
        fullexpr(TokenSpan::new(&tok("(a)"))).unwrap();
        fullexpr(TokenSpan::new(&tok(
            "(((((((((((((((((((a)))))))))))))))))))",
        )))
        .unwrap();
    }

    #[test]
    fn test_unary() {
        fullexpr(TokenSpan::new(&tok("asdf"))).unwrap();
        fullexpr(TokenSpan::new(&tok("~(asdf)"))).unwrap();
        fullexpr(TokenSpan::new(&tok("!asdf"))).unwrap();
        fullexpr(TokenSpan::new(&tok("~!(asdf)"))).unwrap();
    }

    #[test]
    fn test_bitop() {
        fullexpr(TokenSpan::new(&tok("a | b"))).unwrap();
        fullexpr(TokenSpan::new(&tok("a ^ b"))).unwrap();
        fullexpr(TokenSpan::new(&tok("~(a ^ b)"))).unwrap();
        fullexpr(TokenSpan::new(&tok("a ^ !b"))).unwrap();
        fullexpr(TokenSpan::new(&tok("a & !b & c"))).unwrap();
    }

    #[test]
    fn test_call() {
        fullexpr(TokenSpan::new(&tok("a()"))).unwrap();
    }

    #[test]
    fn test_struct_init() {
        fullexpr(TokenSpan::new(&tok("a{}"))).unwrap();
    }
}