// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! Splits a program into tokens, one at a time, on demand.
//!
//! Modeled on `std.zig.Tokenizer`, and for the same reason: the language is a
//! superset of ZON, so anything ZON would tokenize one way has to come out the
//! same way here. What it adds is the interpolation inside a string. A string
//! token runs from its opening quote to its closing one, and `${ ... }` inside
//! it is skipped over by counting braces -- with any string or character
//! literal inside the braces skipped whole, so that `"${@join(xs, "}")}"` is
//! one token. The parser later runs a second tokenizer over just the inside of
//! each `${ }`, which is what `limit` is for.

const std = @import("std");

const Tokenizer = @This();

pub const Tag = enum {
    eof,
    invalid,

    identifier,
    /// `@"..."`, an identifier spelled as a string.
    quoted_identifier,
    /// `@name`, a built-in or host function.
    builtin,
    number,
    char_literal,
    string,
    /// One line of a `\\` multiline string, without its newline.
    multiline_string_line,

    period,
    /// `.{`, which opens every struct and list literal.
    period_brace,
    /// `.?`, which unwraps an optional.
    period_question,
    /// `...`, between the ends of a range in a `switch` case.
    ellipsis3,
    /// `{`, which opens the prongs of a `switch`.
    l_brace,
    r_brace,
    l_paren,
    r_paren,
    l_bracket,
    r_bracket,
    comma,
    semicolon,
    equal,
    equal_equal,
    /// `=>`, between a `switch` prong's cases and its value.
    equal_angle_bracket_right,
    bang,
    bang_equal,
    angle_bracket_left,
    angle_bracket_left_equal,
    angle_bracket_right,
    angle_bracket_right_equal,
    plus,
    plus_plus,
    minus,
    asterisk,
    slash,
    percent,
    pipe,

    keyword_let,
    keyword_if,
    keyword_switch,
    keyword_else,
    keyword_orelse,
    keyword_and,
    keyword_or,
    keyword_true,
    keyword_false,
    keyword_null,

    /// How the token reads in a message: the spelling for punctuation and
    /// keywords, a description for everything else.
    pub fn describe(tag: Tag) []const u8 {
        return switch (tag) {
            .eof => "end of input",
            .invalid => "invalid token",
            .identifier, .quoted_identifier => "an identifier",
            .builtin => "a builtin",
            .number => "a number",
            .char_literal => "a character literal",
            .string, .multiline_string_line => "a string",
            .period => "'.'",
            .period_brace => "'.{'",
            .period_question => "'.?'",
            .ellipsis3 => "'...'",
            .l_brace => "'{'",
            .r_brace => "'}'",
            .l_paren => "'('",
            .r_paren => "')'",
            .l_bracket => "'['",
            .r_bracket => "']'",
            .comma => "','",
            .semicolon => "';'",
            .equal => "'='",
            .equal_equal => "'=='",
            .equal_angle_bracket_right => "'=>'",
            .bang => "'!'",
            .bang_equal => "'!='",
            .angle_bracket_left => "'<'",
            .angle_bracket_left_equal => "'<='",
            .angle_bracket_right => "'>'",
            .angle_bracket_right_equal => "'>='",
            .plus => "'+'",
            .plus_plus => "'++'",
            .minus => "'-'",
            .asterisk => "'*'",
            .slash => "'/'",
            .percent => "'%'",
            .pipe => "'|'",
            .keyword_let => "'let'",
            .keyword_if => "'if'",
            .keyword_switch => "'switch'",
            .keyword_else => "'else'",
            .keyword_orelse => "'orelse'",
            .keyword_and => "'and'",
            .keyword_or => "'or'",
            .keyword_true => "'true'",
            .keyword_false => "'false'",
            .keyword_null => "'null'",
        };
    }
};

pub const Token = struct {
    tag: Tag,
    start: u32,
    end: u32,
};

