import Foundation

public enum IPv6Transport: UInt8, Sendable {
    case tcp = 6
    case udp = 17
    case icmpv6 = 58
}

public struct IPv6PacketInfo: Sendable {
    public let source: MeshAddress
    public let destination: MeshAddress
    public let transport: IPv6Transport
    public let sourcePort: UInt16?
    public let destinationPort: UInt16?
    public let tcpFlags: UInt8?
    public let icmpType: UInt8?
    public let transportOffset: Int
}

public enum IPv6PacketParser {
    public static let maximumPacketSize = 1280

    public static func parse(_ packet: Data) throws -> IPv6PacketInfo {
        guard packet.count >= 40, packet.count <= maximumPacketSize, packet[0] >> 4 == 6 else { throw UMNError.invalidFrame }
        let payloadLength = Int(read16(packet, 4))
        guard payloadLength == packet.count - 40 else { throw UMNError.invalidFrame }
        let source = try MeshAddress(bytes: packet.subdata(in: 8..<24))
        let destinationBytes = packet.subdata(in: 24..<40)
        guard destinationBytes.first != 0xff else { throw UMNError.invalidFrame }
        let destination = try MeshAddress(bytes: destinationBytes)

        var next = packet[6]
        var offset = 40
        var extensions = 0
        while [UInt8(0), 43, 60].contains(next) {
            guard extensions < 8, offset + 2 <= packet.count else { throw UMNError.invalidFrame }
            let following = packet[offset]
            let length: Int
            length = (Int(packet[offset + 1]) + 1) * 8
            guard length >= 8, offset + length <= packet.count else { throw UMNError.invalidFrame }
            if next == 43, packet[offset + 3] != 0 { throw UMNError.invalidFrame }
            if next == 0 || next == 60 { try validateOptions(packet, range: (offset + 2)..<(offset + length)) }
            next = following; offset += length; extensions += 1
        }
        guard next != 44, let transport = IPv6Transport(rawValue: next) else { throw UMNError.invalidFrame }

        switch transport {
        case .tcp:
            guard offset + 20 <= packet.count else { throw UMNError.invalidFrame }
            let headerLength = Int(packet[offset + 12] >> 4) * 4
            guard headerLength >= 20, offset + headerLength <= packet.count else { throw UMNError.invalidFrame }
            return .init(source: source, destination: destination, transport: transport,
                         sourcePort: read16(packet, offset), destinationPort: read16(packet, offset + 2),
                         tcpFlags: packet[offset + 13], icmpType: nil, transportOffset: offset)
        case .udp:
            guard offset + 8 <= packet.count else { throw UMNError.invalidFrame }
            let length = Int(read16(packet, offset + 4))
            guard length >= 8, offset + length == packet.count else { throw UMNError.invalidFrame }
            return .init(source: source, destination: destination, transport: transport,
                         sourcePort: read16(packet, offset), destinationPort: read16(packet, offset + 2),
                         tcpFlags: nil, icmpType: nil, transportOffset: offset)
        case .icmpv6:
            guard offset + 8 <= packet.count else { throw UMNError.invalidFrame }
            return .init(source: source, destination: destination, transport: transport,
                         sourcePort: nil, destinationPort: nil, tcpFlags: nil,
                         icmpType: packet[offset], transportOffset: offset)
        }
    }

    private static func read16(_ data: Data, _ offset: Int) -> UInt16 {
        UInt16(data[offset]) << 8 | UInt16(data[offset + 1])
    }
    private static func validateOptions(_ data: Data, range: Range<Int>) throws {
        var index = range.lowerBound
        while index < range.upperBound {
            if data[index] == 0 { index += 1; continue } // Pad1
            guard index + 2 <= range.upperBound else { throw UMNError.invalidFrame }
            let length = Int(data[index + 1])
            guard index + 2 + length <= range.upperBound else { throw UMNError.invalidFrame }
            index += 2 + length
        }
    }
}

public final class NativePacketFirewall: @unchecked Sendable {
    private struct Flow: Hashable {
        let kind: FirewallProtocol
        let local: MeshAddress
        let remote: MeshAddress
        let localPort: UInt16
        let remotePort: UInt16
    }
    private struct Entry { var lastSeen: Date; var established: Bool }
    private let rules: MeshFirewall
    private let lock = NSLock()
    private var flows: [Flow: Entry] = [:]
    private var echoWindows: [MeshAddress: (Date, Int)] = [:]
    public let maximumFlows: Int

