//! loadgen — websocket load generator for stream.waow.tech
//!
//! holds N concurrent /subscribe connections from one process and reports
//! aggregate throughput. scenarios are knob combinations:
//!
//!   hold          --conns=200                          (steady fan-out load)
//!   connect storm --conns=500 --ramp-ms=0              (all dials at once)
//!   churn         --conns=100 --lifetime-s=20          (connect/disconnect loop)
//!   slow reader   --conns=10 --slow-ms=200             (backpressure via stalled reads)
//!   replay        --conns=20 --rewind-s=3600           (archive-tail catch-up reads)

const std = @import("std");
const websocket = @import("websocket");

const Io = std.Io;

const Config = struct {
    host: []const u8 = "stream.waow.tech",
    port: u16 = 443,
    conns: usize = 1,
    duration_s: u64 = 0, // 0 = run until killed
    lifetime_s: u64 = 0, // per-connection lifetime; 0 = hold
    slow_ms: u64 = 0, // sleep per message (slow-reader scenario)
    rewind_s: u64 = 0, // cursor = now_us - rewind (replay scenario)
    ramp_ms: u64 = 20, // delay between connection starts
    report_s: u64 = 10,
    collections: []const u8 = "", // comma-separated wantedCollections
    path: []const u8 = "/subscribe",
};

const Stats = struct {
    connects: std.atomic.Value(u64) = .init(0),
    disconnects: std.atomic.Value(u64) = .init(0),
    errors: std.atomic.Value(u64) = .init(0),
    active: std.atomic.Value(u64) = .init(0),
    messages: std.atomic.Value(u64) = .init(0),
    bytes: std.atomic.Value(u64) = .init(0),
};

var stats: Stats = .{};
var stop = std.atomic.Value(bool).init(false);

pub fn main(init: std.process.Init.Minimal) !void {
    const gpa = std.heap.c_allocator;

    var threaded: Io.Threaded = .init(gpa, .{ .stack_size = 1024 * 1024 });
    defer threaded.deinit();
    const io = threaded.io();

    const cfg = try parseArgs(init);

    std.log.info("loadgen: {d} conns -> wss://{s}:{d}/subscribe (lifetime={d}s slow={d}ms rewind={d}s ramp={d}ms)", .{
        cfg.conns, cfg.host, cfg.port, cfg.lifetime_s, cfg.slow_ms, cfg.rewind_s, cfg.ramp_ms,
    });

    const workers = try gpa.alloc(std.Thread, cfg.conns);
    for (workers, 0..) |*w, i| {
        w.* = try std.Thread.spawn(
            .{ .stack_size = 2 * 1024 * 1024 },
            connectionWorker,
            .{ io, gpa, cfg, i },
        );
        if (cfg.ramp_ms > 0) try io.sleep(.fromMilliseconds(@intCast(cfg.ramp_ms)), .awake);
    }

    const reporter = try std.Thread.spawn(.{}, reportLoop, .{ io, cfg });

    if (cfg.duration_s > 0) {
        try io.sleep(.fromSeconds(@intCast(cfg.duration_s)), .awake);
        stop.store(true, .monotonic);
        report(io, "final");
        std.process.exit(0);
    }
    reporter.join(); // runs forever
}

fn connectionWorker(io: Io, allocator: std.mem.Allocator, cfg: Config, index: usize) void {
    var backoff_s: u64 = 1;
    while (!stop.load(.monotonic)) {
        runConnection(io, allocator, cfg) catch |err| {
            _ = stats.errors.fetchAdd(1, .monotonic);
            std.log.debug("conn {d}: {s}", .{ index, @errorName(err) });
            io.sleep(.fromSeconds(@intCast(backoff_s)), .awake) catch return;
            backoff_s = @min(backoff_s * 2, 30);
            continue;
        };
        backoff_s = 1;
        // clean end of a lifetime (churn) — brief pause, then reconnect
        if (cfg.lifetime_s > 0 and !stop.load(.monotonic)) {
            io.sleep(.fromMilliseconds(250), .awake) catch return;
        }
    }
}

fn runConnection(io: Io, allocator: std.mem.Allocator, cfg: Config) !void {
    var path_buf: [2048]u8 = undefined;
    const path = try buildPath(io, &path_buf, cfg);

    var client = try websocket.Client.init(io, allocator, .{
        .host = cfg.host,
        .port = cfg.port,
        .tls = cfg.port == 443,
        .max_size = 1024 * 1024,
        .buffer_size = 64 * 1024,
    });
    defer client.deinit();

    var host_header_buf: [256]u8 = undefined;
    const host_header = try std.fmt.bufPrint(&host_header_buf, "Host: {s}\r\n", .{cfg.host});
    try client.handshake(path, .{ .headers = host_header });

    _ = stats.connects.fetchAdd(1, .monotonic);
    _ = stats.active.fetchAdd(1, .monotonic);
    defer {
        _ = stats.active.fetchSub(1, .monotonic);
        _ = stats.disconnects.fetchAdd(1, .monotonic);
    }

    // a read timeout lets churn-mode deadline checks run during quiet
    // periods, but timeouts mid-TLS-record can surface EAGAIN as a panic
    // in the 0.16 reader — so only arm it when churn actually needs it
    if (cfg.lifetime_s > 0) try client.readTimeout(1000);

    const deadline_us: ?i64 = if (cfg.lifetime_s > 0)
        nowUs(io) + @as(i64, @intCast(cfg.lifetime_s * std.time.us_per_s))
    else
        null;

    while (true) {
        if (stop.load(.monotonic)) break;
        if (deadline_us) |d| if (nowUs(io) >= d) break;

        const message = client.read() catch |err| switch (err) {
            error.Closed => return,
            else => return err,
        } orelse continue;
        defer client.done(message);

        switch (message.type) {
            .text, .binary => {
                _ = stats.messages.fetchAdd(1, .monotonic);
                _ = stats.bytes.fetchAdd(message.data.len, .monotonic);
                if (cfg.slow_ms > 0) try io.sleep(.fromMilliseconds(@intCast(cfg.slow_ms)), .awake);
            },
            .ping => try client.writeFrame(.pong, @constCast(message.data)),
            .close => {
                client.close(.{}) catch {};
                return;
            },
            .pong => {},
        }
    }
    client.close(.{}) catch {};
}

