const std = @import("std");
const builtin = @import("builtin");
const Server = std.zig.Server.Message;
const Client = std.zig.Client.Message;
const Memory = struct { rss_after_bytes: ?usize = null, cumulative_peak_rss_bytes: ?usize = null };

fn send(io: std.Io, file: std.Io.File, tag: Client.Tag, body: []const u8, deadline: std.Io.Timeout) !void {
    var header: [8]u8 = undefined;
    std.mem.writeInt(u32, header[0..4], @intFromEnum(tag), .little);
    std.mem.writeInt(u32, header[4..8], @intCast(body.len), .little);
    for ([_][]const u8{ &header, body }) |part| {
        var offset: usize = 0;
        while (offset < part.len) {
            const n = try (try io.operateTimeout(.{ .file_write_streaming = .{
                .file = file,
                .data = &.{part[offset..]},
            } }, deadline)).file_write_streaming;
            if (n == 0) return error.EndOfStream;
            offset += n;
        }
    }
}

fn readAll(io: std.Io, file: std.Io.File, bytes: []u8, deadline: std.Io.Timeout) !void {
    var offset: usize = 0;
    while (offset < bytes.len) {
        const n = try (try io.operateTimeout(.{ .file_read_streaming = .{
            .file = file,
            .data = &.{bytes[offset..]},
        } }, deadline)).file_read_streaming;
        if (n == 0) return error.EndOfStream;
        offset += n;
    }
}

fn receive(io: std.Io, file: std.Io.File, gpa: std.mem.Allocator, deadline: std.Io.Timeout) !struct { tag: Server.Tag, body: []u8 } {
    var header: [8]u8 = undefined;
    try readAll(io, file, &header, deadline);
    const tag: Server.Tag = @enumFromInt(std.mem.readInt(u32, header[0..4], .little));
    const len = std.mem.readInt(u32, header[4..8], .little);
    if (len > 16 * 1024 * 1024) return error.MessageTooLarge;
    const bytes = try gpa.alloc(u8, len);
    errdefer gpa.free(bytes);
    try readAll(io, file, bytes, deadline);
    return .{ .tag = tag, .body = bytes };
}

fn wait(io: std.Io, child: *std.process.Child, deadline: std.Io.Timeout) !std.process.Child.Term {
    const Result = union(enum) { exited: std.process.Child.WaitError!std.process.Child.Term, timeout: std.Io.Cancelable!void };
    var results: [2]Result = undefined;
    var select = std.Io.Select(Result).init(io, &results);
    defer select.cancelDiscard();
    try select.concurrent(.exited, std.process.Child.wait, .{ child, io });
    select.async(.timeout, std.Io.Timeout.sleep, .{ deadline, io });
    return switch (try select.await()) {
        .exited => |result| result,
        .timeout => |result| {
            try result;
            return error.Timeout;
        },
    };
}

fn memory(io: std.Io, gpa: std.mem.Allocator, pid: std.process.Child.Id) !Memory {
    if (comptime builtin.os.tag != .linux) return .{};
    var path_buffer: [64]u8 = undefined;
    const path = try std.fmt.bufPrint(&path_buffer, "/proc/{d}/status", .{pid});
    const file = try std.Io.Dir.cwd().openFile(io, path, .{});
    defer file.close(io);
    var reader = file.readerStreaming(io, &.{});
    const status = try reader.interface.allocRemaining(gpa, .limited(8192));
    defer gpa.free(status);
    var value: Memory = .{};
    var lines = std.mem.splitScalar(u8, status, '\n');
    while (lines.next()) |line| {
        if (!std.mem.startsWith(u8, line, "VmRSS:") and !std.mem.startsWith(u8, line, "VmHWM:")) continue;
        var words = std.mem.tokenizeAny(u8, line, " \t");
        const kind = words.next().?;
        const bytes = try std.fmt.parseInt(usize, words.next() orelse return error.InvalidMemoryReport, 10);
        if (std.mem.eql(u8, kind, "VmRSS:")) value.rss_after_bytes = bytes * 1024 else value.cumulative_peak_rss_bytes = bytes * 1024;
    }
    return value;
}

