const std = @import("std");
const lex = @import("lexical_grammar.zig");
const Token = lex.Token;
const TokenType = lex.TokenType;
const ast = @import("ast/core.zig");

pub const ParseError = error{
    UnexpectedToken,
    ExpectedToken,
    ExpectedExpression,
    ExpectedStatement,
    ExpectedIdentifier,
    UnclosedString,
    InvalidNumber,
    InvalidEscape,
    UnterminatedComment,
    ExpectedParameterName,
    ExpectedPropertyName,
    UnterminatedTemplate,
    InvalidAssignmentTarget,
    TooManyTokens,
};

pub const BP = struct {
    pub const comma = 1;
    pub const assign = 2;
    pub const yield = 3;
    pub const conditional = 4;
    pub const coalesce = 5;
    pub const log_or = 6;
    pub const log_and = 7;
    pub const bit_or = 8;
    pub const bit_xor = 9;
    pub const bit_and = 10;
    pub const equality = 11;
    pub const relational = 12;
    pub const shift = 13;
    pub const additive = 14;
    pub const multiplicative = 15;
    pub const exponentiation = 16;
    pub const unary = 17;
    pub const update = 18;
    pub const call = 19;
    pub const member = 20;
};

pub const NodeArena = struct {
    exprs: [8192]ast.Expr = undefined,
    expr_count: usize = 0,
    stmts: [8192]ast.Stmt = undefined,
    stmt_count: usize = 0,
    decls: [1024]ast.Decl = undefined,
    decl_count: usize = 0,
    pats: [1024]ast.Pat = undefined,
    pat_count: usize = 0,
    lits: [1024]ast.Lit = undefined,
    lit_count: usize = 0,
    string_buf: [65536]u8 = undefined,
    string_pos: usize = 0,

    pub fn exprPtr(self: *NodeArena, data: ast.ExprType, loc: ast.Loc) *const ast.Expr {
        const idx = self.expr_count;
        self.exprs[idx] = .{ .loc = loc, .data = data };
        self.expr_count += 1;
        return &self.exprs[idx];
    }

    pub fn stmtPtr(self: *NodeArena, data: ast.StmtType, loc: ast.Loc) *const ast.Stmt {
        const idx = self.stmt_count;
        self.stmts[idx] = .{ .loc = loc, .data = data };
        self.stmt_count += 1;
        return &self.stmts[idx];
    }

    pub fn declPtr(self: *NodeArena, data: ast.DeclType, loc: ast.Loc) *const ast.Decl {
        const idx = self.decl_count;
        self.decls[idx] = .{ .loc = loc, .data = data };
        self.decl_count += 1;
        return &self.decls[idx];
    }

    pub fn patPtr(self: *NodeArena, data: ast.PatType, loc: ast.Loc) *const ast.Pat {
        const idx = self.pat_count;
        self.pats[idx] = .{ .loc = loc, .data = data };
        self.pat_count += 1;
        return &self.pats[idx];
    }

    pub fn litPtr(self: *NodeArena, data: ast.LitType, loc: ast.Loc) *const ast.Lit {
        const idx = self.lit_count;
        self.lits[idx] = .{ .loc = loc, .data = data };
        self.lit_count += 1;
        return &self.lits[idx];
    }
};

