//! Test runner for parameterized system tests
const ArrayList = std.ArrayList;
const ChildProcess = std.ChildProcess;
const LogLevel = std.log.Level;
const Progress = std.Progress;
const Server = std.zig.Server;
const allocator = std.testing.allocator;
const allocator_ref = &std.testing.allocator_instance;
const argsWithAllocator = std.process.argsWithAllocator;
const assert = std.debug.assert;
const builtin = @import("builtin");
const concat = std.mem.concat;
const cwd = std.fs.cwd;
const dumpStackTrace = std.debug.dumpStackTrace;
const endsWith = std.mem.endsWith;
const eql = std.mem.eql;
const exit = std.process.exit;
const expectEqual = std.testing.expectEqual;
const expectEqualStrings = std.testing.expectEqualStrings;
const getStdIn = std.io.getStdIn;
const getStdOut = std.io.getStdOut;
const joinPath = std.fs.path.join;
const log_level = &std.testing.log_level;
const parseTarget = std.zig.CrossTarget.parse;
const print = std.debug.print;
const startsWith = std.mem.startsWith;
const std = @import("std");
const target = builtin.target;
const test_io_mode = builtin.test_io_mode;
const tmpDir = std.testing.tmpDir;
const tokenizeScalar = std.mem.tokenizeScalar;
const trimRight = std.mem.trimRight;
const zig_version_string = builtin.zig_version_string;

/// Options override.
pub const std_options = struct {
    pub const io_mode = test_io_mode;
    pub const logFn = struct {
        pub fn log(comptime message_level: LogLevel,
                   comptime scope: @Type(.EnumLiteral),
                   comptime format: []const u8, args: anytype) void {
            const fmt = "[" ++ @tagName(scope) ++ "] ("
                ++ @tagName(message_level) ++ "): " ++ format ++ "\n";
            if (@intFromEnum(message_level) <= @intFromEnum(log_level.*))
                print(fmt, args);
        }
    }.log;
};

/// Print usage and exit with an error status.
fn printUsage() noreturn {
    print("usage: zig build test -- [-r QBE] [FILE.ssa]... [DIR]...", .{});
    exit(1);
}

/// QBE implementation to be tested.
var bin: []const u8 = undefined;
/// Reference implementation.
var bin_ref: ?[]const u8 = null;
/// Zig as a C compiler.
var zig: []const u8 = undefined;
/// Target name in Zig/Clang format.
var zig_target: []const u8 = undefined;
/// Target name in QBE format.
var qbe_target: []const u8 = undefined;
/// External executor for our implementation and output programs.
var emulator: ?[]const u8 = null;

/// Parse command-line arguments.
fn parseArgs() ![]const []const u8 {
    var listen = false;
    defer if (!listen)
        @panic("test not run in server mode");
    var paths = ArrayList([]const u8).init(allocator);
    errdefer paths.deinit();

    var args = try argsWithAllocator(allocator);
    _ = args.next().?; // path to test binary
    while (args.next()) |arg| {
        if (eql(u8, arg, "--listen=-")) {
            listen = true;
        } else if (eql(u8, arg, "-b")) {
            bin = try allocator.dupe(u8, args.next().?);
        } else if (eql(u8, arg, "-r")) {
            bin_ref = try allocator.dupe(u8, args.next().?);
        } else if (eql(u8, arg, "-z")) {
            zig = try allocator.dupe(u8, args.next().?);
        } else if (eql(u8, arg, "-t")) {
            const triple = args.next().?;
            zig_target = try concat(allocator, u8, &.{ "--target=", triple });
            const cross_target = try parseTarget(.{ .arch_os_abi = triple });
            const os_tag = cross_target.os_tag orelse target.os.tag;
            const cpu_arch = cross_target.cpu_arch orelse target.cpu.arch;
            qbe_target = switch (os_tag) {
                .macos => switch (cpu_arch) {
                    .aarch64 => "arm64_apple",
                    .x86_64 => "amd64_apple",
                    else => unreachable,
                },
                else => switch (cpu_arch) {
                    .aarch64 => "arm64",
                    .riscv64 => "rv64",
                    .x86_64 => "amd64_sysv",
                    else => unreachable,
                },
            };
        } else if (eql(u8, arg, "-e")) {
            emulator = try allocator.dupe(u8, args.next().?);
        } else if (endsWith(u8, arg, ".ssa")) {
            const file = cwd().openFile(arg, .{}) catch |err| {
                print("cannot open '{s}': {}\n", .{ arg, err });
                printUsage();
            };
            defer file.close();
            try paths.append(try allocator.dupe(u8, arg));
        } else {
            var parent = cwd().openIterableDir(arg, .{}) catch |err| {
                print("cannot open '{s}': {}\n", .{ arg, err });
                printUsage();
            };
            defer parent.close();
            var it = parent.iterate();
            while (try it.next()) |file| {
                if (file.kind != .file) continue;
                if (startsWith(u8, file.name, "_")) continue;
                if (!endsWith(u8, file.name, ".ssa")) continue;
                try paths.append(try joinPath(allocator, &.{ arg, file.name }));
            }
        }
    }

    if (paths.items.len == 0) {
        print("no test to run.\n", .{});
        printUsage();
    }
    return paths.toOwnedSlice();
}

