// SPDX-FileCopyrightText: © 2026 Francesc Muñoz
// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! The tasks that take the samples, and the `Store` they put them in. Each
//! kind of sample lives in an arena of its own; a new sample is built in a
//! fresh arena with no lock held, swapped in under the lock, and the old
//! arena freed after. The interface holds the same lock while it draws, so
//! nothing it is drawing can be freed from under it.

const std = @import("std");
const builtin = @import("builtin");
const Allocator = std.mem.Allocator;
const Io = std.Io;

const exec = @import("exec.zig");
const model = @import("model.zig");
const Lookups = @import("lookups.zig").Lookups;
const fetchPublicIp = @import("lookups.zig").fetchPublicIp;
const connections = @import("collectors/connections.zig");
const latency = @import("collectors/latency.zig");
const system = @import("collectors/system.zig");
const throughput = @import("collectors/throughput.zig");
const traffic = @import("collectors/traffic.zig");

const linux = builtin.os.tag == .linux;
comptime {
    if (!linux and !builtin.os.tag.isDarwin()) @compileError("ziptop runs on Linux and macOS");
}

/// The samples kept for the charts, at one a sample: as many as the widest
/// chart could show, two to a cell.
pub const history_size = 600;

/// The refresh intervals `[` and `]` step between, in milliseconds.
pub const rates = [_]u32{ 500, 1_000, 2_000 };

/// The newest `history_size` values of something.
pub const Series = struct {
    buffer: [history_size]f64 = undefined,
    len: usize = 0,

    pub fn push(self: *Series, value: f64) void {
        if (self.len == history_size) {
            std.mem.copyForwards(f64, self.buffer[0 .. history_size - 1], self.buffer[1..]);
            self.len -= 1;
        }
        self.buffer[self.len] = value;
        self.len += 1;
    }

    pub fn items(self: *const Series) []const f64 {
        return self.buffer[0..self.len];
    }
};

fn Piece(comptime T: type) type {
    return struct {
        arena: ?std.heap.ArenaAllocator = null,
        value: T,

        /// Takes the new value and its arena, returning the old arena for the
        /// caller to free once the lock is released.
        fn swap(self: *@This(), arena: std.heap.ArenaAllocator, value: T) ?std.heap.ArenaAllocator {
            const old = self.arena;
            self.arena = arena;
            self.value = value;
            return old;
        }

        fn deinit(self: *@This()) void {
            if (self.arena) |*a| a.deinit();
        }
    };
}

pub const Store = struct {
    mutex: Io.Mutex = .init,

    throughput: Piece(model.Throughput) = .{ .value = .{} },
    traffic: Piece(model.Traffic) = .{ .value = .{} },
    connections: Piece([]const model.Connection) = .{ .value = &.{} },
    system: Piece(model.System) = .{ .value = .{} },
    latency: Piece(model.Latency) = .{ .value = .{} },

    public_ip_buffer: [64]u8 = undefined,
    public_ip: []const u8 = "",

    rx_history: Series = .{},
    tx_history: Series = .{},
    gateway_history: Series = .{},
    internet_history: Series = .{},
    dns_history: Series = .{},
    /// -1 marks a sample in which that target was not probed, so that the
    /// average loss counts only the bursts that were sent.
    gateway_loss: Series = .{},
    internet_loss: Series = .{},
    dns_loss: Series = .{},

    /// An index into `rates`, which the interface changes and the samplers
    /// read on every round.
    rate_index: std.atomic.Value(u8) = .init(0),

    pub fn deinit(self: *Store) void {
        self.throughput.deinit();
        self.traffic.deinit();
        self.connections.deinit();
        self.system.deinit();
        self.latency.deinit();
    }

    pub fn intervalMs(self: *const Store) u32 {
        return rates[@min(self.rate_index.load(.monotonic), rates.len - 1)];
    }

    pub fn lock(self: *Store, io: Io) void {
        self.mutex.lockUncancelable(io);
    }

    pub fn unlock(self: *Store, io: Io) void {
        self.mutex.unlock(io);
    }
};

/// How a sampler tells the interface there is something new to draw.
pub const Wake = struct {
    context: *anyopaque,
    wake: *const fn (context: *anyopaque) void,
};

