// ============================================================================
// AST Node Types
// ============================================================================

use crate::model::{
    BinaryOperator, Block, ConstExpr, Declaration, Dependency, Expression, Name, Program,
    Statement, TypeId, UnaryOperator, Variable,
};

use super::lexer::{Token, TokenType};

// ============================================================================
// Parser
// ============================================================================

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

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

    fn error(&self, msg: &str) -> String {
        let token = self.current();
        format!(
            "Parser error at line {}, col {}: {}",
            token.line, token.col, msg
        )
    }

    fn current(&self) -> &Token {
        self.tokens
            .get(self.pos)
            .unwrap_or_else(|| self.tokens.last().unwrap())
    }

    fn peek(&self, offset: usize) -> &Token {
        self.tokens
            .get(self.pos + offset)
            .unwrap_or_else(|| self.tokens.last().unwrap())
    }

    fn advance(&mut self) -> &Token {
        if self.pos < self.tokens.len() - 1 {
            self.pos += 1;
        }
        self.current()
    }

    fn expect(&mut self, expected: TokenType) -> Result<Token, String> {
        let token = self.current().clone();
        if std::mem::discriminant(&token.token_type) != std::mem::discriminant(&expected)
        {
            return Err(self.error(&format!(
                "Expected {:?}, got {:?}",
                expected, token.token_type
            )));
        }
        self.advance();
        Ok(token)
    }

    pub fn parse(&mut self) -> Result<Program, String> {
        let mut declarations = Vec::new();

        while !matches!(self.current().token_type, TokenType::Eof) {
            declarations.push(self.parse_declaration()?);
        }

        Ok(Program { declarations })
    }

    fn parse_declaration(&mut self) -> Result<Declaration, String> {
        // check for an import
        if let TokenType::Include = self.current().token_type {
            self.advance();

            let name =
                if let TokenType::Identifier(id) = self.current().clone().token_type {
                    Some(id)
                } else {
                    None
                }
                .ok_or(String::from("Expected identifier"))?;

            self.advance();
            self.expect(TokenType::Colon)?;

            let path = if let TokenType::String(id) = self.current().clone().token_type {
                Some(id)
            } else {
                None
            }
            .ok_or(String::from("Expected string literal"))?;

            self.advance();
            return Ok(Declaration::Dependency(Dependency { name, path }));
        }

        self.expect(TokenType::Int)?;

        let name = match &self.current().token_type {
            TokenType::Identifier(s) => s.clone(),
            _ => return Err(self.error("Expected identifier")),
        };
        self.advance();

        match &self.current().token_type {
            TokenType::LParen => {
                // Function declaration
                self.advance();
                let mut params = Vec::<Variable>::new();

                if !matches!(self.current().token_type, TokenType::RParen) {
                    self.expect(TokenType::Int)?;

                    match &self.current().token_type {
                        TokenType::Identifier(s) => {
                            params.push(Variable {
                                name: s.clone(),
                                type_id: TypeId::U32,
                            });
                            self.advance();
                        }
                        _ => return Err(self.error("Expected parameter name")),
                    }

                    while matches!(self.current().token_type, TokenType::Comma) {
                        self.advance();
                        self.expect(TokenType::Int)?;

                        match &self.current().token_type {
                            TokenType::Identifier(s) => {
                                params.push(Variable {
                                    name: s.clone(),
                                    type_id: TypeId::U32,
                                });
                                self.advance();
                            }
                            _ => return Err(self.error("Expected parameter name")),
                        }
                    }
                }

                self.expect(TokenType::RParen)?;
                let body = self.parse_block()?;

                Ok(Declaration::Function {
                    name,
                    params,
                    body,
                    return_type: TypeId::U32,
                })
            }
            _ => {
                // Variable declaration
                let init = if matches!(self.current().token_type, TokenType::Assign) {
                    self.advance();

                    if let TokenType::Number(n) = self.current().token_type {
                        self.advance();
                        Some(ConstExpr::Number(n))
                    } else {
                        return Err(self
                            .error("Expected constant in global variable declaration"));
                    }
                } else {
                    None
                };

                self.expect(TokenType::Semicolon)?;
                Ok(Declaration::Variable {
                    var: Variable {
                        name,
                        type_id: TypeId::U32,
                    },
                    init,
                    is_const: false,
                })
            }
        }
    }

    fn parse_block(&mut self) -> Result<Block, String> {
        self.expect(TokenType::LBrace)?;
        let mut statements = Vec::new();

        while !matches!(self.current().token_type, TokenType::RBrace) {
            statements.push(self.parse_statement()?);
        }

        self.expect(TokenType::RBrace)?;
        Ok(statements)
    }

    fn parse_statement(&mut self) -> Result<Statement, String> {
        match &self.current().token_type {
            TokenType::LBrace => Ok(Statement::Block(self.parse_block()?)),
            TokenType::If => self.parse_if_stmt(),
            TokenType::While => self.parse_while_stmt(),
            TokenType::Return => self.parse_return_stmt(),
            TokenType::Identifier(name) => {
                let name = name.clone();

                // peek ahead for open paren (func call expr)
                if matches!(self.peek(1).token_type, TokenType::LParen) {
                    let expr = self.parse_expression()?; // a function call expr
                    self.expect(TokenType::Semicolon)?;
                    return Ok(Statement::Expression { expr });
                }

                self.advance(); // advance past identifier

                // assignment expression
                if matches!(self.current().token_type, TokenType::Assign) {
                    self.advance();
                    let expr = self.parse_expression()?;

                    self.expect(TokenType::Semicolon)?;
                    Ok(Statement::Assign {
                        varname: name,
                        value: expr,
                    })
                }
                // var expression
                else {
                    self.expect(TokenType::Semicolon)?;
                    Ok(Statement::Expression {
                        expr: Expression::Variable {
                            name: Name {
                                name,
                                namespace: None,
                            },
                            expr_type: None,
                        },
                    })
                }
            }
            TokenType::Int => {
                // Local variable declaration
                self.advance();
                let name = match &self.current().token_type {
                    TokenType::Identifier(s) => s.clone(),
                    _ => return Err(self.error("Expected variable name")),
                };
                self.advance();

                let init = if matches!(self.current().token_type, TokenType::Assign) {
                    self.advance();
                    Some(self.parse_expression()?)
                } else {
                    None
                };

                self.expect(TokenType::Semicolon)?;

                // Convert to assignment expression statement
                let expr = if let Some(init_expr) = init {
                    Statement::Assign {
                        varname: name,
                        value: init_expr,
                    }
                } else {
                    Statement::Assign {
                        varname: name,
                        value: Expression::Empty,
                    }
                };

                Ok(expr)
            }
            _ => {
                let expr = if matches!(self.current().token_type, TokenType::Semicolon) {
                    Expression::Empty
                } else {
                    self.parse_expression()?
                };

                self.expect(TokenType::Semicolon)?;
                Ok(Statement::Expression { expr })
            }
        }
    }

    fn parse_if_stmt(&mut self) -> Result<Statement, String> {
        self.expect(TokenType::If)?;
        self.expect(TokenType::LParen)?;
        let condition = self.parse_expression()?;
        self.expect(TokenType::RParen)?;
        let then_stmt = self.parse_block()?;

        let else_stmt = if matches!(self.current().token_type, TokenType::Else) {
            self.advance();
            self.parse_block()?
        } else {
            Vec::new()
        };

        Ok(Statement::If {
            condition,
            then_stmt,
            else_stmt,
        })
    }

    fn parse_while_stmt(&mut self) -> Result<Statement, String> {
        self.expect(TokenType::While)?;
        self.expect(TokenType::LParen)?;
        let condition = self.parse_expression()?;
        self.expect(TokenType::RParen)?;
        let body = self.parse_block()?;

        Ok(Statement::While { condition, body })
    }

    fn parse_return_stmt(&mut self) -> Result<Statement, String> {
        self.expect(TokenType::Return)?;

        let expr = if matches!(self.current().token_type, TokenType::Semicolon) {
            None
        } else {
            Some(self.parse_expression()?)
        };

        self.expect(TokenType::Semicolon)?;
        Ok(Statement::Return(expr))
    }

    fn parse_expression(&mut self) -> Result<Expression, String> {
        self.parse_comparison()
    }

    fn parse_comparison(&mut self) -> Result<Expression, String> {
        let mut expr = self.parse_additive()?;

        while let Some(op) = match &self.current().token_type {
            TokenType::Eq => Some(BinaryOperator::Eq),
            TokenType::Ne => Some(BinaryOperator::Ne),
            TokenType::Lt => Some(BinaryOperator::Lt),
            TokenType::Gt => Some(BinaryOperator::Gt),
            TokenType::Le => Some(BinaryOperator::Le),
            TokenType::Ge => Some(BinaryOperator::Ge),
            _ => None,
        } {
            self.advance();
            let right = Box::new(self.parse_additive()?);
            expr = Expression::Binary {
                op,
                left: Box::new(expr),
                right,
                type_id: None,
            };
        }

        Ok(expr)
    }

    fn parse_additive(&mut self) -> Result<Expression, String> {
        let mut expr = self.parse_multiplicative()?;

        while let Some(op) = match &self.current().token_type {
            TokenType::Plus => Some(BinaryOperator::Add),
            TokenType::Minus => Some(BinaryOperator::Sub),
            _ => None,
        } {
            self.advance();
            let right = Box::new(self.parse_multiplicative()?);
            expr = Expression::Binary {
                op,
                left: Box::new(expr),
                right,
                type_id: None,
            };
        }

        Ok(expr)
    }

    fn parse_multiplicative(&mut self) -> Result<Expression, String> {
        let mut expr = self.parse_unary()?;

        while let Some(op) = match &self.current().token_type {
            TokenType::Star => Some(BinaryOperator::Mul),
            TokenType::Slash => Some(BinaryOperator::Div),
            _ => None,
        } {
            self.advance();
            let right = Box::new(self.parse_unary()?);
            expr = Expression::Binary {
                op,
                left: Box::new(expr),
                right,
                type_id: None,
            };
        }

        Ok(expr)
    }

    fn parse_unary(&mut self) -> Result<Expression, String> {
        let op = match &self.current().token_type {
            TokenType::Plus => Some(UnaryOperator::Plus),
            TokenType::Minus => Some(UnaryOperator::Minus),
            _ => None,
        };

        if let Some(op) = op {
            self.advance();
            let operand = Box::new(self.parse_unary()?);
            return Ok(Expression::Unary {
                op,
                operand,
                type_id: None,
            });
        }

        self.parse_primary()
    }

    fn parse_primary(&mut self) -> Result<Expression, String> {
        match &self.current().token_type.clone() {
            TokenType::Number(n) => {
                let value = *n;
                self.advance();
                Ok(Expression::Number {
                    value: value as isize,
                    type_id: None,
                })
            }
            TokenType::Identifier(name) => {
                let name = name.clone();
                self.advance();

                if matches!(self.current().token_type, TokenType::LParen) {
                    // Function call
                    self.advance();
                    let mut args = Vec::new();

                    if !matches!(self.current().token_type, TokenType::RParen) {
                        args.push(self.parse_expression()?);

                        while matches!(self.current().token_type, TokenType::Comma) {
                            self.advance();
                            args.push(self.parse_expression()?);
                        }
                    }

                    self.expect(TokenType::RParen)?;
                    Ok(Expression::Call {
                        name: Name {
                            name,
                            namespace: None,
                        },
                        args,
                        type_id: None,
                    })
                } else {
                    Ok(Expression::Variable {
                        name: Name {
                            name,
                            namespace: None,
                        },
                        expr_type: None,
                    })
                }
            }
            TokenType::LParen => {
                self.advance();
                let expr = self.parse_expression()?;
                self.expect(TokenType::RParen)?;
                Ok(expr)
            }
            _ => Err(self.error(&format!(
                "Unexpected token: {:?}",
                self.current().token_type
            ))),
        }
    }
}
