const std = @import("std");

pub const Loc = struct {
    start: usize,
    end: usize,
};

pub const TokenType = enum {
    // Keywords
    await,
    @"break",
    case,
    @"catch",
    class,
    @"const",
    @"continue",
    debugger,
    default,
    delete,
    do,
    @"else",
    @"enum",
    @"export",
    extends,
    false,
    finally,
    @"for",
    function,
    @"if",
    import,
    in,
    instanceof,
    let,
    new,
    null,
    @"return",
    super,
    @"switch",
    this,
    throw,
    true,
    @"try",
    typeof,
    @"var",
    void,
    @"while",
    with,
    yield,
    using,

    // Identifiers and literals
    identifier,
    private_identifier,
    number,
    string,
    bigint,
    regex,
    template_head,
    template_middle,
    template_tail,
    no_sub_template,

    // Punctuators
    lbrace,
    rbrace,
    lparen,
    rparen,
    lbracket,
    rbracket,
    dot,
    semicolon,
    comma,
    @"...",
    @"<",
    @">",
    @"<=",
    @">=",
    @"==",
    @"!=",
    @"===",
    @"!==",
    @"+",
    @"-",
    @"*",
    @"%",
    @"**",
    @"++",
    @"--",
    @"<<",
    @">>",
    @">>>",
    @"&",
    @"|",
    @"^",
    @"!",
    @"~",
    @"&&",
    @"||",
    @"??",
    @"?",
    @":",
    @"=",
    @"+=",
    @"-=",
    @"*=",
    @"%=",
    @"**=",
    @"<<=",
    @">>=",
    @">>>=",
    @"&=",
    @"|=",
    @"^=",
    @"&&=",
    @"||=",
    @"??=",
    @"=>",
    @"/",
    @"/=",
    @"?.",
    @".",
    @"#",

    // Special
    eof,
    unknown,

    pub fn isKeyword(self: TokenType) bool {
        return switch (self) {
            .await,
            .@"break",
            .case,
            .@"catch",
            .class,
            .@"const",
            .@"continue",
            .debugger,
            .default,
            .delete,
            .do,
            .@"else",
            .@"enum",
            .@"export",
            .extends,
            .false,
            .finally,
            .@"for",
            .function,
            .@"if",
            .import,
            .in,
            .instanceof,
            .let,
            .new,
            .null,
            .@"return",
            .super,
            .@"switch",
            .this,
            .throw,
            .true,
            .@"try",
            .typeof,
            .@"var",
            .void,
            .@"while",
            .with,
            .yield,
            .using,
            => true,
            else => false,
        };
    }

    pub fn isIdentifier(self: TokenType) bool {
        return self == .identifier or self.isKeyword();
    }
};

pub const Token = struct {
    kind: TokenType,
    loc: Loc,
    slice: []const u8,
};

