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

//! Traffic by service: connections grouped by remote port and protocol and
//! joined with the per-connection rates the traffic collector keeps, so that
//! "https is 90% of my traffic" falls out of data already at hand.

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

const addr = @import("addr.zig");
const model = @import("model.zig");

const Known = struct {
    tcp: bool,
    port: u16,
    name: []const u8,
};

/// Well-known remote ports. TCP and UDP on 443 are told apart, which is
/// https against QUIC for free.
const known = [_]Known{
    .{ .tcp = true, .port = 443, .name = "https" },
    .{ .tcp = false, .port = 443, .name = "quic" },
    .{ .tcp = true, .port = 80, .name = "http" },
    .{ .tcp = true, .port = 53, .name = "dns" },
    .{ .tcp = false, .port = 53, .name = "dns" },
    .{ .tcp = true, .port = 853, .name = "dot" },
    .{ .tcp = false, .port = 853, .name = "doq" },
    .{ .tcp = true, .port = 22, .name = "ssh" },
    .{ .tcp = false, .port = 123, .name = "ntp" },
    .{ .tcp = true, .port = 21, .name = "ftp" },
    .{ .tcp = true, .port = 25, .name = "smtp" },
    .{ .tcp = true, .port = 587, .name = "smtp" },
    .{ .tcp = true, .port = 465, .name = "smtps" },
    .{ .tcp = true, .port = 143, .name = "imap" },
    .{ .tcp = true, .port = 993, .name = "imaps" },
    .{ .tcp = true, .port = 110, .name = "pop3" },
    .{ .tcp = true, .port = 995, .name = "pop3s" },
    .{ .tcp = true, .port = 5222, .name = "xmpp" },
    .{ .tcp = false, .port = 5353, .name = "mdns" },
    .{ .tcp = false, .port = 1900, .name = "ssdp" },
    .{ .tcp = false, .port = 3478, .name = "stun" },
    .{ .tcp = true, .port = 3478, .name = "stun" },
    .{ .tcp = false, .port = 51820, .name = "wireguard" },
    .{ .tcp = true, .port = 1194, .name = "openvpn" },
    .{ .tcp = false, .port = 1194, .name = "openvpn" },
    .{ .tcp = false, .port = 500, .name = "ipsec" },
    .{ .tcp = false, .port = 4500, .name = "ipsec" },
    .{ .tcp = true, .port = 8080, .name = "http-alt" },
    .{ .tcp = true, .port = 8443, .name = "https-alt" },
    .{ .tcp = true, .port = 5432, .name = "postgres" },
    .{ .tcp = true, .port = 3306, .name = "mysql" },
    .{ .tcp = true, .port = 6379, .name = "redis" },
    .{ .tcp = true, .port = 27017, .name = "mongodb" },
};

fn lookup(tcp: bool, port: u16) ?[]const u8 {
    for (known) |k| {
        if (k.tcp == tcp and k.port == port) return k.name;
    }
    return null;
}

/// The service on a port when the protocol is not known, which is the case
/// for the traffic collector's per-process connections. TCP wins a tie, since
/// TCP is what both platforms can meter.
pub fn name(port: u16) []const u8 {
    return lookup(true, port) orelse lookup(false, port) orelse "";
}

pub const ServiceRate = struct {
    /// Empty when the port is not a well-known one.
    name: []const u8,
    protocol: []const u8,
    port: u16,
    rate: f64,
    rx_rate: f64,
    tx_rate: f64,
    connections: usize,
};

fn isTcp(protocol: []const u8) bool {
    return std.ascii.eqlIgnoreCase(protocol, "tcp");
}

fn busiestFirst(_: void, a: ServiceRate, b: ServiceRate) bool {
    if (a.rate != b.rate) return a.rate > b.rate;
    if (a.connections != b.connections) return a.connections > b.connections;
    return a.port < b.port;
}