const keywords = std.StaticStringMap(Tag).initComptime(.{
    .{ "let", .keyword_let },
    .{ "if", .keyword_if },
    .{ "switch", .keyword_switch },
    .{ "else", .keyword_else },
    .{ "orelse", .keyword_orelse },
    .{ "and", .keyword_and },
    .{ "or", .keyword_or },
    .{ "true", .keyword_true },
    .{ "false", .keyword_false },
    .{ "null", .keyword_null },
});

/// How deeply strings may nest inside interpolations inside strings. Skipping
/// them is recursive, and a fuzzer will happily write ten thousand of them.
pub const max_string_nesting = 32;

source: []const u8,
index: u32,
/// Tokenizing stops here, as if the source ended. The parser sets it to the
/// closing brace of an interpolation to tokenize what is inside.
limit: u32,

pub fn init(source: []const u8) Tokenizer {
    return .{ .source = source, .index = 0, .limit = @intCast(source.len) };
}

fn peekByte(self: *const Tokenizer, offset: u32) u8 {
    const i = self.index + offset;
    return if (i < self.limit) self.source[i] else 0;
}

pub fn next(self: *Tokenizer) Token {
    self.skipSpaceAndComments();
    const start = self.index;
    if (start >= self.limit) return .{ .tag = .eof, .start = start, .end = start };

    const c = self.source[start];
    self.index += 1;
    const tag: Tag = switch (c) {
        'a'...'z', 'A'...'Z', '_' => blk: {
            while (self.index < self.limit and isIdentChar(self.source[self.index])) self.index += 1;
            break :blk keywords.get(self.source[start..self.index]) orelse .identifier;
        },
        '0'...'9' => blk: {
            self.scanNumber(start);
            break :blk .number;
        },
        '"' => blk: {
            const end = scanString(self.source, self.index, self.limit, 0) orelse {
                self.index = self.lineEnd();
                break :blk .invalid;
            };
            self.index = end;
            break :blk .string;
        },
        '\'' => blk: {
            const end = scanChar(self.source, self.index, self.limit) orelse {
                self.index = self.lineEnd();
                break :blk .invalid;
            };
            self.index = end;
            break :blk .char_literal;
        },
        '\\' => blk: {
            if (self.peekByte(0) != '\\') break :blk .invalid;
            self.index = self.lineEnd();
            break :blk .multiline_string_line;
        },
        '@' => blk: {
            if (self.peekByte(0) == '"') {
                const end = scanString(self.source, self.index + 1, self.limit, max_string_nesting) orelse {
                    self.index = self.lineEnd();
                    break :blk .invalid;
                };
                self.index = end;
                break :blk .quoted_identifier;
            }
            const name_start = self.index;
            while (self.index < self.limit and isIdentChar(self.source[self.index])) self.index += 1;
            if (self.index == name_start) break :blk .invalid;
            break :blk .builtin;
        },
        '.' => switch (self.peekByte(0)) {
            '{' => self.take(1, .period_brace),
            '?' => self.take(1, .period_question),
            '.' => if (self.peekByte(1) == '.') self.take(2, .ellipsis3) else .period,
            else => .period,
        },
        '{' => .l_brace,
        '}' => .r_brace,
        '(' => .l_paren,
        ')' => .r_paren,
        '[' => .l_bracket,
        ']' => .r_bracket,
        ',' => .comma,
        ';' => .semicolon,
        '|' => .pipe,
        '*' => .asterisk,
        '/' => .slash,
        '%' => .percent,
        '-' => .minus,
        '=' => switch (self.peekByte(0)) {
            '=' => self.take(1, .equal_equal),
            '>' => self.take(1, .equal_angle_bracket_right),
            else => .equal,
        },
        '!' => if (self.peekByte(0) == '=') self.take(1, .bang_equal) else .bang,
        '<' => if (self.peekByte(0) == '=') self.take(1, .angle_bracket_left_equal) else .angle_bracket_left,
        '>' => if (self.peekByte(0) == '=') self.take(1, .angle_bracket_right_equal) else .angle_bracket_right,
        '+' => if (self.peekByte(0) == '+') self.take(1, .plus_plus) else .plus,
        else => .invalid,
    };
    return .{ .tag = tag, .start = start, .end = self.index };
}