fn queryTests(io: std.Io, child: *std.process.Child, gpa: std.mem.Allocator, deadline: std.Io.Timeout) ![]u8 {
    const version = try receive(io, child.stdout.?, gpa, deadline);
    defer gpa.free(version.body);
    if (version.tag != .zig_version or !std.mem.eql(u8, version.body, builtin.zig_version_string))
        return error.DifferentZigVersion;
    try send(io, child.stdin.?, .query_test_metadata, "", deadline);
    const metadata = try receive(io, child.stdout.?, gpa, deadline);
    errdefer gpa.free(metadata.body);
    if (metadata.tag != .test_metadata or metadata.body.len < 8) return error.InvalidMetadata;
    const strings = std.mem.readInt(u32, metadata.body[0..4], .little);
    const count = std.mem.readInt(u32, metadata.body[4..8], .little);
    const table_bytes = try std.math.mul(usize, count, 8);
    if (metadata.body.len != 8 + table_bytes + strings) return error.InvalidMetadata;
    return metadata.body;
}

fn testName(metadata: []const u8, index: usize) ![]const u8 {
    const count = std.mem.readInt(u32, metadata[4..8], .little);
    if (index >= count) return error.InvalidMetadata;
    const offset = std.mem.readInt(u32, metadata[8 + index * 4 ..][0..4], .little);
    const strings = metadata[8 + @as(usize, count) * 8 ..];
    if (offset >= strings.len) return error.InvalidMetadata;
    const end = std.mem.indexOfScalar(u8, strings[offset..], 0) orelse return error.InvalidMetadata;
    return strings[offset..][0..end];
}

fn checkFilter(io: std.Io, gpa: std.mem.Allocator, filter: []const u8, binaries: []const []const u8) !void {
    var matches: usize = 0;
    for (binaries) |binary| {
        var child = try std.process.spawn(io, .{
            .argv = &.{ binary, "--listen=-" },
            .stdin = .pipe,
            .stdout = .pipe,
            .stderr = .inherit,
        });
        errdefer child.kill(io);
        const deadline: std.Io.Timeout = (std.Io.Timeout{ .duration = .{ .clock = .awake, .raw = .fromSeconds(30) } }).toDeadline(io);
        const metadata = try queryTests(io, &child, gpa, deadline);
        defer gpa.free(metadata);
        const count = std.mem.readInt(u32, metadata[4..8], .little);
        for (0..count) |index| {
            const name = try testName(metadata, index);
            if (std.mem.indexOf(u8, name, ".test.") == null) if (std.mem.lastIndexOf(u8, name, ".test_")) |at| {
                const suffix = name[at + ".test_".len ..];
                const unnamed = for (suffix) |c| {
                    if (!std.ascii.isDigit(c)) break false;
                } else true;
                if (suffix.len > 0 and unnamed) continue;
            };
            if (std.mem.indexOf(u8, name, filter) != null) matches += 1;
        }
        try send(io, child.stdin.?, .exit, "", deadline);
        const result = try wait(io, &child, deadline);
        if (result != .exited or result.exited != 0) return error.TestRunnerFailed;
    }
    if (matches == 0) {
        std.debug.print("test filter '{s}' matched no named tests in {d} test program(s)\n", .{ filter, binaries.len });
        return error.NoMatchingTests;
    }
    std.debug.print("test filter '{s}': {d} named match(es) in {d} test program(s)\n", .{ filter, matches, binaries.len });
}