pub const Sampler = struct {
    gpa: Allocator,
    io: Io,
    store: *Store,
    lookups: *Lookups,
    wake: Wake,
    meter: throughput.Meter,
    engine: traffic.Engine,

    pub fn init(gpa: Allocator, io: Io, store: *Store, lookups: *Lookups, wake: Wake) Sampler {
        return .{
            .gpa = gpa,
            .io = io,
            .store = store,
            .lookups = lookups,
            .wake = wake,
            .meter = .init(gpa),
            .engine = .init(gpa),
        };
    }

    pub fn deinit(self: *Sampler) void {
        self.meter.deinit();
        self.engine.deinit();
    }

    pub fn start(self: *Sampler, group: *Io.Group) !void {
        try group.concurrent(self.io, fastLoop, .{self});
        try group.concurrent(self.io, latencyLoop, .{self});
        try group.concurrent(self.io, publicIpLoop, .{self});
    }

    fn seconds(self: *Sampler) f64 {
        const ns: f64 = @floatFromInt(Io.Clock.awake.now(self.io).nanoseconds);
        return ns / std.time.ns_per_s;
    }

    fn sleepRest(self: *Sampler, started: f64) Io.Cancelable!void {
        const interval: f64 = @floatFromInt(self.store.intervalMs());
        const spent = (self.seconds() - started) * 1000;
        const rest = interval - spent;
        if (rest > 0) try self.io.sleep(.fromMilliseconds(@intFromFloat(rest)), .awake);
    }

    fn fastLoop(self: *Sampler) Io.Cancelable!void {
        while (true) {
            const started = self.seconds();
            self.sampleFast() catch |err| switch (err) {
                error.Canceled => return error.Canceled,
                // Out of memory for one sample is a sample skipped.
                error.OutOfMemory => {},
            };
            self.wake.wake(self.wake.context);
            try self.sleepRest(started);
        }
    }

    fn sampleFast(self: *Sampler) (Io.Cancelable || Allocator.Error)!void {
        const io = self.io;
        const lookup = self.lookups.lookup();

        var tp_arena: std.heap.ArenaAllocator = .init(self.gpa);
        errdefer tp_arena.deinit();
        const tp_out = if (linux)
            try exec.readFile(tp_arena.allocator(), io, throughput.linux_path)
        else
            try exec.run(tp_arena.allocator(), io, &throughput.darwin_command);
        const counters = if (linux)
            try throughput.parseLinux(tp_arena.allocator(), tp_out)
        else
            try throughput.parseDarwin(tp_arena.allocator(), tp_out);
        const tp = try self.meter.sample(tp_arena.allocator(), counters, self.seconds());

        var tr_arena: std.heap.ArenaAllocator = .init(self.gpa);
        errdefer tr_arena.deinit();
        const tr_out = try exec.run(tr_arena.allocator(), io, if (linux) &traffic.linux_command else &traffic.darwin_command);
        const raw: ?traffic.Raw = if (linux)
            try traffic.parseLinux(tr_arena.allocator(), tr_out)
        else
            try traffic.parseDarwin(tr_arena.allocator(), tr_out);
        const tr: model.Traffic = if (raw) |r|
            try self.engine.sample(tr_arena.allocator(), r, self.seconds(), lookup)
        else
            .{};

        var cn_arena: std.heap.ArenaAllocator = .init(self.gpa);
        errdefer cn_arena.deinit();
        const cn_out = try exec.run(cn_arena.allocator(), io, if (linux) &connections.linux_command else &connections.darwin_command);
        const cn = if (linux)
            try connections.parseLinux(cn_arena.allocator(), cn_out)
        else
            try connections.parseDarwin(cn_arena.allocator(), cn_out);
        for (cn) |*c| {
            if (c.remote.len > 0 and std.mem.eql(u8, c.state, "ESTABLISHED")) {
                c.remote_display = try lookup.display(cn_arena.allocator(), c.remote);
            }
        }

        var sy_arena: std.heap.ArenaAllocator = .init(self.gpa);
        errdefer sy_arena.deinit();
        const sy = try self.sampleSystem(sy_arena.allocator());

        var old: [4]?std.heap.ArenaAllocator = undefined;
        {
            self.store.lock(io);
            defer self.store.unlock(io);
            old = .{
                self.store.throughput.swap(tp_arena, tp),
                self.store.traffic.swap(tr_arena, tr),
                self.store.connections.swap(cn_arena, cn),
                self.store.system.swap(sy_arena, sy),
            };
            self.store.rx_history.push(tp.rx_rate);
            self.store.tx_history.push(tp.tx_rate);
        }
        for (&old) |*a| if (a.*) |*arena| arena.deinit();
    }

    fn sampleSystem(self: *Sampler, arena: Allocator) (Io.Cancelable || Allocator.Error)!model.System {
        const io = self.io;
        if (linux) {
            const route = system.parseLinuxRoute(try exec.run(arena, io, &system.linux_route_command));
            var local_ip = route.local_ip;
            if (local_ip.len == 0) if (route.interface) |dev| {
                local_ip = system.parseLinuxAddress(try exec.run(arena, io, &.{ "ip", "-4", "addr", "show", "dev", dev }));
            };
            return .{ .gateway = route.gateway, .local_ip = local_ip };
        } else {
            const route = system.parseDarwinRoute(try exec.run(arena, io, &system.darwin_route_command));
            const local = try exec.run(arena, io, &.{ "ipconfig", "getifaddr", route.interface orelse "en0" });
            return .{ .gateway = route.gateway, .local_ip = std.mem.trim(u8, local, " \t\r\n") };
        }
    }

    /// Asked until it is answered, then never again: the public address
    /// changes rarely enough that a restart is the way to see a new one.
    fn publicIpLoop(self: *Sampler) Io.Cancelable!void {
        var buffer: [64]u8 = undefined;
        while (true) {
            const ip = try fetchPublicIp(self.gpa, self.io, &buffer);
            if (ip.len > 0) {
                self.store.lock(self.io);
                defer self.store.unlock(self.io);
                @memcpy(self.store.public_ip_buffer[0..ip.len], ip);
                self.store.public_ip = self.store.public_ip_buffer[0..ip.len];
                break;
            }
            try self.io.sleep(.fromSeconds(30), .awake);
        }
        self.wake.wake(self.wake.context);
    }

    fn dnsServer(self: *Sampler, buffer: []u8) Io.Cancelable![]const u8 {
        var arena: std.heap.ArenaAllocator = .init(self.gpa);
        defer arena.deinit();
        const server = if (linux)
            latency.parseResolvConf(try exec.readFile(arena.allocator(), self.io, "/etc/resolv.conf"))
        else
            latency.parseScutil(try exec.run(arena.allocator(), self.io, &.{ "scutil", "--dns" }));
        const n = @min(server.len, buffer.len);
        @memcpy(buffer[0..n], server[0..n]);
        return buffer[0..n];
    }

    fn latencyLoop(self: *Sampler) Io.Cancelable!void {
        var dns_buffer: [64]u8 = undefined;
        const dns = try self.dnsServer(&dns_buffer);
        var gateway_buffer: [64]u8 = undefined;

        while (true) {
            const started = self.seconds();
            const gateway = blk: {
                self.store.lock(self.io);
                defer self.store.unlock(self.io);
                const g = self.store.system.value.gateway;
                const n = @min(g.len, gateway_buffer.len);
                @memcpy(gateway_buffer[0..n], g[0..n]);
                break :blk gateway_buffer[0..n];
            };

            // The DNS server is often the gateway, and pinging the same
            // address twice says nothing more.
            const targets = [3][]const u8{ gateway, latency.internet_host, if (std.mem.eql(u8, dns, gateway)) "" else dns };
            var times: [3][latency.burst]?f64 = @splat(@splat(null));
            {
                var group: Io.Group = .init;
                defer group.cancel(self.io);
                for (targets, &times) |host, *row| {
                    if (host.len == 0) continue;
                    for (row, 0..) |*slot, i| {
                        group.concurrent(self.io, probe, .{ self, host, i, slot }) catch {
                            try probe(self, host, i, slot);
                        };
                    }
                }
                try group.await(self.io);
            }

            var bursts: [3]?latency.Burst = undefined;
            for (targets, times, &bursts) |host, row, *b| b.* = if (host.len == 0) null else .of(&row);

            var arena: std.heap.ArenaAllocator = .init(self.gpa);
            const sample: model.Latency = .{
                .gateway_ms = if (bursts[0]) |b| b.ms else null,
                .internet_ms = if (bursts[1]) |b| b.ms else null,
                .dns_ms = if (bursts[2]) |b| b.ms else null,
                .gateway_loss = if (bursts[0]) |b| b.loss else null,
                .internet_loss = if (bursts[1]) |b| b.loss else null,
                .dns_loss = if (bursts[2]) |b| b.loss else null,
                .dns_server = arena.allocator().dupe(u8, dns) catch "",
            };
            var old: ?std.heap.ArenaAllocator = null;
            {
                self.store.lock(self.io);
                defer self.store.unlock(self.io);
                old = self.store.latency.swap(arena, sample);
                self.store.gateway_history.push(sample.gateway_ms orelse 0);
                self.store.internet_history.push(sample.internet_ms orelse 0);
                self.store.dns_history.push(sample.dns_ms orelse 0);
                self.store.gateway_loss.push(sample.gateway_loss orelse -1);
                self.store.internet_loss.push(sample.internet_loss orelse -1);
                self.store.dns_loss.push(sample.dns_loss orelse -1);
            }
            if (old) |*a| a.deinit();
            self.wake.wake(self.wake.context);
            try self.sleepRest(started);
        }
    }

    /// The `i`th ping of a burst, started `i` spacings after the first.
    fn probe(self: *Sampler, host: []const u8, i: usize, out: *?f64) Io.Cancelable!void {
        if (i > 0) try self.io.sleep(.fromMilliseconds(@intCast(i * latency.burst_spacing_ms)), .awake);
        var arena: std.heap.ArenaAllocator = .init(self.gpa);
        defer arena.deinit();
        const argv = latency.pingCommand(host);
        out.* = latency.parsePing(try exec.run(arena.allocator(), self.io, &argv));
    }
};

const testing = std.testing;

test "Series keeps the newest values" {
    var s: Series = .{};
    for (0..history_size + 5) |i| s.push(@floatFromInt(i));
    try testing.expectEqual(@as(usize, history_size), s.items().len);
    try testing.expectEqual(@as(f64, 5), s.items()[0]);
    try testing.expectEqual(@as(f64, history_size + 4), s.items()[history_size - 1]);
}