pub const Lexer = struct {
    source: []const u8,
    pos: usize,

    const Self = @This();

    pub fn init(source: []const u8) Self {
        return .{ .source = source, .pos = 0 };
    }

    pub fn tokenize(source: []const u8) []const Token {
        var lexer = Self.init(source);
        var tokens: [4096]Token = undefined;
        var count: usize = 0;

        while (true) {
            const tok = lexer.next();
            tokens[count] = tok;
            count += 1;
            if (tok.kind == .eof) break;
            if (count >= tokens.len - 1) {
                tokens[count] = .{ .kind = .eof, .loc = .{ .start = lexer.pos, .end = lexer.pos }, .slice = "" };
                break;
            }
        }

        return tokens[0..count];
    }

    fn ch(self: Self) ?u8 {
        if (self.pos >= self.source.len) return null;
        return self.source[self.pos];
    }

    fn chAt(self: Self, offset: usize) ?u8 {
        const idx = self.pos + offset;
        if (idx >= self.source.len) return null;
        return self.source[idx];
    }

    fn advance(self: *Self) void {
        if (self.pos < self.source.len) self.pos += 1;
    }

    fn isWhiteSpace(c: u8) bool {
        return switch (c) {
            ' ', '\t', 0x0B, 0x0C => true,
            else => false,
        };
    }

    fn isLineTerminator(c: u8) bool {
        return switch (c) {
            '\n', '\r' => true,
            0x2028, 0x2029 => true,
            else => false,
        };
    }

    fn isDigit(c: u8) bool {
        return c >= '0' and c <= '9';
    }

    fn isHexDigit(c: u8) bool {
        return switch (c) {
            '0'...'9', 'a'...'f', 'A'...'F' => true,
            else => false,
        };
    }

    fn isOctalDigit(c: u8) bool {
        return c >= '0' and c <= '7';
    }

    fn isIdStart(c: u8) bool {
        return switch (c) {
            'a'...'z', 'A'...'Z', '_', '$' => true,
            else => c > 0x7F,
        };
    }

    fn isIdContinue(c: u8) bool {
        return isIdStart(c) or isDigit(c) or c == 0x200C or c == 0x200D;
    }

    fn skipWhiteSpace(self: *Self) void {
        while (self.ch()) |c| {
            if (isWhiteSpace(c) or isLineTerminator(c)) {
                self.advance();
            } else {
                break;
            }
        }
    }

    fn skipSingleLineComment(self: *Self) void {
        self.advance();
        self.advance();
        while (self.ch()) |c| {
            if (isLineTerminator(c)) break;
            self.advance();
        }
    }

    fn skipMultiLineComment(self: *Self) void {
        self.advance();
        self.advance();
        while (self.ch()) |c| {
            if (c == '*' and self.chAt(1)) |n| {
                if (n == '/') {
                    self.advance();
                    self.advance();
                    return;
                }
            }
            self.advance();
        }
    }

    fn skipHashbang(self: *Self) void {
        if (self.ch()) |c| {
            if (c == '#') {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '!') {
                        self.advance();
                        while (self.ch()) |ch| {
                            if (isLineTerminator(ch)) break;
                            self.advance();
                        }
                    }
                }
            }
        }
    }

    fn scanString(self: *Self, quote: u8) Token {
        const start = self.pos;
        self.advance();

        while (self.ch()) |c| {
            if (c == quote) {
                self.advance();
                return .{
                    .kind = .string,
                    .loc = .{ .start = start, .end = self.pos },
                    .slice = self.source[start + 1 .. self.pos - 1],
                };
            }
            if (c == '\\') {
                self.advance();
                if (self.ch()) |_| self.advance();
            } else if (isLineTerminator(c) and c != 0x2028 and c != 0x2029) {
                break;
            } else {
                self.advance();
            }
        }

        return .{
            .kind = .string,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start + 1 .. self.pos],
        };
    }

    fn scanTemplate(self: *Self, start_kind: TokenType) Token {
        const start = self.pos;
        if (start_kind == .template_middle or start_kind == .template_head) {
            self.advance();
        }
        return self.scanTemplateRest(start);
    }

    fn scanTemplateRest(self: *Self, start: usize) Token {
        while (self.ch()) |c| {
            if (c == '`') {
                self.advance();
                return .{
                    .kind = .no_sub_template,
                    .loc = .{ .start = start, .end = self.pos },
                    .slice = self.source[start .. self.pos - 1],
                };
            }
            if (c == '$' and self.chAt(1)) |n| {
                if (n == '{') {
                    self.advance();
                    self.advance();
                    return .{
                        .kind = .template_head,
                        .loc = .{ .start = start, .end = self.pos },
                        .slice = self.source[start .. self.pos - 2],
                    };
                }
            }
            if (c == '\\') {
                self.advance();
                if (self.ch()) |_| self.advance();
            } else {
                self.advance();
            }
        }

        return .{
            .kind = .no_sub_template,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanTemplateMiddleOrTail(self: *Self) Token {
        const start = self.pos;
        self.advance();

        while (self.ch()) |c| {
            if (c == '`') {
                self.advance();
                return .{
                    .kind = .template_tail,
                    .loc = .{ .start = start, .end = self.pos },
                    .slice = self.source[start .. self.pos - 1],
                };
            }
            if (c == '$' and self.chAt(1)) |n| {
                if (n == '{') {
                    self.advance();
                    self.advance();
                    return .{
                        .kind = .template_middle,
                        .loc = .{ .start = start, .end = self.pos },
                        .slice = self.source[start .. self.pos - 2],
                    };
                }
            }
            if (c == '\\') {
                self.advance();
                if (self.ch()) |_| self.advance();
            } else {
                self.advance();
            }
        }

        return .{
            .kind = .template_tail,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanNumber(self: *Self) Token {
        const start = self.pos;

        if (self.ch()) |c| {
            if (c == '0') {
                self.advance();
                if (self.ch()) |n| {
                    switch (n) {
                        'x', 'X' => return self.scanHex(start),
                        'b', 'B' => return self.scanBinary(start),
                        'o', 'O' => return self.scanOctal(start),
                        'n' => {
                            self.advance();
                            return .{
                                .kind = .bigint,
                                .loc = .{ .start = start, .end = self.pos },
                                .slice = self.source[start .. self.pos - 1],
                            };
                        },
                        'e', 'E' => return self.scanDecimalAfterInt(start),
                        '.', '0'...'9' => return self.scanDecimalAfterInt(start),
                        else => {},
                    }
                }
            }
        }

        return self.scanDecimal(start);
    }

    fn scanDecimal(self: *Self, start: usize) Token {
        while (self.ch()) |c| {
            if (!isDigit(c) and c != '_') break;
            self.advance();
        }

        if (self.ch()) |c| {
            if (c == '.') {
                self.advance();
                while (self.ch()) |c2| {
                    if (!isDigit(c2) and c2 != '_') break;
                    self.advance();
                }
            }
            if (self.ch()) |c2| {
                if (c2 == 'e' or c2 == 'E') {
                    return self.scanExponent(start);
                }
            }
            if (self.ch()) |c2| {
                if (c2 == 'n') {
                    self.advance();
                    return .{
                        .kind = .bigint,
                        .loc = .{ .start = start, .end = self.pos },
                        .slice = self.source[start .. self.pos - 1],
                    };
                }
            }
        }

        return .{
            .kind = .number,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanDecimalAfterInt(self: *Self, start: usize) Token {
        while (self.ch()) |c| {
            if (!isDigit(c) and c != '_') break;
            self.advance();
        }
        if (self.ch()) |c| {
            if (c == '.') {
                self.advance();
                while (self.ch()) |c2| {
                    if (!isDigit(c2) and c2 != '_') break;
                    self.advance();
                }
            }
            if (self.ch()) |c2| {
                if (c2 == 'e' or c2 == 'E') {
                    return self.scanExponent(start);
                }
            }
            if (self.ch()) |c2| {
                if (c2 == 'n') {
                    self.advance();
                    return .{
                        .kind = .bigint,
                        .loc = .{ .start = start, .end = self.pos },
                        .slice = self.source[start .. self.pos - 1],
                    };
                }
            }
        }
        return .{
            .kind = .number,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanExponent(self: *Self, start: usize) Token {
        self.advance();
        if (self.ch()) |c| {
            if (c == '+' or c == '-') self.advance();
        }
        while (self.ch()) |c| {
            if (!isDigit(c) and c != '_') break;
            self.advance();
        }
        return .{
            .kind = .number,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanHex(self: *Self, start: usize) Token {
        self.advance();
        while (self.ch()) |c| {
            if (!isHexDigit(c) and c != '_') break;
            self.advance();
        }
        if (self.ch()) |c| {
            if (c == 'n') {
                self.advance();
                return .{
                    .kind = .bigint,
                    .loc = .{ .start = start, .end = self.pos },
                    .slice = self.source[start .. self.pos - 1],
                };
            }
        }
        return .{
            .kind = .number,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanBinary(self: *Self, start: usize) Token {
        self.advance();
        while (self.ch()) |c| {
            if (c != '0' and c != '1' and c != '_') break;
            self.advance();
        }
        if (self.ch()) |c| {
            if (c == 'n') {
                self.advance();
                return .{
                    .kind = .bigint,
                    .loc = .{ .start = start, .end = self.pos },
                    .slice = self.source[start .. self.pos - 1],
                };
            }
        }
        return .{
            .kind = .number,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanOctal(self: *Self, start: usize) Token {
        self.advance();
        while (self.ch()) |c| {
            if (!isOctalDigit(c) and c != '_') break;
            self.advance();
        }
        if (self.ch()) |c| {
            if (c == 'n') {
                self.advance();
                return .{
                    .kind = .bigint,
                    .loc = .{ .start = start, .end = self.pos },
                    .slice = self.source[start .. self.pos - 1],
                };
            }
        }
        return .{
            .kind = .number,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanIdentifierOrKeyword(self: *Self) Token {
        const start = self.pos;
        while (self.ch()) |c| {
            if (!isIdContinue(c)) break;
            self.advance();
        }

        const slice = self.source[start..self.pos];
        const kind = keywordFromString(slice);

        return .{
            .kind = kind,
            .loc = .{ .start = start, .end = self.pos },
            .slice = slice,
        };
    }

    fn scanPrivateIdentifier(self: *Self) Token {
        const start = self.pos;
        self.advance();
        while (self.ch()) |c| {
            if (!isIdContinue(c)) break;
            self.advance();
        }
        return .{
            .kind = .private_identifier,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanRegex(self: *Self) Token {
        const start = self.pos;
        self.advance();

        var in_class = false;
        while (self.ch()) |c| {
            if (c == '\\') {
                self.advance();
                if (self.ch()) |_| self.advance();
            } else if (c == '[') {
                in_class = true;
                self.advance();
            } else if (c == ']') {
                in_class = false;
                self.advance();
            } else if (c == '/' and !in_class) {
                self.advance();
                break;
            } else if (isLineTerminator(c)) {
                break;
            } else {
                self.advance();
            }
        }

        while (self.ch()) |c| {
            if (!isIdContinue(c)) break;
            self.advance();
        }

        return .{
            .kind = .regex,
            .loc = .{ .start = start, .end = self.pos },
            .slice = self.source[start..self.pos],
        };
    }

    fn scanPunctuator(self: *Self) Token {
        const start = self.pos;
        const c = self.ch().?;

        switch (c) {
            '{' => {
                self.advance();
                return simple(.lbrace, start);
            },
            '}' => {
                self.advance();
                return simple(.rbrace, start);
            },
            '(' => {
                self.advance();
                return simple(.lparen, start);
            },
            ')' => {
                self.advance();
                return simple(.rparen, start);
            },
            '[' => {
                self.advance();
                return simple(.lbracket, start);
            },
            ']' => {
                self.advance();
                return simple(.rbracket, start);
            },
            ';' => {
                self.advance();
                return simple(.semicolon, start);
            },
            ',' => {
                self.advance();
                return simple(.comma, start);
            },
            ':' => {
                self.advance();
                return simple(.@":", start);
            },
            '?' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '?') {
                        self.advance();
                        if (self.ch()) |n2| {
                            if (n2 == '=') {
                                self.advance();
                                return simple(.@"??=", start);
                            }
                        }
                        return simple(.@"??", start);
                    }
                    if (n == '.') {
                        if (self.chAt(1)) |n2| {
                            if (!isDigit(n2)) {
                                self.advance();
                                return simple(.@"?.", start);
                            }
                        } else {
                            self.advance();
                            return simple(.@"?.", start);
                        }
                    }
                }
                return simple(.@"?", start);
            },
            '~' => {
                self.advance();
                return simple(.@"~", start);
            },
            '.' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '.' and self.chAt(1)) |n2| {
                        if (n2 == '.') {
                            self.advance();
                            self.advance();
                            return simple(.@"...", start);
                        }
                    }
                }
                return simple(.@".", start);
            },
            '#' => {
                self.advance();
                if (self.ch()) |n| {
                    if (isIdStart(n)) {
                        return self.scanPrivateIdentifier();
                    }
                }
                return simple(.@"#", start);
            },
            '+' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '+') {
                        self.advance();
                        return simple(.@"++", start);
                    }
                    if (n == '=') {
                        self.advance();
                        return simple(.@"+=", start);
                    }
                }
                return simple(.@"+", start);
            },
            '-' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '-') {
                        self.advance();
                        return simple(.@"--", start);
                    }
                    if (n == '=') {
                        self.advance();
                        return simple(.@"-=", start);
                    }
                }
                return simple(.@"-", start);
            },
            '*' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '*') {
                        self.advance();
                        if (self.ch()) |n2| {
                            if (n2 == '=') {
                                self.advance();
                                return simple(.@"**=", start);
                            }
                        }
                        return simple(.@"**", start);
                    }
                    if (n == '=') {
                        self.advance();
                        return simple(.@"*=", start);
                    }
                }
                return simple(.@"*", start);
            },
            '%' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '=') {
                        self.advance();
                        return simple(.@"%=", start);
                    }
                }
                return simple(.@"%", start);
            },
            '=' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '=') {
                        self.advance();
                        if (self.ch()) |n2| {
                            if (n2 == '=') {
                                self.advance();
                                return simple(.@"===", start);
                            }
                        }
                        return simple(.@"==", start);
                    }
                    if (n == '>') {
                        self.advance();
                        return simple(.@"=>", start);
                    }
                }
                return simple(.@"=", start);
            },
            '!' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '=') {
                        self.advance();
                        if (self.ch()) |n2| {
                            if (n2 == '=') {
                                self.advance();
                                return simple(.@"!==", start);
                            }
                        }
                        return simple(.@"!=", start);
                    }
                }
                return simple(.@"!", start);
            },
            '<' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '<') {
                        self.advance();
                        if (self.ch()) |n2| {
                            if (n2 == '=') {
                                self.advance();
                                return simple(.@"<<=", start);
                            }
                        }
                        return simple(.@"<<", start);
                    }
                    if (n == '=') {
                        self.advance();
                        return simple(.@"<=", start);
                    }
                }
                return simple(.@"<", start);
            },
            '>' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '>') {
                        self.advance();
                        if (self.ch()) |n2| {
                            if (n2 == '>') {
                                self.advance();
                                if (self.ch()) |n3| {
                                    if (n3 == '=') {
                                        self.advance();
                                        return simple(.@">>>=", start);
                                    }
                                }
                                return simple(.@">>>", start);
                            }
                            if (n2 == '=') {
                                self.advance();
                                return simple(.@">>=", start);
                            }
                        }
                        return simple(.@">>", start);
                    }
                    if (n == '=') {
                        self.advance();
                        return simple(.@">=", start);
                    }
                }
                return simple(.@">", start);
            },
            '&' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '&') {
                        self.advance();
                        if (self.ch()) |n2| {
                            if (n2 == '=') {
                                self.advance();
                                return simple(.@"&&=", start);
                            }
                        }
                        return simple(.@"&&", start);
                    }
                    if (n == '=') {
                        self.advance();
                        return simple(.@"&=", start);
                    }
                }
                return simple(.@"&", start);
            },
            '|' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '|') {
                        self.advance();
                        if (self.ch()) |n2| {
                            if (n2 == '=') {
                                self.advance();
                                return simple(.@"||=", start);
                            }
                        }
                        return simple(.@"||", start);
                    }
                    if (n == '=') {
                        self.advance();
                        return simple(.@"|=", start);
                    }
                }
                return simple(.@"|", start);
            },
            '^' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '=') {
                        self.advance();
                        return simple(.@"^=", start);
                    }
                }
                return simple(.@"^", start);
            },
            '`' => {
                self.advance();
                return self.scanTemplateRest(start);
            },
            '\'' => return self.scanString('\''),
            '"' => return self.scanString('"'),
            '/' => {
                self.advance();
                if (self.ch()) |n| {
                    if (n == '/') {
                        self.skipSingleLineComment();
                        return self.next();
                    }
                    if (n == '*') {
                        self.skipMultiLineComment();
                        return self.next();
                    }
                    if (n == '=') {
                        self.advance();
                        return simple(.@"/=", start);
                    }
                }
                return simple(.@"/", start);
            },
            else => {
                if (isLineTerminator(c) or isWhiteSpace(c)) {
                    self.skipWhiteSpace();
                    return self.next();
                }
                if (isDigit(c)) return self.scanNumber();
                if (isIdStart(c)) return self.scanIdentifierOrKeyword();
                if (c == '\\') {
                    self.advance();
                    return self.scanIdentifierOrKeyword();
                }
                self.advance();
                return simple(.unknown, start);
            },
        }
    }

    fn simple(kind: TokenType, start: usize) Token {
        _ = start;
        return .{ .kind = kind, .loc = undefined, .slice = "" };
    }

    fn simpleWithSlice(comptime kind: TokenType, start: usize, slice: []const u8) Token {
        return .{ .kind = kind, .loc = undefined, .slice = slice };
    }

    pub fn next(self: *Self) Token {
        const start = self.pos;
        if (self.ch()) |c| {
            _ = c;
            // Handle hashbang at beginning
            if (self.pos == 0 and self.ch()) |c2| {
                if (c2 == '#') {
                    self.skipHashbang();
                    return self.next();
                }
            }
            return self.scanPunctuator();
        }
        return simple(.eof, self.pos);
    }

    fn keywordFromString(s: []const u8) TokenType {
        const keywords = std.ComptimeStringMap(TokenType, .{
            .{ "await", .await },
            .{ "break", .@"break" },
            .{ "case", .case },
            .{ "catch", .@"catch" },
            .{ "class", .class },
            .{ "const", .@"const" },
            .{ "continue", .@"continue" },
            .{ "debugger", .debugger },
            .{ "default", .default },
            .{ "delete", .delete },
            .{ "do", .do },
            .{ "else", .@"else" },
            .{ "enum", .@"enum" },
            .{ "export", .@"export" },
            .{ "extends", .extends },
            .{ "false", .false },
            .{ "finally", .finally },
            .{ "for", .@"for" },
            .{ "function", .function },
            .{ "if", .@"if" },
            .{ "import", .import },
            .{ "in", .in },
            .{ "instanceof", .instanceof },
            .{ "let", .let },
            .{ "new", .new },
            .{ "null", .null },
            .{ "return", .@"return" },
            .{ "super", .super },
            .{ "switch", .@"switch" },
            .{ "this", .this },
            .{ "throw", .throw },
            .{ "true", .true },
            .{ "try", .@"try" },
            .{ "typeof", .typeof },
            .{ "var", .@"var" },
            .{ "void", .void },
            .{ "while", .@"while" },
            .{ "with", .with },
            .{ "yield", .yield },
            .{ "using", .using },
        });

        if (keywords.get(s)) |k| return k;
        return .identifier;
    }

    pub fn tokenizeComptime(comptime source: []const u8) []const Token {
        comptime {
            return tokenize(source);
        }
    }
};