const std = @import("../std.zig");
const assert = std.debug.assert;

pub const ParseError = error{
    OutOfMemory,
    InvalidLiteral,
};

pub const Base = enum(u8) { decimal = 10, hex = 16, binary = 2, octal = 8 };
pub const FloatBase = enum(u8) { decimal = 10, hex = 16 };

pub const Result = union(enum) {
    /// Result fits if it fits in u64
    int: u64,
    /// Result is an int that doesn't fit in u64. Payload is the base, if it is
    /// not `.decimal` then the slice has a two character prefix.
    big_int: Base,
    /// Result is a float. Payload is the base, if it is not `.decimal` then
    /// the slice has a two character prefix.
    float: FloatBase,
    failure: Error,
};

pub const Error = union(enum) {
    /// The number has leading zeroes.
    leading_zero,
    /// Expected a digit after base prefix.
    digit_after_base,
    /// The base prefix is in uppercase.
    upper_case_base: usize,
    /// Float literal has an invalid base prefix.
    invalid_float_base: usize,
    /// Repeated '_' digit separator.
    repeated_underscore: usize,
    /// '_' digit separator after special character (+-.)
    invalid_underscore_after_special: usize,
    /// Invalid digit for the specified base.
    invalid_digit: struct { i: usize, base: Base },
    /// Invalid digit for an exponent.
    invalid_digit_exponent: usize,
    /// Float literal has multiple periods.
    duplicate_period,
    /// Float literal has multiple exponents.
    duplicate_exponent: usize,
    /// Exponent comes directly after '_' digit separator.
    exponent_after_underscore: usize,
    /// Special character (+-.) comes directly after underscore.
    special_after_underscore: usize,
    /// Number ends in special character (+-.)
    trailing_special: usize,
    /// Number ends in '_' digit separator.
    trailing_underscore: usize,
    /// Character not in [0-9a-zA-Z.+-_]
    invalid_character: usize,
    /// [+-] not immediately after [pPeE]
    invalid_exponent_sign: usize,
    /// Period comes after exponent.
    period_after_exponent: usize,
};