fn buildPath(io: Io, buf: *[2048]u8, cfg: Config) ![]const u8 {
    var w: std.Io.Writer = .fixed(buf);
    try w.writeAll(cfg.path);
    var has_param = false;

    var it = std.mem.splitScalar(u8, cfg.collections, ',');
    while (it.next()) |col| {
        if (col.len == 0) continue;
        try w.writeByte(if (!has_param) '?' else '&');
        try w.print("wantedCollections={s}", .{col});
        has_param = true;
    }

    if (cfg.rewind_s > 0) {
        const cursor = nowUs(io) - @as(i64, @intCast(cfg.rewind_s * std.time.us_per_s));
        try w.writeByte(if (!has_param) '?' else '&');
        try w.print("cursor={d}", .{cursor});
    }

    return w.buffered();
}

fn reportLoop(io: Io, cfg: Config) void {
    while (!stop.load(.monotonic)) {
        io.sleep(.fromSeconds(@intCast(cfg.report_s)), .awake) catch return;
        report(io, "tick");
    }
}

var last_messages: u64 = 0;
var last_bytes: u64 = 0;
var last_report_us: i64 = 0;

fn report(io: Io, label: []const u8) void {
    const now = nowUs(io);
    const messages = stats.messages.load(.monotonic);
    const bytes = stats.bytes.load(.monotonic);

    const dt_us = if (last_report_us == 0) @as(i64, 0) else now - last_report_us;
    const dt_s = @as(f64, @floatFromInt(dt_us)) / std.time.us_per_s;
    const msg_rate = if (dt_s > 0) @as(f64, @floatFromInt(messages - last_messages)) / dt_s else 0;
    const byte_rate = if (dt_s > 0) @as(f64, @floatFromInt(bytes - last_bytes)) / dt_s else 0;

    last_messages = messages;
    last_bytes = bytes;
    last_report_us = now;

    std.log.info("{s}: active={d} connects={d} disconnects={d} errors={d} msgs={d} ({d:.0}/s) bytes={d} ({d:.0}/s)", .{
        label,
        stats.active.load(.monotonic),
        stats.connects.load(.monotonic),
        stats.disconnects.load(.monotonic),
        stats.errors.load(.monotonic),
        messages,
        msg_rate,
        bytes,
        byte_rate,
    });
}

fn nowUs(io: Io) i64 {
    return Io.Timestamp.now(io, .real).toMicroseconds();
}

fn parseArgs(init: std.process.Init.Minimal) !Config {
    var cfg: Config = .{};
    var arg_it = init.args.iterate();
    _ = arg_it.next(); // program name
    while (arg_it.next()) |arg| {
        if (flagValue(arg, "--host=")) |v| {
            cfg.host = v;
        } else if (flagValue(arg, "--port=")) |v| {
            cfg.port = try std.fmt.parseInt(u16, v, 10);
        } else if (flagValue(arg, "--conns=")) |v| {
            cfg.conns = try std.fmt.parseInt(usize, v, 10);
        } else if (flagValue(arg, "--duration-s=")) |v| {
            cfg.duration_s = try std.fmt.parseInt(u64, v, 10);
        } else if (flagValue(arg, "--lifetime-s=")) |v| {
            cfg.lifetime_s = try std.fmt.parseInt(u64, v, 10);
        } else if (flagValue(arg, "--slow-ms=")) |v| {
            cfg.slow_ms = try std.fmt.parseInt(u64, v, 10);
        } else if (flagValue(arg, "--rewind-s=")) |v| {
            cfg.rewind_s = try std.fmt.parseInt(u64, v, 10);
        } else if (flagValue(arg, "--ramp-ms=")) |v| {
            cfg.ramp_ms = try std.fmt.parseInt(u64, v, 10);
        } else if (flagValue(arg, "--report-s=")) |v| {
            cfg.report_s = try std.fmt.parseInt(u64, v, 10);
        } else if (flagValue(arg, "--collections=")) |v| {
            cfg.collections = v;
        } else if (flagValue(arg, "--path=")) |v| {
            cfg.path = v;
        } else {
            std.log.err("unknown flag: {s}", .{arg});
            return error.UnknownFlag;
        }
    }
    return cfg;
}

fn flagValue(arg: []const u8, prefix: []const u8) ?[]const u8 {
    if (std.mem.startsWith(u8, arg, prefix)) return arg[prefix.len..];
    return null;
}