use std::sync::atomic::{AtomicUsize, Ordering};

use crate::tokenizer::{
    Token,
    TokenType::{self, Identifier},
    ZernError, error,
};

#[derive(Debug, Clone)]
pub struct Param {
    pub var_type: Token,
    pub var_name: Token,
}

#[derive(Debug, Clone)]
pub enum Params {
    Normal(Vec<Param>),
    Variadic,
}

#[derive(Debug, Clone)]
pub enum Stmt {
    Expression(Expr),
    Declare {
        name: Token,
        var_type: Option<Token>,
        initializer: Expr,
    },
    Assign {
        left: Expr,
        op: Token,
        value: Expr,
    },
    Destructure {
        targets: Vec<Token>,
        op: Token,
        value: Expr,
    },
    Const {
        name: Token,
        value: Token,
        neg: bool,
    },
    Block(Vec<Stmt>),
    If {
        keyword: Token,
        condition: Expr,
        then_branch: Box<Stmt>,
        else_branch: Box<Stmt>,
    },
    While {
        keyword: Token,
        condition: Expr,
        body: Box<Stmt>,
    },
    For {
        var: Token,
        start: Expr,
        end: Expr,
        body: Box<Stmt>,
    },
    Function {
        name: Token,
        params: Params,
        return_types: Vec<Token>,
        body: Box<Stmt>,
        exported: bool,
    },
    Return {
        keyword: Token,
        exprs: Vec<Expr>,
    },
    Break,
    Continue,
    Extern(Token),
    Struct {
        name: Token,
        fields: Vec<Param>,
    },
    Enum {
        name: Token,
        variants: Vec<Token>,
    },
}

pub static NEXT_EXPR_ID: AtomicUsize = AtomicUsize::new(0);

#[derive(Debug, Clone)]
pub struct Expr {
    pub id: usize,
    pub kind: ExprKind,
}

impl Expr {
    pub fn new(kind: ExprKind) -> Expr {
        NEXT_EXPR_ID.fetch_add(1, Ordering::SeqCst);
        Expr {
            id: NEXT_EXPR_ID.load(Ordering::SeqCst),
            kind,
        }
    }
}

#[derive(Debug, Clone)]
pub enum ExprKind {
    Binary {
        left: Box<Expr>,
        op: Token,
        right: Box<Expr>,
    },
    Logical {
        left: Box<Expr>,
        op: Token,
        right: Box<Expr>,
    },
    Grouping(Box<Expr>),
    Literal(Token),
    Unary {
        op: Token,
        right: Box<Expr>,
    },
    Variable(Token),
    Call {
        callee: Box<Expr>,
        paren: Token,
        args: Vec<Expr>,
    },
    ArrayLiteral(Vec<Expr>),
    Index {
        expr: Box<Expr>,
        bracket: Token,
        index: Box<Expr>,
    },
    AddrOf {
        op: Token,
        expr: Box<Expr>,
    },
    New {
        struct_name: Token,
        use_heap: bool,
    },
    MemberAccess {
        left: Box<Expr>,
        field: Token,
    },
    Cast {
        expr: Box<Expr>,
        type_name: Token,
    },
    MethodCall {
        expr: Box<Expr>,
        method: Token,
        args: Vec<Expr>,
    },
}

pub struct Parser {
    tokens: Vec<Token>,
    current: usize,
    is_inside_function: bool,
}

impl Parser {
    pub fn new(tokens: Vec<Token>) -> Parser {
        Parser {
            tokens,
            current: 0,
            is_inside_function: false,
        }
    }

    pub fn parse(mut self) -> Result<Vec<Stmt>, ZernError> {
        let mut statements = vec![];
        while !self.eof() {
            statements.push(self.declaration()?);
        }
        Ok(statements)
    }

    fn declaration(&mut self) -> Result<Stmt, ZernError> {
        if !self.is_inside_function {
            if self.match_token(&[TokenType::KeywordFunc]) {
                return self.func_declaration(false);
            }
            if self.match_token(&[TokenType::KeywordExport]) {
                self.consume(TokenType::KeywordFunc, "expected 'func' after 'export'")?;
                return self.func_declaration(true);
            }
            if self.match_token(&[TokenType::KeywordExtern]) {
                return self.extern_declaration();
            }
            if self.match_token(&[TokenType::KeywordConst]) {
                return self.const_declaration();
            }
            if self.match_token(&[TokenType::KeywordStruct]) {
                return self.struct_declaration();
            }
            if self.match_token(&[TokenType::KeywordEnum]) {
                return self.enum_declaration();
            }
            return error!(
                self.peek().loc,
                "statements not allowed outside function body"
            );
        }

        self.statement()
    }