    public init(rules: MeshFirewall, maximumFlows: Int = 4096) {
        self.rules = rules; self.maximumFlows = maximumFlows
    }

    /// Records an OS-originated packet. Outbound connections are always permitted.
    public func allowOutbound(_ info: IPv6PacketInfo, now: Date = Date()) -> Bool {
        guard let kind = protocolKind(info), let localPort = info.sourcePort, let remotePort = info.destinationPort else {
            return info.transport == .icmpv6 && [1, 2, 3, 4, 128, 129].contains(info.icmpType ?? 255)
        }
        lock.lock(); defer { lock.unlock() }
        expireLocked(now)
        let key = Flow(kind: kind, local: info.source, remote: info.destination,
                       localPort: localPort, remotePort: remotePort)
        insert(key, entry: .init(lastSeen: now, established: kind == .udp || ((info.tcpFlags ?? 0) & 0x02) != 0))
        return true
    }

    public func allowInbound(_ info: IPv6PacketInfo, packet: Data? = nil, now: Date = Date()) -> Bool {
        if info.transport == .icmpv6 { return allowICMP(info, packet: packet, now: now) }
        guard let kind = protocolKind(info), let remotePort = info.sourcePort, let localPort = info.destinationPort else { return false }
        lock.lock(); defer { lock.unlock() }
        expireLocked(now)
        let key = Flow(kind: kind, local: info.destination, remote: info.source,
                       localPort: localPort, remotePort: remotePort)
        if var entry = flows[key] {
            entry.lastSeen = now
            if kind == .tcp, ((info.tcpFlags ?? 0) & 0x12) == 0x12 { entry.established = true }
            flows[key] = entry
            if kind == .tcp, ((info.tcpFlags ?? 0) & 0x05) != 0 { flows.removeValue(forKey: key) }
            return true
        }
        guard rules.allows(protocol: kind, port: localPort, source: info.source) else { return false }
        // A new TCP flow must begin with SYN and not ACK. UDP is connectionless.
        if kind == .tcp, ((info.tcpFlags ?? 0) & 0x12) != 0x02 { return false }
        insert(key, entry: .init(lastSeen: now, established: kind == .udp))
        return true
    }

    private func allowICMP(_ info: IPv6PacketInfo, packet: Data?, now: Date) -> Bool {
        guard let type = info.icmpType else { return false }
        if (1...4).contains(type), let packet, packet.count >= info.transportOffset + 8 + 48 {
            var quoted = Data(packet[(info.transportOffset + 8)...])
            let actualPayload = quoted.count - 40
            quoted[4] = UInt8(actualPayload >> 8); quoted[5] = UInt8(actualPayload & 255)
            guard let original = try? IPv6PacketParser.parse(quoted),
                  let kind = protocolKind(original), let localPort = original.sourcePort,
                  let remotePort = original.destinationPort else { return false }
            lock.lock(); defer { lock.unlock() }; expireLocked(now)
            return flows[Flow(kind: kind, local: original.source, remote: original.destination,
                              localPort: localPort, remotePort: remotePort)] != nil
        }
        guard type == 128 || type == 129 else { return false }
        lock.lock(); defer { lock.unlock() }
        let current = echoWindows[info.source]
        if let current, now.timeIntervalSince(current.0) < 1 {
            guard current.1 < 10 else { return false }
            echoWindows[info.source] = (current.0, current.1 + 1)
        } else {
            guard echoWindows[info.source] != nil || echoWindows.count < 128 else { return false }
            echoWindows[info.source] = (now, 1)
        }
        return true
    }

    private func protocolKind(_ info: IPv6PacketInfo) -> FirewallProtocol? {
        info.transport == .tcp ? .tcp : (info.transport == .udp ? .udp : nil)
    }
    private func insert(_ flow: Flow, entry: Entry) {
        if flows[flow] == nil, flows.count >= maximumFlows,
           let oldest = flows.min(by: { $0.value.lastSeen < $1.value.lastSeen })?.key { flows.removeValue(forKey: oldest) }
        flows[flow] = entry
    }
    private func expireLocked(_ now: Date) {
        flows = flows.filter { now.timeIntervalSince($0.value.lastSeen) <= ($0.key.kind == .tcp ? 300 : 60) }
        echoWindows = echoWindows.filter { now.timeIntervalSince($0.value.0) < 60 }
    }
}