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

//! The per-country colors shared by the world map, the remote hosts and the
//! inspector. A country's identity is its hue, handed out on first sight and
//! kept for the session; how much traffic it carries is how vivid that hue
//! is, on a log scale and quantized so that cells on the map merge cleanly.

const std = @import("std");

const theme = @import("theme.zig");
const Rgb = theme.Rgb;

/// The palette colors a country can be given, by name, so that a theme
/// switch recolors the map. Upload's pink is left out: markers brighten
/// toward it and would clash.
const hues = [_]enum { cyan, green, orange, magenta, yellow, accent }{ .cyan, .green, .orange, .magenta, .yellow, .accent };

fn hueColor(i: usize) Rgb {
    const p = &theme.current;
    return switch (hues[i]) {
        .cyan => p.cyan,
        .green => p.green,
        .orange => p.orange,
        .magenta => p.magenta,
        .yellow => p.yellow,
        .accent => p.accent,
    };
}

/// The muted base a hue fades toward at low volume.
const base: Rgb = .hex("#414868");
pub const levels = 6;
const min_vividness = 0.75;

/// The order hues are handed out in, shuffled once per session so that the
/// colors differ from run to run, then dealt in turn so that no two countries
/// share one until every hue has been used.
var order: [hues.len]u8 = .{ 0, 1, 2, 3, 4, 5 };

/// Country codes in the order they were first seen; a country's hue is its
/// position here, through `order`.
var assigned: std.ArrayList([2]u8) = .empty;
var assigned_buffer: [256][2]u8 = undefined;

pub fn shuffle(seed: u64) void {
    var prng: std.Random.DefaultPrng = .init(seed);
    prng.random().shuffle(u8, &order);
    assigned = .initBuffer(&assigned_buffer);
}

fn code(country: []const u8) ?[2]u8 {
    if (country.len != 2) return null;
    return .{ country[0], country[1] };
}

fn hueOf(c: [2]u8) Rgb {
    if (assigned.capacity == 0) assigned = .initBuffer(&assigned_buffer);
    const i = for (assigned.items, 0..) |a, i| {
        if (std.mem.eql(u8, &a, &c)) break i;
    } else blk: {
        // Past 256 countries every newcomer shares the last slot, which is
        // more countries than a session talks to.
        if (assigned.items.len < assigned.capacity) assigned.appendAssumeCapacity(c);
        break :blk assigned.items.len - 1;
    };
    return hueColor(order[i % order.len]);
}

pub const Tint = struct {
    level: usize,
    color: Rgb,
    /// The hue brightened, for a marker on the map, so that it shows against
    /// the tinted land around it.
    marker: Rgb,
};

/// A country's hue at its least vivid, for one not among the current
/// remotes -- a quiet connection in the inspector, say.
pub fn fallback(country: []const u8) Rgb {
    const c = code(country) orelse return theme.current.dim;
    return base.mix(hueOf(c), min_vividness);
}

/// The tints of every country among the remote hosts, keyed by country code.
pub const Tints = struct {
    entries: [64]Entry = undefined,
    len: usize = 0,

    const Entry = struct { country: [2]u8, rate: f64, tint: Tint };

    pub fn get(self: *const Tints, country: []const u8) ?Tint {
        const c = code(country) orelse return null;
        for (self.entries[0..self.len]) |e| {
            if (std.mem.eql(u8, &e.country, &c)) return e.tint;
        }
        return null;
    }
};

pub fn compute(remotes: anytype) Tints {
    var tints: Tints = .{};
    for (remotes) |r| {
        const c = code(r.country) orelse continue;
        const rate = r.rx_rate + r.tx_rate;
        for (tints.entries[0..tints.len]) |*e| {
            if (std.mem.eql(u8, &e.country, &c)) {
                e.rate += rate;
                break;
            }
        } else if (tints.len < tints.entries.len) {
            tints.entries[tints.len] = .{ .country = c, .rate = rate, .tint = undefined };
            tints.len += 1;
        }
    }
    var max_rate: f64 = 1;
    for (tints.entries[0..tints.len]) |e| max_rate = @max(max_rate, e.rate);
    for (tints.entries[0..tints.len]) |*e| {
        const t = std.math.log1p(e.rate) / std.math.log1p(max_rate);
        const level: usize = @intFromFloat(@round(t * (levels - 1)));
        const level_f: f64 = @floatFromInt(level);
        const vividness = min_vividness + (1 - min_vividness) * (level_f / (levels - 1));
        const hue = hueOf(e.country);
        e.tint = .{
            .level = level,
            .color = base.mix(hue, vividness),
            .marker = hue.mix(theme.white, 0.6),
        };
    }
    return tints;
}

const testing = std.testing;

test "a country keeps its hue, and the busiest is the most vivid" {
    shuffle(1);
    const R = struct { country: []const u8, rx_rate: f64, tx_rate: f64 };
    const tints = compute(&[_]R{
        .{ .country = "US", .rx_rate = 1_000_000, .tx_rate = 0 },
        .{ .country = "DE", .rx_rate = 10, .tx_rate = 0 },
        .{ .country = "US", .rx_rate = 5, .tx_rate = 5 },
        .{ .country = "", .rx_rate = 5, .tx_rate = 5 },
    });
    try testing.expectEqual(@as(usize, 2), tints.len);
    try testing.expectEqual(@as(usize, levels - 1), tints.get("US").?.level);
    try testing.expect(tints.get("DE").?.level < levels - 1);
    try testing.expect(tints.get("FR") == null);
    // The same country comes back the same color.
    const again = compute(&[_]R{.{ .country = "DE", .rx_rate = 1, .tx_rate = 0 }});
    try testing.expect(fallback("DE").eql(base.mix(hueOf(.{ 'D', 'E' }), min_vividness)));
    _ = again;
}