    fn func_declaration(&mut self, exported: bool) -> Result<Stmt, ZernError> {
        let name = self.consume(TokenType::Identifier, "expected function name")?;
        self.consume(TokenType::LeftBracket, "expected '[' after function name")?;

        let mut is_variadic = false;
        let mut params = vec![];
        if !self.check(&TokenType::RightBracket) {
            if self.match_token(&[TokenType::DoubleDot]) {
                is_variadic = true;
            } else {
                loop {
                    let var_name =
                        self.consume(TokenType::Identifier, "expected parameter name")?;
                    self.consume(TokenType::Colon, "expected ':' after parameter name")?;

                    let var_type =
                        self.consume(TokenType::Identifier, "expected parameter type")?;

                    params.push(Param { var_type, var_name });
                    if !self.match_token(&[TokenType::Comma]) {
                        break;
                    }
                }
            }
        }

        self.consume(TokenType::RightBracket, "expected ']' after arguments")?;
        self.consume(TokenType::Colon, "expected ':' after '['")?;

        let mut return_types = vec![];
        loop {
            return_types.push(self.consume(TokenType::Identifier, "expected return type")?);
            if !self.match_token(&[TokenType::Comma]) {
                break;
            }
        }

        self.is_inside_function = true;
        let body = Box::new(self.block()?);
        self.is_inside_function = false;

        Ok(Stmt::Function {
            name,
            params: if is_variadic {
                Params::Variadic
            } else {
                Params::Normal(params)
            },
            return_types,
            body,
            exported,
        })
    }

    fn struct_declaration(&mut self) -> Result<Stmt, ZernError> {
        let name = self.consume(TokenType::Identifier, "expected struct name")?;

        self.consume(TokenType::Indent, "expected indent after struct name")?;

        let mut fields = vec![];
        while !self.eof() && !self.check(&TokenType::Dedent) {
            let var_name = self.consume(TokenType::Identifier, "expected field name")?;
            self.consume(TokenType::Colon, "expected ':' after field name")?;

            let var_type = self.consume(TokenType::Identifier, "expected field type")?;

            fields.push(Param { var_type, var_name });
        }

        self.consume(TokenType::Dedent, "expected dedent after struct fields")?;

        Ok(Stmt::Struct { name, fields })
    }

    fn enum_declaration(&mut self) -> Result<Stmt, ZernError> {
        let name = self.consume(TokenType::Identifier, "expected enum name")?;

        self.consume(TokenType::Indent, "expected indent after enum name")?;

        let mut variants = vec![];
        while !self.eof() && !self.check(&TokenType::Dedent) {
            variants.push(self.consume(TokenType::Identifier, "expected variant name")?);
        }

        self.consume(TokenType::Dedent, "expected dedent after enum variants")?;

        Ok(Stmt::Enum { name, variants })
    }

    fn const_declaration(&mut self) -> Result<Stmt, ZernError> {
        let name = self.consume(TokenType::Identifier, "expected const name")?;
        self.consume(TokenType::Equal, "expected '=' after const name")?;
        let neg = self.match_token(&[TokenType::Minus]);
        let value = self.consume(TokenType::IntLiteral, "expected a number after '='")?;
        Ok(Stmt::Const { name, value, neg })
    }

    fn extern_declaration(&mut self) -> Result<Stmt, ZernError> {
        Ok(Stmt::Extern(
            self.consume(TokenType::Identifier, "expected extern name")?,
        ))
    }

    fn block(&mut self) -> Result<Stmt, ZernError> {
        self.consume(TokenType::Indent, "expected an indent")?;

        let mut statements = vec![];
        while !self.eof() && !self.match_token(&[TokenType::Dedent]) {
            statements.push(self.declaration()?);
        }

        Ok(Stmt::Block(statements))
    }