/// Allocate and initialize every item to 0.
fn calloc(comptime T: type, n: usize) ![]T {
    const slice = try allocator.alloc(T, n);
    for (slice) |*item|
        item.* = 0;
    return slice;
}

/// Extract embedded test driver and output.
fn extract(ssa: []const u8, drv: []const u8, out: []const u8) !struct {
    driver: bool,
    output: bool,
} {
    const dir = cwd();
    const source = try dir.openFile(ssa, .{});
    defer source.close();
    const driver = try dir.createFile(drv, .{});
    defer driver.close();
    const output = try dir.createFile(out, .{});
    defer output.close();
    var has_driver = false;
    var has_output = false;

    const ir = try source.readToEndAlloc(allocator, 1 << 20);
    defer allocator.free(ir);
    var lines = tokenizeScalar(u8, ir, '\n');
    while (lines.next()) |line|
        if (startsWith(u8, line, "# skip ")) {
            var skipped = tokenizeScalar(u8, line["# skip ".len..], ' ');
            while (skipped.next()) |s|
                if (eql(u8, s, qbe_target))
                    return error.SkipZigTest;
        } else if (startsWith(u8, line, "# >>> ")) {
            const writer = if (eql(u8, line, "# >>> driver")) w: {
                has_driver = true;
                break :w driver.writer();
            } else if (eql(u8, line, "# >>> output")) w: {
                has_output = true;
                break :w output.writer();
            } else continue;
            while (lines.next()) |comment| {
                if (eql(u8, comment, "# <<<"))
                    break;
                if (!startsWith(u8, comment, "# "))
                    return error.InvalidDocTest;
                try writer.print("{s}\n", .{
                    trimRight(u8, comment["# ".len..], "#"),
                });
            }
        };
    return .{ .driver = has_driver, .output = has_output };
}

/// Spawn given command and return the collected stdout.
fn collectOutput(argv: []const []const u8) ![]const u8 {
    const result = try ChildProcess.exec(.{
        .allocator = allocator,
        .argv = argv,
        .max_output_bytes = 1 << 20,
    });
    allocator.free(result.stderr);
    return result.stdout;
}

/// Spawn given command and expect a non-error exit status.
fn spawn(argv: []const []const u8) !void {
    var proc = ChildProcess.init(argv, allocator);
    proc.stderr_behavior = .Ignore;
    try expectEqual(ChildProcess.Term{ .Exited = 0 }, try proc.spawnAndWait());
}

/// Count the number of lines in a buffer and free it.
fn loc(buffer: []const u8) u32 {
    var n: u32 = 0;
    for (buffer) |c| {
        if (c == '\n')
            n += 1;
    } else allocator.free(buffer);
    return n;
}

