const std = @import("std");
const root = @import("../main.zig");
const BenchmarkArgs = root.argument.BenchmarkArgs;
const recognizer = root.recognizer;
const Pex = @import("../pex/root.zig");

fn write_result(
    writer: anytype,
    is_tty: bool,
    name: []const u8,
    index: usize,
    result: recognizer.Result,
    elapsed: u64,
    throughput: f64,
) !void {
    if (is_tty) {
        try writer.print("{s}[{}] {s}\x1b[0m ({:.2}, {d:.2} tps, {})\n", .{
            name,
            index,
            if (result.is_accepted) "\x1b[32maccept"
            else "\x1b[31mreject",
            std.fmt.fmtDuration(elapsed),
            throughput,
            result.info,
        });
    } else {
        try writer.print("{s}[{}] {s} ({:.2}, {d:.2} tps, {})\n", .{
            name,
            index,
            if (result.is_accepted) "accept"
            else "reject",
            std.fmt.fmtDuration(elapsed),
            throughput,
            result.info,
        });
    }
}

fn write_result_as_csv(
    writer: anytype,
    index: usize,
    input_length: usize,
    result: recognizer.Result,
    elapsed: u64,
    throughput: f64,
) !void {
    try writer.print("{},{},{d:.2},{d:.2},{:.2},{}\n", .{
        index,
        input_length,
        elapsed,
        throughput,
        result.info.max_heap_size,
        result.info.number_of_nodes,
    });
}

fn write_summary(
    writer: anytype,
    is_tty: bool,
    total: usize,
    accepted: usize,
    total_time: u64,
    throughput: struct { f64, f64, f64 },
) !void {
    const min, const max, const avg = throughput;

    if (is_tty) {
        try writer.print("finished {} ({} \x1b[32maccepted\x1b[0m {} \x1b[31mrejected\x1b[0m)", .{
            total,
            accepted,
            total - accepted,
        });
    } else {
        try writer.print("finished {} ({} accepted {} rejected)", .{
            total,
            accepted,
            total - accepted,
        });
    }

    try writer.print(" in {} (min: {d:.5} tps, max: {d:.2} tps, average: {d:.2} tps)\n", .{
        std.fmt.fmtDuration(total_time),
        min,
        max,
        avg,
    });
}

fn min_max_avg(values: []f64) struct { f64, f64, f64 } {
    var min: f64 = -1;
    var max: f64 = 0;
    var sum: f64 = 0;

    for (values) |value| {
        min = if (min < 0) value else @min(min, value);
        max = @max(max, value);
        sum += value;
    }

    return .{ min, max, sum / @as(f64, @floatFromInt(values.len)) };
}

pub fn benchmark(args: *BenchmarkArgs, allocator: std.mem.Allocator) !void {
    const stdout = std.io.getStdOut();
    const writer = stdout.writer();

    if (args.csv) {
        try writer.print("index,input length,duration [ns],throughput,max heap size [bytes],number of nodes\n", .{});
    }

    var bufreader = std.io.bufferedReader(args.input.file.reader());
    var reader = bufreader.reader();
    var index: usize = 0;
    const stderr = std.io.getStdErr();

    var throughputs = std.ArrayList(f64).init(allocator);
    defer throughputs.deinit();

    if (args.input.file.isTty()) {
        try stderr.writeAll("> ");
    }

    var read_arena = std.heap.ArenaAllocator.init(allocator);
    defer read_arena.deinit();
    var number_of_accepted: usize = 0;
    var total_time: u64 = 0;
    var timer = try std.time.Timer.start();

    const pex: Pex = try .compile(&args.grammar, allocator);

    while (try reader.readUntilDelimiterOrEofAlloc(
            read_arena.allocator(),
            '\n',
            std.math.maxInt(usize)
    )) |buffer| {

        const trimmed = std.mem.trim(u8, buffer, &std.ascii.whitespace);

        timer.reset();
        const result = try recognizer.check(
            &pex,
            trimmed,
            allocator
        );

        const elapsed = timer.read();
        const throughput = @as(f64, @floatFromInt(trimmed.len)) / (@as(f64, @floatFromInt(elapsed)) / std.time.ns_per_s);
        try throughputs.append(throughput);

        if (args.csv) {
            try write_result_as_csv(
                writer,
                index,
                trimmed.len,
                result,
                elapsed,
                throughput,
            );
        } else {
            try write_result(
                writer,
                stdout.isTty(),
                args.input.name,
                index,
                result,
                elapsed,
                throughput,
            );
        }

        index += 1;
        number_of_accepted += @intFromBool(result.is_accepted);
        total_time += elapsed;

        if (args.input.file.isTty()) {
            try stderr.writeAll("> ");
        }
    }

    if (!args.csv) {
        try write_summary(
            writer,
            stdout.isTty(),
            index,
            number_of_accepted,
            total_time,
            min_max_avg(throughputs.items)
        );
    }
}