pub fn main(init: std.process.Init) !void {
    const io = init.io;
    const gpa = init.gpa;
    const args = try init.minimal.args.toSlice(gpa);
    defer gpa.free(args);
    if (args.len > 1 and std.mem.eql(u8, args[1], "--check-filter")) {
        if (args.len < 4) return error.ExpectedFilterAndTestBinaries;
        return checkFilter(io, gpa, args[2], args[3..]);
    }
    if (args.len < 2) return error.ExpectedTestBinaryAndOptionalFilter;
    var filter: ?[]const u8 = null;
    var timeout_ms: u32 = 30_000;
    for (args[2..]) |arg| {
        if (std.mem.startsWith(u8, arg, "--timeout-ms=")) {
            timeout_ms = try std.fmt.parseInt(u32, arg["--timeout-ms=".len..], 10);
            if (timeout_ms == 0) return error.InvalidTimeout;
        } else if (filter == null) {
            filter = arg;
        } else return error.ExpectedTestBinaryAndOptionalFilter;
    }
    const timeout: std.Io.Timeout = .{ .duration = .{ .clock = .awake, .raw = .fromMilliseconds(timeout_ms) } };
    const process_begin = std.Io.Clock.Timestamp.now(io, .awake);
    var child = try std.process.spawn(io, .{
        .argv = &.{ args[1], "--listen=-" },
        .stdin = .pipe,
        .stdout = .pipe,
        .stderr = .inherit,
        .request_resource_usage_statistics = true,
    });
    errdefer child.kill(io);
    var output_buffer: [4096]u8 = undefined;
    var output = std.Io.File.stdout().writerStreaming(io, &output_buffer);
    const startup_timeout: std.Io.Timeout = .{ .duration = .{ .clock = .awake, .raw = .fromMilliseconds(@max(timeout_ms, 1500)) } };
    const startup_deadline = startup_timeout.toDeadline(io);
    const metadata = try queryTests(io, &child, gpa, startup_deadline);
    defer gpa.free(metadata);
    const count = std.mem.readInt(u32, metadata[4..8], .little);
    var passed: usize = 0;
    var failed: usize = 0;
    var skipped: usize = 0;
    var request_to_result_ns: i128 = 0;
    var observed_peak_rss: ?usize = null;
    for (0..count) |i| {
        const name = try testName(metadata, i);
        if (filter) |needle| if (std.mem.indexOf(u8, name, needle) == null) continue;
        var id: [4]u8 = undefined;
        std.mem.writeInt(u32, &id, @intCast(i), .little);
        const begin = std.Io.Clock.Timestamp.now(io, .awake);
        const deadline: std.Io.Timeout = .{ .deadline = begin.addDuration(timeout.duration) };
        errdefer std.debug.print("unit-profile: running {s}\n", .{name});
        try send(io, child.stdin.?, .run_test, &id, deadline);
        const started = try receive(io, child.stdout.?, gpa, deadline);
        defer gpa.free(started.body);
        if (started.tag != .test_started or started.body.len != 0) return error.UnexpectedReply;
        const result = try receive(io, child.stdout.?, gpa, deadline);
        defer gpa.free(result.body);
        const elapsed = begin.untilNow(io).raw.nanoseconds;
        if (result.tag != .test_results or result.body.len != @sizeOf(Server.TestResults)) return error.UnexpectedReply;
        if (std.mem.readInt(u32, result.body[0..4], .little) != i) return error.UnexpectedTestIndex;
        const bits = std.mem.readInt(u64, result.body[4..12], .little);
        if (bits & 3 == 3) return error.InvalidTestStatus;
        const flags: Server.TestResults.Flags = @bitCast(bits);
        const bad = flags.status == .fail or flags.log_err_count != 0 or flags.leak_count != 0;
        if (bad) failed += 1 else if (flags.status == .skip) skipped += 1 else passed += 1;
        request_to_result_ns += elapsed;
        const usage = memory(io, gpa, child.id.?) catch Memory{};
        if (usage.cumulative_peak_rss_bytes) |peak| observed_peak_rss = @max(observed_peak_rss orelse 0, peak);
        try std.json.Stringify.value(.{
            .kind = "case",
            .name = name,
            .request_to_result_ns = elapsed,
            .rss_after_bytes = usage.rss_after_bytes,
            .cumulative_peak_rss_bytes = usage.cumulative_peak_rss_bytes,
            .status = if (bad) Server.TestResults.Status.fail else flags.status,
            .log_errors = flags.log_err_count,
            .leaks = flags.leak_count,
        }, .{}, &output.interface);
        try output.interface.writeByte('\n');
        try output.interface.flush();
    }
    const exit_deadline = timeout.toDeadline(io);
    try send(io, child.stdin.?, .exit, "", exit_deadline);
    const result = try wait(io, &child, exit_deadline);
    const process_wall_ns = process_begin.untilNow(io).raw.nanoseconds;
    if (child.resource_usage_statistics.getMaxRss()) |peak| observed_peak_rss = @max(observed_peak_rss orelse 0, peak);
    try std.json.Stringify.value(.{
        .kind = "summary",
        .passed = passed,
        .failed = failed,
        .skipped = skipped,
        .request_to_result_ns = request_to_result_ns,
        .process_wall_ns = process_wall_ns,
        .process_peak_rss_bytes = observed_peak_rss,
        .exit = result,
    }, .{}, &output.interface);
    try output.interface.writeByte('\n');
    try output.interface.flush();
    if (result != .exited or result.exited != 0) return error.TestRunnerFailed;
    if (passed + skipped + failed == 0) return error.NoMatchingTests;
    if (failed != 0) return error.TestFailed;
}