fn take(self: *Tokenizer, n: u32, tag: Tag) Tag {
    self.index += n;
    return tag;
}

fn lineEnd(self: *const Tokenizer) u32 {
    var i = self.index;
    while (i < self.limit and self.source[i] != '\n') i += 1;
    return i;
}

fn skipSpaceAndComments(self: *Tokenizer) void {
    while (self.index < self.limit) {
        switch (self.source[self.index]) {
            ' ', '\t', '\r', '\n' => self.index += 1,
            '/' => {
                if (self.peekByte(1) != '/') return;
                self.index = self.lineEnd();
            },
            else => return,
        }
    }
}

/// A number, as Zig spells one: any run of letters, digits, `_` and `.` that
/// starts with a digit, with a sign allowed straight after an exponent. What
/// it means is decided later by `std.zig.parseNumberLiteral`; here it only
/// has to end in the right place. A `.` is taken only when a digit follows it,
/// so that `xs[0].name` is a number, a `]`, and a field access.
fn scanNumber(self: *Tokenizer, start: u32) void {
    const hex = self.source[start] == '0' and (self.peekByte(0) == 'x' or self.peekByte(0) == 'X');
    while (self.index < self.limit) {
        const c = self.source[self.index];
        if (isIdentChar(c)) {
            self.index += 1;
        } else if (c == '.' and std.ascii.isDigit(self.peekByte(1))) {
            self.index += 1;
        } else if (c == '.' and hex and std.ascii.isHex(self.peekByte(1))) {
            self.index += 1;
        } else if ((c == '+' or c == '-') and self.index > start) {
            const prev = self.source[self.index - 1];
            const exponent = if (hex) prev == 'p' or prev == 'P' else prev == 'e' or prev == 'E';
            if (!exponent) return;
            self.index += 1;
        } else return;
    }
}

fn isIdentChar(c: u8) bool {
    return std.ascii.isAlphanumeric(c) or c == '_';
}

/// Where the string whose content starts at `i` ends: the index just past its
/// closing quote, or null if it has none before a newline or `limit`. An
/// interpolation inside it is skipped whole. `depth` is how many strings this
/// one is already inside, and an interpolation is refused once that reaches
/// `max_string_nesting` -- which is also how a quoted identifier, `@"..."`,
/// which cannot be interpolated, is scanned.
pub fn scanString(source: []const u8, start: u32, limit: u32, depth: u32) ?u32 {
    var i = start;
    while (i < limit) {
        switch (source[i]) {
            '"' => return i + 1,
            '\n' => return null,
            '\\' => i += 2,
            '$' => {
                if (i + 1 < limit and source[i + 1] == '{') {
                    if (depth >= max_string_nesting) return null;
                    const close = skipInterpolation(source, i + 2, limit, depth + 1) orelse return null;
                    i = close + 1;
                } else i += 1;
            },
            else => i += 1,
        }
    }
    return null;
}

/// The index of the `}` that closes an interpolation whose inside starts at
/// `i`, or null if it is not closed before a newline or `limit`.
pub fn skipInterpolation(source: []const u8, start: u32, limit: u32, depth: u32) ?u32 {
    var i = start;
    var braces: u32 = 1;
    while (i < limit) {
        switch (source[i]) {
            '\n' => return null,
            '{' => {
                braces += 1;
                i += 1;
            },
            '}' => {
                braces -= 1;
                if (braces == 0) return i;
                i += 1;
            },
            '"' => i = scanString(source, i + 1, limit, depth) orelse return null,
            '\'' => i = scanChar(source, i + 1, limit) orelse return null,
            else => i += 1,
        }
    }
    return null;
}