/// Every service with a connection open, busiest first.
pub fn aggregate(
    arena: Allocator,
    connections: []const model.Connection,
    traffic: *const model.Traffic,
) Allocator.Error![]ServiceRate {
    var services: std.ArrayList(ServiceRate) = .empty;
    for (connections) |c| {
        if (c.remote.len == 0) continue;
        const port = addr.port(c.remote) orelse continue;
        const tcp = isTcp(c.protocol);
        const r = traffic.rateOf(c.local, c.remote) orelse model.Rates{};
        const entry = for (services.items) |*s| {
            if (s.port == port and isTcp(s.protocol) == tcp) break s;
        } else blk: {
            const s = try services.addOne(arena);
            s.* = .{
                .name = lookup(tcp, port) orelse "",
                .protocol = c.protocol,
                .port = port,
                .rate = 0,
                .rx_rate = 0,
                .tx_rate = 0,
                .connections = 0,
            };
            break :blk s;
        };
        entry.rate += r.sum();
        entry.rx_rate += r.rx;
        entry.tx_rate += r.tx;
        entry.connections += 1;
    }
    std.mem.sort(ServiceRate, services.items, {}, busiestFirst);
    return services.items;
}

const testing = std.testing;

fn conn(overrides: struct {
    protocol: []const u8 = "TCP",
    local: []const u8 = "192.168.1.2:50000",
    remote: []const u8 = "1.2.3.4:443",
}) model.Connection {
    return .{
        .command = "app",
        .pid = 1,
        .protocol = overrides.protocol,
        .local = overrides.local,
        .remote = overrides.remote,
        .remote_display = overrides.remote,
        .state = "ESTABLISHED",
    };
}

fn trafficWith(gpa: Allocator, rates: []const struct { []const u8, model.Rates }) !model.Traffic {
    var traffic: model.Traffic = .{};
    for (rates) |r| try traffic.connection_rates.put(gpa, r[0], r[1]);
    return traffic;
}

test "tcp/443 and udp/443 are different services" {
    var arena: std.heap.ArenaAllocator = .init(testing.allocator);
    defer arena.deinit();
    const traffic: model.Traffic = .{};
    const services = try aggregate(arena.allocator(), &.{
        conn(.{}),
        conn(.{ .protocol = "UDP", .local = "192.168.1.2:50001", .remote = "5.6.7.8:443" }),
    }, &traffic);
    try testing.expectEqual(@as(usize, 2), services.len);
    try testing.expect(std.mem.eql(u8, services[0].name, "https") or std.mem.eql(u8, services[1].name, "https"));
    try testing.expect(std.mem.eql(u8, services[0].name, "quic") or std.mem.eql(u8, services[1].name, "quic"));
}

test "rates and connections are summed per service" {
    var arena: std.heap.ArenaAllocator = .init(testing.allocator);
    defer arena.deinit();
    const traffic = try trafficWith(arena.allocator(), &.{
        .{ "192.168.1.2:50000|1.2.3.4:443", .{ .rx = 100, .tx = 20 } },
        .{ "192.168.1.2:50001|5.6.7.8:443", .{ .rx = 50, .tx = 5 } },
    });
    const services = try aggregate(arena.allocator(), &.{
        conn(.{}),
        conn(.{ .local = "192.168.1.2:50001", .remote = "5.6.7.8:443" }),
    }, &traffic);
    try testing.expectEqual(@as(usize, 1), services.len);
    try testing.expectEqual(@as(f64, 175), services[0].rate);
    try testing.expectEqual(@as(f64, 150), services[0].rx_rate);
    try testing.expectEqual(@as(f64, 25), services[0].tx_rate);
    try testing.expectEqual(@as(usize, 2), services[0].connections);
}

test "an unknown port has no name, and the busiest service sorts first" {
    var arena: std.heap.ArenaAllocator = .init(testing.allocator);
    defer arena.deinit();
    const traffic = try trafficWith(arena.allocator(), &.{
        .{ "192.168.1.2:50001|5.6.7.8:60123", .{ .rx = 900, .tx = 0 } },
    });
    const services = try aggregate(arena.allocator(), &.{
        conn(.{}),
        conn(.{ .local = "192.168.1.2:50001", .remote = "5.6.7.8:60123" }),
    }, &traffic);
    try testing.expectEqualStrings("", services[0].name);
    try testing.expectEqual(@as(u16, 60123), services[0].port);
    try testing.expectEqualStrings("https", services[1].name);
}

test "a connection without a remote end is skipped" {
    var arena: std.heap.ArenaAllocator = .init(testing.allocator);
    defer arena.deinit();
    const traffic: model.Traffic = .{};
    try testing.expectEqual(@as(usize, 0), (try aggregate(arena.allocator(), &.{conn(.{ .remote = "" })}, &traffic)).len);
}

test "name prefers tcp" {
    try testing.expectEqualStrings("https", name(443));
    try testing.expectEqualStrings("ntp", name(123));
    try testing.expectEqualStrings("", name(60123));
}