import Foundation
import IrdP2PCrypto

/// Which peers get a data channel when `RoomOptions.maxDirectPeers` caps the mesh:
/// each peer's budget is the `limit` members nearest to it by XOR distance on the
/// 16-byte peer id, and a pair connects only when each is within the other's budget.
/// The rule is symmetric, so both ends derive the same answer from the same roster
/// without negotiating, and nobody offers to a peer that would refuse.
enum DirectBudget {
    /// Byte-wise XOR of the two ids; distances compare lexicographically (the first
    /// differing byte decides), so distinct peers never tie.
    static func distance(_ a: PeerId, _ b: PeerId) -> [UInt8] {
        zip(Bytes.fromHex(a.hex) ?? Data(), Bytes.fromHex(b.hex) ?? Data()).map { $0 ^ $1 }
    }

    /// The `limit` members nearest to `own` (itself excluded).
    static func nearest(to own: PeerId, among members: Set<PeerId>, limit: Int) -> Set<PeerId> {
        let others = members.subtracting([own]).map { (id: $0, d: distance(own, $0)) }
        return Set(others.sorted { $0.d.lexicographicallyPrecedes($1.d) }.prefix(max(0, limit)).map(\.id))
    }

    /// Whether `a` and `b` may connect directly among `members` (the verified roster, both included).
    static func allows(_ a: PeerId, _ b: PeerId, members: Set<PeerId>, limit: Int?) -> Bool {
        guard let limit else { return true }
        let all = members.union([a, b])
        return nearest(to: a, among: all, limit: limit).contains(b) && nearest(to: b, among: all, limit: limit).contains(a)
    }
}