    fn statement(&mut self) -> Result<Stmt, ZernError> {
        if self.check_ahead(&TokenType::Colon) {
            let name = self.consume(TokenType::Identifier, "expected variable name")?;
            self.consume(TokenType::Colon, "expected ':'")?;

            let var_type = if self.match_token(&[TokenType::Equal]) {
                None
            } else {
                let var_type = self.consume(TokenType::Identifier, "expected variable type")?;
                self.consume(TokenType::Equal, "expected '=' after varaible type")?;
                Some(var_type)
            };

            let initializer = self.expression()?;
            Ok(Stmt::Declare {
                name,
                var_type,
                initializer,
            })
        } else if self.match_token(&[TokenType::KeywordIf]) {
            self.if_statement()
        } else if self.match_token(&[TokenType::KeywordWhile]) {
            self.while_statement()
        } else if self.match_token(&[TokenType::KeywordFor]) {
            self.for_statement()
        } else if self.match_token(&[TokenType::KeywordReturn]) {
            let keyword = self.previous().clone();
            let mut exprs = vec![];
            if !self.check(&TokenType::Dedent) {
                loop {
                    exprs.push(self.expression()?);
                    if !self.match_token(&[TokenType::Comma]) {
                        break;
                    }
                }
            }
            Ok(Stmt::Return { keyword, exprs })
        } else if self.match_token(&[TokenType::KeywordBreak]) {
            Ok(Stmt::Break)
        } else if self.match_token(&[TokenType::KeywordContinue]) {
            Ok(Stmt::Continue)
        } else if self.match_token(&[TokenType::Tilde]) {
            let mut targets = vec![];
            loop {
                targets.push(self.consume(Identifier, "expected an identifier")?);
                if !self.match_token(&[TokenType::Comma]) {
                    break;
                }
            }
            let op = self.consume(TokenType::Colon, "expected ':'")?;
            self.consume(TokenType::Equal, "expected '=' after ':'")?;
            let value = self.expression()?;
            Ok(Stmt::Destructure { targets, op, value })
        } else {
            let expr = self.expression()?;
            if self.match_token(&[TokenType::Equal]) {
                let op = self.previous().clone();
                let value = self.expression()?;
                Ok(Stmt::Assign {
                    left: expr,
                    op,
                    value,
                })
            } else if self.match_token(&[TokenType::PlusEqual, TokenType::MinusEqual]) {
                let op = self.previous().clone();
                let right = self.expression()?;
                let binary_token = Token {
                    token_type: match op.token_type {
                        TokenType::PlusEqual => TokenType::Plus,
                        TokenType::MinusEqual => TokenType::Minus,
                        _ => unreachable!(),
                    },
                    lexeme: match op.token_type {
                        TokenType::PlusEqual => String::from("+"),
                        TokenType::MinusEqual => String::from("-"),
                        _ => unreachable!(),
                    },
                    loc: op.loc.clone(),
                };
                Ok(Stmt::Assign {
                    left: expr.clone(),
                    op: Token {
                        token_type: TokenType::Equal,
                        lexeme: String::from("="),
                        loc: op.loc,
                    },
                    value: Expr::new(ExprKind::Binary {
                        left: Box::new(expr),
                        op: binary_token,
                        right: Box::new(right),
                    }),
                })
            } else {
                Ok(Stmt::Expression(expr))
            }
        }
    }

    fn if_statement(&mut self) -> Result<Stmt, ZernError> {
        let keyword = self.previous().clone();
        let condition = self.expression()?;
        let then_branch = self.block()?;
        let else_branch = if self.match_token(&[TokenType::KeywordElse]) {
            if self.match_token(&[TokenType::KeywordIf]) {
                Box::new(self.if_statement()?)
            } else {
                Box::new(self.block()?)
            }
        } else {
            Box::new(Stmt::Block(vec![]))
        };
        Ok(Stmt::If {
            keyword,
            condition,
            then_branch: Box::new(then_branch),
            else_branch,
        })
    }

    fn while_statement(&mut self) -> Result<Stmt, ZernError> {
        let keyword = self.previous().clone();
        let condition = self.expression()?;
        let body = self.block()?;
        Ok(Stmt::While {
            keyword,
            condition,
            body: Box::new(body),
        })
    }