/// Compile given QBE IR and check its exit status and (optionally) stdout.
fn check(ssa: []const u8) anyerror!void {
    var tmp = tmpDir(.{});
    defer tmp.cleanup();
    const dir = try tmp.dir.realpathAlloc(allocator, ".");
    defer allocator.free(dir);
    const drv = try joinPath(allocator, &.{ dir, "driver.c" });
    defer allocator.free(drv);
    const out = try joinPath(allocator, &.{ dir, "output" });
    defer allocator.free(out);
    const extracted = try extract(ssa, drv, out);

    const s = try joinPath(allocator, &.{ dir, "a.s" });
    defer allocator.free(s);
    const ir_loc = b: {
        const ir = try collectOutput(&.{ bin, "-t", qbe_target, ssa });
        errdefer allocator.free(ir);
        try tmp.parent_dir.writeFile(s, ir);
        break :b loc(ir);
    };
    if (bin_ref) |qbe| {
        const loc_diff = @as(i64, ir_loc) - loc(try collectOutput(&.{
            qbe, "-t", qbe_target, ssa
        }));
        if (loc_diff != 0)
            print("'{s}': {:1} lines asm\n", .{ ssa, loc_diff });
    }

    const exe = try joinPath(allocator, &.{ dir, "a.out" });
    defer allocator.free(exe);
    try spawn(if (extracted.driver)
        &.{ zig, "cc", zig_target, "-g", "-o", exe, drv, s }
    else
        &.{ zig, "cc", zig_target, "-g", "-o", exe, s });

    const run_argv: []const []const u8 = if (emulator) |executor|
        &.{ executor, exe, "a", "b", "c" }
    else
        &.{ exe, "a", "b", "c" };
    if (extracted.output) {
        const expected = try tmp.dir.readFileAlloc(allocator, out, 1 << 20);
        defer allocator.free(expected);
        const actual = try collectOutput(run_argv);
        defer allocator.free(actual);
        try expectEqualStrings(expected, actual);
    } else try spawn(run_argv);
}

/// Collect paths to test QBE IR files and feed them to the fn check
/// instead of running builtin.test_functions.
pub fn main() !void {
    const paths = try parseArgs();
    defer {
        for (paths) |ssa| allocator.free(ssa);
        allocator.free(paths);
        // Global variables
        allocator.free(bin);
        if (bin_ref) |string| allocator.free(string);
        allocator.free(zig);
        allocator.free(zig_target);
        if (emulator) |string| allocator.free(string);
        assert(allocator_ref.deinit() == .ok);
    }
    var server = try Server.init(.{
        .gpa = allocator,
        .in = getStdIn(),
        .out = getStdOut(),
        .zig_version = zig_version_string,
    });
    defer server.deinit();

    while (server.receiveMessage()) |message| {
        switch (message.tag) {
            .query_test_metadata => {
                const parent_allocator = allocator_ref.*;
                allocator_ref.* = .{};
                defer allocator_ref.* = parent_allocator;
                defer assert(allocator_ref.deinit() == .ok);

                var string_bytes = ArrayList(u8).init(allocator);
                defer string_bytes.deinit();
                const names = try allocator.alloc(u32, paths.len);
                defer allocator.free(names);
                for (paths, names) |ssa, *name| {
                    name.* = @intCast(string_bytes.items.len);
                    try string_bytes.ensureUnusedCapacity(ssa.len + 1);
                    string_bytes.appendSliceAssumeCapacity(ssa);
                    string_bytes.appendAssumeCapacity(0);
                }

                const async_frame_sizes = try calloc(u32, paths.len);
                defer allocator.free(async_frame_sizes);
                const expected_panic_msgs = try calloc(u32, paths.len);
                defer allocator.free(expected_panic_msgs);

                try server.serveTestMetadata(.{
                    .names = names,
                    .async_frame_sizes = async_frame_sizes,
                    .expected_panic_msgs = expected_panic_msgs,
                    .string_bytes = string_bytes.items,
                });
            },
            .run_test => {
                const parent_allocator = allocator_ref.*;
                allocator_ref.* = .{};
                defer allocator_ref.* = parent_allocator;

                const index = try server.receiveBody_u32();
                var fail = false;
                var skip = false;
                check(paths[index]) catch |err| switch (err) {
                    error.SkipZigTest => skip = true,
                    else => {
                        fail = true;
                        if (@errorReturnTrace()) |trace|
                            dumpStackTrace(trace.*);
                    },
                };
                try server.serveTestResults(.{
                    .index = index,
                    .flags = .{
                        .fail = fail,
                        .skip = skip,
                        .leak = allocator_ref.deinit() == .leak,
                    },
                });
            },
            .exit => break,
            else => |tag| {
                print("unsupported message: {}", .{ tag });
                exit(1);
            },
        }
    } else |err| return err;
}