/// Parse Zig number literal accepted by fmt.parseInt, fmt.parseFloat and big_int.setString.
/// Valid for any number_literal token bytes.
pub fn parseNumberLiteral(bytes: []const u8) Result {
    // This is enforced by the tokenizer.
    assert(bytes.len > 0 and std.ascii.isDigit(bytes[0]));
    var i: usize = 0;
    var base: u8 = 10;
    if (bytes.len >= 2 and bytes[0] == '0') switch (bytes[1]) {
        'b' => {
            base = 2;
            i = 2;
        },
        'o' => {
            base = 8;
            i = 2;
        },
        'x' => {
            base = 16;
            i = 2;
        },
        'B', 'O', 'X' => return .{ .failure = .{ .upper_case_base = 1 } },
        '.', 'e', 'E' => {},
        else => return .{ .failure = .leading_zero },
    };
    if (bytes.len == 2 and base != 10) return .{ .failure = .digit_after_base };

    var x: u64 = 0;
    var overflow = false;
    var underscore = false;
    var period = false;
    var special: u8 = 0;
    var exponent = false;
    var float = false;
    while (i < bytes.len) : (i += 1) {
        const c = bytes[i];
        switch (c) {
            '_' => {
                if (i == 2 and base != 10) return .{ .failure = .{ .invalid_underscore_after_special = i } };
                if (special != 0) return .{ .failure = .{ .invalid_underscore_after_special = i } };
                if (underscore) return .{ .failure = .{ .repeated_underscore = i } };
                underscore = true;
                continue;
            },
            'e', 'E' => if (base == 10) {
                float = true;
                if (exponent) return .{ .failure = .{ .duplicate_exponent = i } };
                if (underscore) return .{ .failure = .{ .exponent_after_underscore = i } };
                special = c;
                exponent = true;
                continue;
            },
            'p', 'P' => if (base == 16) {
                if (i == 2) {
                    return .{ .failure = .{ .digit_after_base = {} } };
                }
                float = true;
                if (exponent) return .{ .failure = .{ .duplicate_exponent = i } };
                if (underscore) return .{ .failure = .{ .exponent_after_underscore = i } };
                special = c;
                exponent = true;
                continue;
            },
            '.' => {
                if (exponent) return .{ .failure = .{ .period_after_exponent = i } };
                float = true;
                if (base != 10 and base != 16) return .{ .failure = .{ .invalid_float_base = 1 } };
                if (period) return .{ .failure = .duplicate_period };
                period = true;
                if (underscore) return .{ .failure = .{ .special_after_underscore = i } };
                special = c;
                continue;
            },
            '+', '-' => {
                switch (special) {
                    'p', 'P' => {},
                    'e', 'E' => if (base != 10) return .{ .failure = .{ .invalid_exponent_sign = i } },
                    else => return .{ .failure = .{ .invalid_exponent_sign = i } },
                }
                special = c;
                continue;
            },
            else => {},
        }
        const digit = switch (c) {
            '0'...'9' => c - '0',
            'A'...'Z' => c - 'A' + 10,
            'a'...'z' => c - 'a' + 10,
            else => return .{ .failure = .{ .invalid_character = i } },
        };
        if (digit >= base) return .{ .failure = .{ .invalid_digit = .{ .i = i, .base = @as(Base, @fromBackingInt(@intCast(base))) } } };
        if (exponent and digit >= 10) return .{ .failure = .{ .invalid_digit_exponent = i } };
        underscore = false;
        special = 0;

        if (float) continue;
        if (x != 0) {
            const res = @mulWithOverflow(x, base);
            if (res[1] != 0) overflow = true;
            x = res[0];
        }
        const res = @addWithOverflow(x, digit);
        if (res[1] != 0) overflow = true;
        x = res[0];
    }
    if (underscore) return .{ .failure = .{ .trailing_underscore = bytes.len - 1 } };
    if (special != 0) return .{ .failure = .{ .trailing_special = bytes.len - 1 } };

    if (float) return .{ .float = @as(FloatBase, @fromBackingInt(@intCast(base))) };
    if (overflow) return .{ .big_int = @as(Base, @fromBackingInt(@intCast(base))) };
    return .{ .int = x };
}

test parseNumberLiteral {
    try std.testing.expectEqual(Result{ .float = .decimal }, parseNumberLiteral("3E2"));
    try std.testing.expectEqual(Result{ .int = 0x3E2 }, parseNumberLiteral("0x3E2"));
    try std.testing.expectEqual(Result{ .float = .hex }, parseNumberLiteral("0x3p2"));
    try std.testing.expectEqual(Result{ .failure = .{ .period_after_exponent = 3 } }, parseNumberLiteral("3E2.5"));
    try std.testing.expectEqual(Result{ .failure = .{ .period_after_exponent = 2 } }, parseNumberLiteral("3E.5"));
    try std.testing.expectEqual(Result{ .failure = .{ .period_after_exponent = 3 } }, parseNumberLiteral("3E1."));
    try std.testing.expectEqual(Result{ .failure = .{ .invalid_float_base = 1 } }, parseNumberLiteral("0o3.1"));
    try std.testing.expectEqual(Result{ .failure = .{ .invalid_digit = .{ .i = 3, .base = .octal } } }, parseNumberLiteral("0o3e1"));
}

/// Returns an error if `parseNumberLiteral` returns `.float` but `parseFloat` fails.
/// AstGen relies on `parseFloat` being unable to fail after calling `parseNumberLiteral`.
fn checkFloat(bytes: []const u8) !void {
    // Number literals must start with a digit
    if (bytes.len == 0 or !std.ascii.isDigit(bytes[0])) return;

    switch (parseNumberLiteral(bytes)) {
        .float => {
            _ = try std.fmt.parseFloat(f128, bytes);
        },
        else => {},
    }
}

test "parseNumberLiteral float validation" {
    const Context = struct {
        fn testOne(_: @This(), smith: *std.testing.Smith) anyerror!void {
            var buf: [256]u8 = undefined;
            const bytes = buf[0..smith.slice(&buf)];
            try checkFloat(bytes);
        }
    };
    return std.testing.fuzz(Context{}, Context.testOne, .{});
}