    fn for_statement(&mut self) -> Result<Stmt, ZernError> {
        let var = self.consume(TokenType::Identifier, "expected variable name after 'for'")?;
        self.consume(TokenType::KeywordIn, "expected 'in' after variable name")?;
        let start = self.expression()?;
        self.consume(TokenType::DoubleDot, "expected '..' after the number")?;
        let end = self.expression()?;

        let body = self.block()?;
        Ok(Stmt::For {
            var,
            start,
            end,
            body: Box::new(body),
        })
    }

    fn expression(&mut self) -> Result<Expr, ZernError> {
        self.or_and()
    }

    fn or_and(&mut self) -> Result<Expr, ZernError> {
        let mut expr = self.equality()?;

        while self.match_token(&[TokenType::LogicalOr, TokenType::LogicalAnd]) {
            let op = self.previous().clone();
            let right = self.equality()?;
            expr = Expr::new(ExprKind::Logical {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            })
        }

        Ok(expr)
    }

    fn equality(&mut self) -> Result<Expr, ZernError> {
        let mut expr = self.comparison()?;

        while self.match_token(&[TokenType::DoubleEqual, TokenType::NotEqual]) {
            let op = self.previous().clone();
            let right = self.comparison()?;
            expr = Expr::new(ExprKind::Binary {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            })
        }

        Ok(expr)
    }

    fn comparison(&mut self) -> Result<Expr, ZernError> {
        let mut expr = self.term()?;

        while self.match_token(&[
            TokenType::Greater,
            TokenType::GreaterEqual,
            TokenType::LessEqual,
            TokenType::Less,
        ]) {
            let op = self.previous().clone();
            let right = self.term()?;
            expr = Expr::new(ExprKind::Binary {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            })
        }

        Ok(expr)
    }

    fn term(&mut self) -> Result<Expr, ZernError> {
        let mut expr = self.factor()?;

        while self.match_token(&[
            TokenType::Plus,
            TokenType::Minus,
            TokenType::Xor,
            TokenType::BitAnd,
            TokenType::BitOr,
        ]) {
            let op = self.previous().clone();
            let right = self.factor()?;
            expr = Expr::new(ExprKind::Binary {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            })
        }

        Ok(expr)
    }

    fn factor(&mut self) -> Result<Expr, ZernError> {
        let mut expr = self.cast()?;

        while self.match_token(&[
            TokenType::Star,
            TokenType::Slash,
            TokenType::Mod,
            TokenType::ShiftLeft,
            TokenType::ShiftRight,
        ]) {
            let op = self.previous().clone();
            let right = self.unary()?;
            expr = Expr::new(ExprKind::Binary {
                left: Box::new(expr),
                op,
                right: Box::new(right),
            })
        }

        Ok(expr)
    }

    fn cast(&mut self) -> Result<Expr, ZernError> {
        let mut expr = self.unary()?;

        while self.match_token(&[TokenType::KeywordAs]) {
            let type_name = self.consume(TokenType::Identifier, "expected type after 'as'")?;
            expr = Expr::new(ExprKind::Cast {
                expr: Box::new(expr),
                type_name,
            })
        }

        Ok(expr)
    }

    fn unary(&mut self) -> Result<Expr, ZernError> {
        if self.match_token(&[TokenType::Xor]) {
            let op = self.previous().clone();
            let right = self.unary()?;
            return Ok(Expr::new(ExprKind::AddrOf {
                op,
                expr: Box::new(right),
            }));
        }
        if self.match_token(&[TokenType::Bang, TokenType::Minus]) {
            let op = self.previous().clone();
            let right = self.unary()?;
            return Ok(Expr::new(ExprKind::Unary {
                op,
                right: Box::new(right),
            }));
        }

        self.call()
    }