pub const Parser = struct {
    tokens: []const Token,
    pos: usize,
    source: []const u8,
    arena: *NodeArena,

    pub fn init(tokens: []const Token, source: []const u8, arena: *NodeArena) @This() {
        return .{ .tokens = tokens, .pos = 0, .source = source, .arena = arena };
    }

    pub fn tok(self: *@This()) Token {
        return self.tokens[self.pos];
    }

    pub fn peek(self: *@This()) TokenType {
        return self.tokens[self.pos].kind;
    }

    pub fn peekAt(self: *@This(), offset: usize) TokenType {
        const idx = self.pos + offset;
        if (idx >= self.tokens.len) return .eof;
        return self.tokens[idx].kind;
    }

    pub fn advance(self: *@This()) void {
        self.pos += 1;
    }

    pub fn advanceTok(self: *@This()) Token {
        const t = self.tokens[self.pos];
        self.pos += 1;
        return t;
    }

    pub fn expect(self: *@This(), kind: TokenType) ParseError!void {
        if (self.peek() != kind) return ParseError.ExpectedToken;
        self.advance();
    }

    pub fn expectTok(self: *@This(), kind: TokenType) ParseError!Token {
        if (self.peek() != kind) return ParseError.ExpectedToken;
        return self.advanceTok();
    }

    pub fn match(self: *@This(), kind: TokenType) bool {
        if (self.peek() == kind) {
            _ = self.advanceTok();
            return true;
        }
        return false;
    }

    pub fn matchKeyword(self: *@This(), kw: TokenType) bool {
        if (self.peek() == kw) {
            _ = self.advanceTok();
            return true;
        }
        return false;
    }

    pub fn isLineTerminatorBeforeNext(self: *@This()) bool {
        if (self.pos == 0) return false;
        const prev_token = self.tokens[self.pos - 1];
        const next_token = self.tokens[self.pos];
        var i = prev_token.loc.end;
        while (i < next_token.loc.start) : (i += 1) {
            if (i >= self.source.len) break;
            const c = self.source[i];
            switch (c) {
                '\n', '\r' => return true,
                else => {},
            }
        }
        return false;
    }

    pub fn as_identifier(self: *@This(), t: Token) []const u8 {
        return self.source[t.loc.start..t.loc.end];
    }

    pub fn tokenSlice(self: *@This(), t: Token) []const u8 {
        return self.source[t.loc.start..t.loc.end];
    }

    pub fn parseForDeclaration(self: *@This()) ParseError!*const ast.Decl {
        const kind_tok = self.advanceTok();
        var decls: [128]ast.VarDeclarator = undefined;
        var count: usize = 0;

        while (true) {
            const pat = try ast.Pat.parse(self);
            var decl_init: ?*const ast.Expr = null;
            if (self.match(.@"=")) {
                decl_init = try ast.Expr.parse(self, 0);
            }
            decls[count] = .{ .loc = pat.loc, .id = pat, .init = decl_init };
            count += 1;
            if (!self.match(.comma)) break;
        }

        const kind: ast.VarKind = switch (kind_tok.kind) {
            .@"var" => .@"var",
            .let => .let,
            .@"const" => .@"const",
            else => return ParseError.UnexpectedToken,
        };

        return self.arena.declPtr(.{ .@"var" = .{ .loc = .{ .start = kind_tok.loc.start, .end = decls[count - 1].loc.end }, .kind = kind, .decls = decls[0..count] } }, .{ .start = kind_tok.loc.start, .end = decls[count - 1].loc.end });
    }

    pub fn parseParams(self: *@This()) ParseError![]const ast.Pat {
        try self.expect(.lparen);
        var params: [128]ast.Pat = undefined;
        var count: usize = 0;

        if (self.peek() != .rparen) {
            while (true) {
                if (self.peek() == .@"...") {
                    self.advance();
                    const pat = try ast.Pat.parse(self);
                    params[count] = ast.Pat{ .loc = pat.loc, .data = .{ .rest = .{ .loc = pat.loc, .arg = self.arena.patPtr(pat.data, pat.loc) } } };
                    count += 1;
                    break;
                }
                const pat = try ast.Pat.parse(self);
                if (self.match(.@"=")) {
                    const expr = try ast.Expr.parse(self, 0);
                    params[count] = ast.Pat{ .loc = .{ .start = pat.loc.start, .end = expr.loc.end }, .data = .{ .assign = .{ .loc = .{ .start = pat.loc.start, .end = expr.loc.end }, .left = self.arena.patPtr(pat.data, pat.loc), .right = expr } } };
                } else {
                    params[count] = pat;
                }
                count += 1;
                if (!self.match(.comma)) break;
                if (self.peek() == .@"...") {
                    self.advance();
                    const rest_pat = try ast.Pat.parse(self);
                    params[count] = ast.Pat{ .loc = rest_pat.loc, .data = .{ .rest = .{ .loc = rest_pat.loc, .arg = self.arena.patPtr(rest_pat.data, rest_pat.loc) } } };
                    count += 1;
                    break;
                }
            }
        }

        try self.expect(.rparen);
        return params[0..count];
    }

    pub fn parseFunctionBody(self: *@This()) ParseError!ast.BlockStmt {
        return ast.BlockStmt.parse(self);
    }

    pub fn parsePrimaryExpr(self: *@This()) ParseError!*const ast.Expr {
        const start = self.tok().loc.start;

        switch (self.peek()) {
            .this => {
                const t = self.advanceTok();
                return self.arena.exprPtr(.{ .this = {} }, .{ .start = start, .end = t.loc.end });
            },
            .super => {
                const t = self.advanceTok();
                return self.arena.exprPtr(.{ .super = {} }, .{ .start = start, .end = t.loc.end });
            },
            .null => {
                const t = self.advanceTok();
                return self.arena.exprPtr(.{ .null = {} }, .{ .start = start, .end = t.loc.end });
            },
            .true => {
                const t = self.advanceTok();
                return self.arena.exprPtr(.{ .bool = true }, .{ .start = start, .end = t.loc.end });
            },
            .false => {
                const t = self.advanceTok();
                return self.arena.exprPtr(.{ .bool = false }, .{ .start = start, .end = t.loc.end });
            },
            .number => {
                const t = self.advanceTok();
                const val = std.fmt.parseFloat(f64, self.tokenSlice(t)) catch 0.0;
                return self.arena.exprPtr(.{ .number = val }, .{ .start = start, .end = t.loc.end });
            },
            .string => {
                const t = self.advanceTok();
                return self.arena.exprPtr(.{ .string = self.tokenSlice(t) }, .{ .start = start, .end = t.loc.end });
            },
            .bigint => {
                const t = self.advanceTok();
                return self.arena.exprPtr(.{ .number = 0 }, .{ .start = start, .end = t.loc.end });
            },
            .regex => {
                const t = self.advanceTok();
                const slice = self.tokenSlice(t);
                const slash_pos = std.mem.indexOfScalar(u8, slice[1..], '/') orelse 0;
                const pattern = slice[1 .. slash_pos + 1];
                const flags = slice[slash_pos + 2 ..];
                return self.arena.exprPtr(
                    .{ .regex = .{ .pattern = pattern, .flags = flags } },
                    .{ .start = start, .end = t.loc.end },
                );
            },
            .identifier => {
                const t = self.advanceTok();
                const name = self.tokenSlice(t);
                return self.arena.exprPtr(.{ .identifier = name }, .{ .start = start, .end = t.loc.end });
            },
            .private_identifier => {
                const t = self.advanceTok();
                const name = self.tokenSlice(t);
                return self.arena.exprPtr(.{ .private_ident = name }, .{ .start = start, .end = t.loc.end });
            },
            .lbrace => return ast.ObjectExpr.parse(self),
            .lbracket => return ast.ArrayExpr.parse(self),
            .lparen => {
                self.advance();
                if (self.peek() == .rparen) {
                    self.advance();
                    if (self.peek() == .@"=>") {
                        return ast.ArrowFn.parse(self, .{ .loc = .{ .start = start, .end = self.tok().loc.start }, .params = &.{} });
                    }
                    return self.arena.exprPtr(.{ .identifier = "" }, .{ .start = start, .end = self.tok().loc.start });
                }
                if (self.peek() == .@"...") {
                    self.advance();
                    const arg = try ast.Expr.parse(self, 0);
                    try self.expect(.rparen);
                    if (self.peek() == .@"=>") {
                        return ast.ArrowFn.parse(self, .{ .loc = .{ .start = start, .end = arg.loc.end }, .params = &.{} });
                    }
                    const loc: ast.Loc = .{ .start = start, .end = self.tok().loc.start };
                    return self.arena.exprPtr(.{ .spread = .{ .loc = loc, .arg = arg } }, loc);
                }
                const expr = try ast.Expr.parse(self, 0);
                if (self.peek() == .comma) {
                    var params: [128]ast.Pat = undefined;
                    var count: usize = 0;
                    params[count] = ast.Pat{ .loc = expr.loc, .data = .{ .ident = expr.data.identifier } };
                    count += 1;
                    while (self.match(.comma)) {
                        const p = try ast.Pat.parse(self);
                        params[count] = p;
                        count += 1;
                    }
                    try self.expect(.rparen);
                    if (self.peek() == .@"=>") {
                        return ast.ArrowFn.parse(self, .{ .loc = .{ .start = start, .end = self.tok().loc.start }, .params = params[0..count] });
                    }
                    return self.arena.exprPtr(.{ .identifier = "" }, .{ .start = start, .end = self.tok().loc.start });
                }
                try self.expect(.rparen);
                if (self.peek() == .@"=>") {
                    const pat = ast.Pat{ .loc = expr.loc, .data = .{ .ident = expr.data.identifier } };
                    return ast.ArrowFn.parse(self, .{ .loc = .{ .start = start, .end = self.tok().loc.start }, .params = &.{pat} });
                }
                return self.arena.exprPtr(.{ .parenthesized = expr }, .{ .start = start, .end = self.tok().loc.start });
            },
            .function => return ast.FnExpr.parse(self),
            .class => return ast.ClassExpr.parse(self),
            .async => {
                self.advance();
                if (self.peek() == .function) {
                    return ast.FnExpr.parse(self);
                }
                const id_tok = try self.expectTok(.identifier);
                const name = self.tokenSlice(id_tok);
                if (self.peek() == .@"=>") {
                    const pat = ast.Pat{ .loc = .{ .start = start, .end = id_tok.loc.end }, .data = .{ .ident = name } };
                    return ast.ArrowFn.parse(self, .{ .loc = .{ .start = start, .end = id_tok.loc.end }, .params = &.{pat} });
                }
                return self.arena.exprPtr(.{ .identifier = name }, .{ .start = start, .end = id_tok.loc.end });
            },
            .new => {
                self.advance();
                const callee = try ast.Expr.parse(self, BP.member);
                var args: [256]ast.Expr = undefined;
                var arg_count: usize = 0;
                if (self.peek() == .lparen) {
                    self.advance();
                    if (self.peek() != .rparen) {
                        while (true) {
                            const arg = try ast.Expr.parse(self, 0);
                            args[arg_count] = arg.*;
                            arg_count += 1;
                            if (!self.match(.comma)) break;
                        }
                    }
                    try self.expect(.rparen);
                }
                const loc: ast.Loc = .{ .start = start, .end = if (arg_count > 0) args[arg_count - 1].loc.end else callee.loc.end };
                return self.arena.exprPtr(.{ .new = .{ .loc = loc, .callee = callee, .args = args[0..arg_count] } }, loc);
            },
            .delete, .typeof, .void, .@"!", .@"~", .@"+", .@"-" => {
                const op = self.unaryOpFromToken();
                self.advance();
                const arg = try ast.Expr.parse(self, BP.unary);
                const loc: ast.Loc = .{ .start = start, .end = arg.loc.end };
                return self.arena.exprPtr(.{ .unary = .{ .loc = loc, .op = op, .arg = arg, .prefix = true } }, loc);
            },
            .@"++" => {
                self.advance();
                const arg = try ast.Expr.parse(self, BP.update);
                const loc: ast.Loc = .{ .start = start, .end = arg.loc.end };
                return self.arena.exprPtr(.{ .update = .{ .loc = loc, .op = .@"++", .arg = arg, .prefix = true } }, loc);
            },
            .@"--" => {
                self.advance();
                const arg = try ast.Expr.parse(self, BP.update);
                const loc: ast.Loc = .{ .start = start, .end = arg.loc.end };
                return self.arena.exprPtr(.{ .update = .{ .loc = loc, .op = .@"--", .arg = arg, .prefix = true } }, loc);
            },
            .await => {
                self.advance();
                const arg = try ast.Expr.parse(self, BP.unary);
                const loc: ast.Loc = .{ .start = start, .end = arg.loc.end };
                return self.arena.exprPtr(.{ .await = .{ .loc = loc, .arg = arg } }, loc);
            },
            .yield => {
                self.advance();
                var delegate = false;
                var arg: ?*const ast.Expr = null;
                if (!self.isLineTerminatorBeforeNext()) {
                    if (self.match(.@"*")) {
                        delegate = true;
                    }
                    if (self.peek() != .rbrace and self.peek() != .rparen and self.peek() != .rbracket and self.peek() != .semicolon and self.peek() != .comma and self.peek() != .@":" and self.peek() != .eof) {
                        arg = try ast.Expr.parse(self, 0);
                    }
                }
                const loc: ast.Loc = .{ .start = start, .end = if (arg) |a| a.loc.end else start + 5 };
                return self.arena.exprPtr(.{ .yield = .{ .loc = loc, .arg = arg, .delegate = delegate } }, loc);
            },
            .template_head, .no_sub_template => return self.arena.exprPtr(.{ .template = try ast.TemplateLit.parse(self) }, .{ .start = start, .end = self.tok().loc.start }),
            .import => {
                self.advance();
                try self.expect(.lparen);
                _ = try ast.Expr.parse(self, 0);
                try self.expect(.rparen);
                return self.arena.exprPtr(.{ .identifier = "import" }, .{ .start = start, .end = self.tok().loc.start });
            },
            .@"#" => {
                self.advance();
                const id_tok = try self.expectTok(.identifier);
                return self.arena.exprPtr(.{ .private_ident = self.tokenSlice(id_tok) }, .{ .start = start, .end = id_tok.loc.end });
            },
            else => {
                if (self.peek().isIdentifier()) {
                    const t = self.advanceTok();
                    const name = self.tokenSlice(t);
                    if (self.peek() == .@"=>") {
                        const pat = ast.Pat{ .loc = .{ .start = t.loc.start, .end = t.loc.end }, .data = .{ .ident = name } };
                        return ast.ArrowFn.parse(self, .{ .loc = .{ .start = t.loc.start, .end = t.loc.end }, .params = &.{pat} });
                    }
                    return self.arena.exprPtr(.{ .identifier = name }, .{ .start = t.loc.start, .end = t.loc.end });
                }
                return ParseError.ExpectedExpression;
            },
        }
    }

    pub fn parsePropKey(self: *@This()) ParseError!ast.PropKey {
        switch (self.peek()) {
            .identifier, .let, .static, .get, .set, .async => {
                const t = self.advanceTok();
                return ast.PropKey{ .ident = self.tokenSlice(t) };
            },
            .string => {
                const t = self.advanceTok();
                return ast.PropKey{ .string = self.tokenSlice(t) };
            },
            .number => {
                const t = self.advanceTok();
                const val = std.fmt.parseFloat(f64, self.tokenSlice(t)) catch 0.0;
                return ast.PropKey{ .number = val };
            },
            .lbracket => {
                self.advance();
                const expr = try ast.Expr.parse(self, 0);
                try self.expect(.rbracket);
                return ast.PropKey{ .computed = expr };
            },
            .private_identifier => {
                const t = self.advanceTok();
                return ast.PropKey{ .private = self.tokenSlice(t) };
            },
            else => return ParseError.ExpectedPropertyName,
        }
    }

    pub fn bp(_: *@This(), kind: TokenType) u32 {
        return switch (kind) {
            .comma => BP.comma,
            .@"=", .@"+=", .@"-=", .@"*=", .@"/=", .@"%=", .@"**=", .@"<<=", .@">>=", .@">>>=", .@"&=", .@"|=", .@"^=", .@"&&=", .@"||=", .@"??=" => BP.assign,
            .@"?" => BP.conditional,
            .@"??" => BP.coalesce,
            .@"||" => BP.log_or,
            .@"&&" => BP.log_and,
            .@"|" => BP.bit_or,
            .@"^" => BP.bit_xor,
            .@"&" => BP.bit_and,
            .@"==", .@"!=", .@"===", .@"!==" => BP.equality,
            .@"<", .@">", .@"<=", .@">=", .in, .instanceof => BP.relational,
            .@"<<", .@">>", .@">>>" => BP.shift,
            .@"+", .@"-" => BP.additive,
            .@"*", .@"/", .@"%" => BP.multiplicative,
            .@"**" => BP.exponentiation,
            .lparen => BP.call,
            .lbracket, .@".", .@"?." => BP.member,
            .@"++", .@"--" => BP.update,
            .template_head, .no_sub_template => BP.member,
            else => 0,
        };
    }

    pub fn unaryOpFromToken(self: *@This()) ast.UnaryOp {
        return switch (self.tok().kind) {
            .@"+" => .@"+",
            .@"-" => .@"-",
            .@"!" => .@"!",
            .@"~" => .@"~",
            .typeof => .typeof,
            .void => .void,
            .delete => .delete,
            else => .@"+",
        };
    }

    pub fn assignOpFromToken(self: *@This()) ast.AssignOp {
        return switch (self.tokens[self.pos - 1].kind) {
            .@"=" => .@"=",
            .@"+=" => .@"+=",
            .@"-=" => .@"-=",
            .@"*=" => .@"*=",
            .@"/=" => .@"/=",
            .@"%=" => .@"%=",
            .@"**=" => .@"**=",
            .@"<<=" => .@"<<=",
            .@">>=" => .@">>=",
            .@">>>=" => .@">>>=",
            .@"&=" => .@"&=",
            .@"|=" => .@"|=",
            .@"^=" => .@"^=",
            .@"&&=" => .@"&&=",
            .@"||=" => .@"||=",
            .@"??=" => .@"??=",
            else => .@"=",
        };
    }

    pub fn semicolon(self: *@This()) void {
        if (self.peek() == .semicolon) {
            self.advance();
        }
    }

    pub fn statementPointer(self: *@This()) ParseError!*const ast.Stmt {
        const s = try ast.Stmt.parse(self);
        return self.arena.stmtPtr(s.data, s.loc);
    }
};