import Foundation

public enum FirewallProtocol: String, Codable, CaseIterable, Sendable {
    case overlay, tcp, udp
}

public struct FirewallRule: Codable, Hashable, Sendable {
    public let protocolKind: FirewallProtocol
    public let port: UInt16
    public let source: MeshAddress?
    public init(protocolKind: FirewallProtocol = .overlay, port: UInt16, source: MeshAddress?) {
        self.protocolKind = protocolKind; self.port = port; self.source = source
    }

    private enum CodingKeys: String, CodingKey { case protocolKind = "protocol", port, source }
    public init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        protocolKind = try c.decodeIfPresent(FirewallProtocol.self, forKey: .protocolKind) ?? .overlay
        port = try c.decode(UInt16.self, forKey: .port)
        source = try c.decodeIfPresent(MeshAddress.self, forKey: .source)
    }
}

public final class MeshFirewall: @unchecked Sendable {
    private var rules: Set<FirewallRule>
    private let lock = NSLock()
    public init(rules: Set<FirewallRule> = []) { self.rules = rules }

    public func allows(protocol protocolKind: FirewallProtocol = .overlay, port: UInt16, source: MeshAddress) -> Bool {
        lock.lock(); defer { lock.unlock() }
        return rules.contains(FirewallRule(protocolKind: protocolKind, port: port, source: nil)) ||
            rules.contains(FirewallRule(protocolKind: protocolKind, port: port, source: source))
    }
    public func allow(protocol protocolKind: FirewallProtocol = .overlay, port: UInt16, source: MeshAddress?) {
        lock.lock(); rules.insert(.init(protocolKind: protocolKind, port: port, source: source)); lock.unlock()
    }
    public func revoke(protocol protocolKind: FirewallProtocol = .overlay, port: UInt16, source: MeshAddress?) {
        lock.lock(); rules.remove(.init(protocolKind: protocolKind, port: port, source: source)); lock.unlock()
    }
    public func allRules() -> [FirewallRule] {
        lock.lock(); defer { lock.unlock() }
        return rules.sorted { ($0.protocolKind.rawValue, $0.port, $0.source?.description ?? "") <
            ($1.protocolKind.rawValue, $1.port, $1.source?.description ?? "") }
    }
}