    fn call(&mut self) -> Result<Expr, ZernError> {
        let mut expr = self.primary()?;

        loop {
            if self.peek().loc.line != self.previous().loc.line {
                break;
            }

            if self.match_token(&[TokenType::LeftParen]) {
                let mut args = vec![];
                if !self.check(&TokenType::RightParen) {
                    loop {
                        args.push(self.expression()?);
                        if !self.match_token(&[TokenType::Comma]) {
                            break;
                        }
                    }
                }

                let paren = self.consume(TokenType::RightParen, "expected ')' after arguments")?;

                expr = Expr::new(ExprKind::Call {
                    callee: Box::new(expr),
                    paren,
                    args,
                })
            } else if self.match_token(&[TokenType::LeftBracket]) {
                let index = self.expression()?;
                let bracket = self.consume(TokenType::RightBracket, "expected ']' after index")?;
                expr = Expr::new(ExprKind::Index {
                    expr: Box::new(expr),
                    bracket,
                    index: Box::new(index),
                })
            } else if self.match_token(&[TokenType::Arrow]) {
                if self.check(&TokenType::Identifier) && self.check_ahead(&TokenType::LeftParen) {
                    let method = self.consume(TokenType::Identifier, "expected method name")?;
                    self.consume(TokenType::LeftParen, "expected '('")?;
                    let mut args = vec![];
                    if !self.check(&TokenType::RightParen) {
                        loop {
                            args.push(self.expression()?);
                            if !self.match_token(&[TokenType::Comma]) {
                                break;
                            }
                        }
                    }
                    self.consume(TokenType::RightParen, "expected ')'")?;
                    expr = Expr::new(ExprKind::MethodCall {
                        expr: Box::new(expr),
                        method,
                        args,
                    });
                } else {
                    let field =
                        self.consume(TokenType::Identifier, "expected field name after '->'")?;
                    expr = Expr::new(ExprKind::MemberAccess {
                        left: Box::new(expr),
                        field,
                    });
                }
            } else {
                break;
            }
        }

        Ok(expr)
    }

    fn primary(&mut self) -> Result<Expr, ZernError> {
        if self.match_token(&[
            TokenType::IntLiteral,
            TokenType::CharLiteral,
            TokenType::StringLiteral,
            TokenType::True,
            TokenType::False,
        ]) {
            Ok(Expr::new(ExprKind::Literal(self.previous().clone())))
        } else if self.match_token(&[TokenType::LeftParen]) {
            let expr = self.expression()?;
            self.consume(TokenType::RightParen, "expected ')' after expression")?;
            Ok(Expr::new(ExprKind::Grouping(Box::new(expr))))
        } else if self.match_token(&[TokenType::LeftBracket]) {
            let mut xs = vec![];
            if !self.check(&TokenType::RightBracket) {
                loop {
                    xs.push(self.expression()?);
                    if !self.match_token(&[TokenType::Comma]) {
                        break;
                    }
                }
            }
            self.consume(TokenType::RightBracket, "expected ']' after values")?;

            Ok(Expr::new(ExprKind::ArrayLiteral(xs)))
        } else if self.match_token(&[TokenType::KeywordNew]) {
            let use_heap = self.match_token(&[TokenType::Star]);
            let struct_name =
                self.consume(TokenType::Identifier, "expected struct name after 'new'")?;
            Ok(Expr::new(ExprKind::New {
                struct_name,
                use_heap,
            }))
        } else if self.match_token(&[TokenType::Identifier]) {
            Ok(Expr::new(ExprKind::Variable(self.previous().clone())))
        } else {
            error!(
                self.peek().loc,
                format!("expected expression, got '{}'", self.peek().lexeme)
            )
        }
    }

    fn consume(&mut self, token_type: TokenType, message: &str) -> Result<Token, ZernError> {
        if self.check(&token_type) {
            self.current += 1;
            Ok(self.previous().clone())
        } else {
            error!(self.previous().loc, format!("{}", message))
        }
    }

    fn match_token(&mut self, token_types: &[TokenType]) -> bool {
        for x in token_types {
            if self.check(x) {
                self.current += 1;
                return true;
            }
        }
        false
    }

    fn check_ahead(&self, token_type: &TokenType) -> bool {
        if self.current + 1 >= self.tokens.len() {
            false
        } else {
            self.tokens[self.current + 1].token_type == *token_type
        }
    }

    fn check(&self, token_type: &TokenType) -> bool {
        if self.eof() {
            false
        } else {
            self.peek().token_type == *token_type
        }
    }

    fn peek(&self) -> &Token {
        &self.tokens[self.current]
    }

    fn previous(&self) -> &Token {
        &self.tokens[self.current - 1]
    }

    fn eof(&self) -> bool {
        self.peek().token_type == TokenType::Eof
    }
}