/// The index just past the closing quote of a character literal whose content
/// starts at `i`.
fn scanChar(source: []const u8, start: u32, limit: u32) ?u32 {
    var i = start;
    while (i < limit) {
        switch (source[i]) {
            '\'' => return i + 1,
            '\n' => return null,
            '\\' => i += 2,
            else => i += 1,
        }
    }
    return null;
}

fn expectTokens(source: []const u8, expected: []const Tag) !void {
    var t: Tokenizer = .init(source);
    for (expected) |tag| {
        const tok = t.next();
        std.testing.expectEqual(tag, tok.tag) catch |err| {
            std.debug.print("at {d}..{d}: '{s}'\n", .{ tok.start, tok.end, source[tok.start..tok.end] });
            return err;
        };
    }
    try std.testing.expectEqual(Tag.eof, t.next().tag);
}

test "punctuation and keywords" {
    try expectTokens(
        \\let x = .{ .a = 1 } ++ ctx.y.?; // trailing comment
        \\if (a and !b or c != d) e orelse f
    , &.{
        .keyword_let, .identifier,      .equal,      .period_brace,   .period,     .identifier,
        .equal,       .number,          .r_brace,    .plus_plus,      .identifier, .period,
        .identifier,  .period_question, .semicolon,  .keyword_if,     .l_paren,    .identifier,
        .keyword_and, .bang,            .identifier, .keyword_or,     .identifier, .bang_equal,
        .identifier,  .r_paren,         .identifier, .keyword_orelse, .identifier,
    });
}

test "the punctuation of a switch" {
    try expectTokens("switch (x) { 1...5, .a => |v| v, else => 0 } 1..2", &.{
        .keyword_switch, .l_paren,                   .identifier,   .r_paren,                   .l_brace,
        .number,         .ellipsis3,                 .number,       .comma,                     .period,
        .identifier,     .equal_angle_bracket_right, .pipe,         .identifier,                .pipe,
        .identifier,     .comma,                     .keyword_else, .equal_angle_bracket_right, .number,
        .r_brace,        .number,                    .period,       .period,                    .number,
    });
}

test "numbers end in the right place" {
    try expectTokens("1_000 0x1F 1.5e-3 0x1.8p+1 xs[0].a 1+2", &.{
        .number,     .number,     .number, .number,
        .identifier, .l_bracket,  .number, .r_bracket,
        .period,     .identifier, .number, .plus,
        .number,
    });
}

test "an interpolation with a string inside it is one token" {
    try expectTokens(
        \\"a${@join(xs, "}")}b" "${"${"x"}"}" '}'
    , &.{ .string, .string, .char_literal });
}

test "unterminated strings and interpolations are invalid" {
    try expectTokens("\"abc", &.{.invalid});
    try expectTokens("\"a${b\"", &.{.invalid});
    try expectTokens("\"a\nb\"", &.{ .invalid, .identifier, .invalid });
}

test "builtins, quoted identifiers and multiline strings" {
    try expectTokens(
        \\@lower(@"a b")
        \\\\line one
        \\\\line ${two}
    , &.{ .builtin, .l_paren, .quoted_identifier, .r_paren, .multiline_string_line, .multiline_string_line });
}

test "limit stops tokenizing early" {
    var t: Tokenizer = .init("a b c");
    t.limit = 3;
    try std.testing.expectEqual(Tag.identifier, t.next().tag);
    try std.testing.expectEqual(Tag.identifier, t.next().tag);
    try std.testing.expectEqual(Tag.eof, t.next().tag);
}

test "string nesting is bounded" {
    var src: std.ArrayList(u8) = .empty;
    defer src.deinit(std.testing.allocator);
    for (0..max_string_nesting + 2) |_| try src.appendSlice(std.testing.allocator, "\"${");
    for (0..max_string_nesting + 2) |_| try src.appendSlice(std.testing.allocator, "}\"");
    try expectTokens(src.items, &.{.invalid});
}