import Foundation
import Darwin
import UltraMeshCore
import UMNSystem

struct TestFailure: Error, LocalizedError {
    let errorDescription: String?
    init(_ message: String) { errorDescription = message }
}

var passed = 0
func check(_ condition: @autoclosure () throws -> Bool, _ name: String) throws {
    guard try condition() else { throw TestFailure(name) }
    passed += 1; print("ok \(passed) - \(name)")
}

func rejects(_ name: String, _ operation: () throws -> Void) throws {
    do { try operation() } catch {
        passed += 1; print("ok \(passed) - \(name)")
        return
    }
    throw TestFailure(name)
}

func ipv6Packet(source: MeshAddress, destination: MeshAddress, next: UInt8, payload: Data) -> Data {
    var packet = Data(repeating: 0, count: 40)
    packet[0] = 0x60; packet[4] = UInt8(payload.count >> 8); packet[5] = UInt8(payload.count & 255)
    packet[6] = next; packet[7] = 64
    packet.replaceSubrange(8..<24, with: source.bytes); packet.replaceSubrange(24..<40, with: destination.bytes)
    packet.append(payload); return packet
}
func tcpHeader(source: UInt16, destination: UInt16, flags: UInt8) -> Data {
    var data = Data(repeating: 0, count: 20)
    data[0] = UInt8(source >> 8); data[1] = UInt8(source & 255)
    data[2] = UInt8(destination >> 8); data[3] = UInt8(destination & 255)
    data[12] = 0x50; data[13] = flags; return data
}
func udpHeader(source: UInt16, destination: UInt16, payload: Data = Data()) -> Data {
    let length = 8 + payload.count; var data = Data(repeating: 0, count: 8)
    data[0] = UInt8(source >> 8); data[1] = UInt8(source & 255)
    data[2] = UInt8(destination >> 8); data[3] = UInt8(destination & 255)
    data[4] = UInt8(length >> 8); data[5] = UInt8(length & 255); data.append(payload); return data
}
func dnsQuery(_ name: String, type: UInt16) -> Data {
    var data = Data([0x12, 0x34, 0x01, 0x00, 0, 1, 0, 0, 0, 0, 0, 0])
    for label in name.split(separator: ".") { data.append(UInt8(label.utf8.count)); data.append(contentsOf: label.utf8) }
    data.append(0); data.append(UInt8(type >> 8)); data.append(UInt8(type & 255)); data.append(contentsOf: [0, 1]); return data
}

do {
    let alice = try NodeIdentity(); let bob = try NodeIdentity(); let carol = try NodeIdentity()
    try check(try MeshAddress(alice.record.address.description) == alice.record.address, "mesh address round trip")
    try check(alice.record.address.bytes.first == 0xfd, "mesh address uses fd prefix")

    let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
    defer { try? FileManager.default.removeItem(at: directory) }
    let identityURL = directory.appendingPathComponent("identity.plist")
    let addressURL = directory.appendingPathComponent("address")
    let stored1 = try NodeIdentity.loadOrCreate(in: directory)
    let stored2 = try NodeIdentity.loadOrCreate(in: directory)
    try check(stored1.record == stored2.record, "identity and address persist across daemon-style reloads")
    let attributes = try FileManager.default.attributesOfItem(atPath: identityURL.path)
    try check((attributes[.posixPermissions] as? NSNumber)?.intValue == 0o600, "identity permissions")
    let addressAttributes = try FileManager.default.attributesOfItem(atPath: addressURL.path)
    try check((addressAttributes[.posixPermissions] as? NSNumber)?.intValue == 0o600, "address permissions")
    let storedAddressText = try String(contentsOf: addressURL, encoding: .utf8)
        .trimmingCharacters(in: .whitespacesAndNewlines)
    try check(try MeshAddress(storedAddressText) == stored1.record.address, "stored address is valid IPv6 matching identity")

    let legacyDirectory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
    defer { try? FileManager.default.removeItem(at: legacyDirectory) }
    let legacyIdentity = try NodeIdentity.loadOrCreate(at: legacyDirectory.appendingPathComponent("identity.plist"))
    try check(!FileManager.default.fileExists(atPath: legacyDirectory.appendingPathComponent("address").path),
              "legacy path API creates only identity")
    let upgradedIdentity = try NodeIdentity.loadOrCreate(in: legacyDirectory)
    try check(upgradedIdentity.record == legacyIdentity.record &&
              FileManager.default.fileExists(atPath: legacyDirectory.appendingPathComponent("address").path),
              "existing installation gains address sidecar without identity change")

    let corruptDirectory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
    defer { try? FileManager.default.removeItem(at: corruptDirectory) }
    try FileManager.default.createDirectory(at: corruptDirectory, withIntermediateDirectories: true)
    let corruptURL = corruptDirectory.appendingPathComponent("identity.plist")
    let corruptBytes = Data("not an identity".utf8)
    try corruptBytes.write(to: corruptURL)
    try rejects("corrupt identity fails without replacement") { _ = try NodeIdentity.loadOrCreate(in: corruptDirectory) }
    try check(try Data(contentsOf: corruptURL) == corruptBytes, "corrupt identity remains unchanged")

    let incompatibleDirectory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
    defer { try? FileManager.default.removeItem(at: incompatibleDirectory) }
    try FileManager.default.createDirectory(at: incompatibleDirectory, withIntermediateDirectories: true)
    let incompatibleURL = incompatibleDirectory.appendingPathComponent("identity.plist")
    let incompatibleBytes = try PropertyListSerialization.data(
        fromPropertyList: ["version": 999, "signing": Data(repeating: 1, count: 32),
                           "agreement": Data(repeating: 2, count: 32)], format: .binary, options: 0)
    try incompatibleBytes.write(to: incompatibleURL)
    try rejects("incompatible identity version fails without replacement") {
        _ = try NodeIdentity.loadOrCreate(in: incompatibleDirectory)
    }
    try check(try Data(contentsOf: incompatibleURL) == incompatibleBytes, "incompatible identity remains unchanged")

    let orphanDirectory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
    defer { try? FileManager.default.removeItem(at: orphanDirectory) }
    try FileManager.default.createDirectory(at: orphanDirectory, withIntermediateDirectories: true)
    let orphanAddressURL = orphanDirectory.appendingPathComponent("address")
    let orphanBytes = Data("\(alice.record.address)\n".utf8)
    try orphanBytes.write(to: orphanAddressURL)
    try rejects("address without identity fails without creating identity") {
        _ = try NodeIdentity.loadOrCreate(in: orphanDirectory)
    }
    try check(!FileManager.default.fileExists(atPath: orphanDirectory.appendingPathComponent("identity.plist").path) &&
              (try Data(contentsOf: orphanAddressURL)) == orphanBytes, "orphan address state remains unchanged")

    let mismatchDirectory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
    defer { try? FileManager.default.removeItem(at: mismatchDirectory) }
    let mismatchIdentity = try NodeIdentity.loadOrCreate(in: mismatchDirectory)
    let mismatchAddressURL = mismatchDirectory.appendingPathComponent("address")
    let mismatchedBytes = Data("\(bob.record.address)\n".utf8)
    try mismatchedBytes.write(to: mismatchAddressURL, options: .atomic)
    try rejects("address and identity mismatch fails") { _ = try NodeIdentity.loadOrCreate(in: mismatchDirectory) }
    try check(try Data(contentsOf: mismatchAddressURL) == mismatchedBytes &&
              mismatchIdentity.record.address != bob.record.address, "mismatched address remains unchanged")

    let unreadableDirectory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
    defer { try? FileManager.default.removeItem(at: unreadableDirectory) }
    _ = try NodeIdentity.loadOrCreate(in: unreadableDirectory)
    let unreadableURL = unreadableDirectory.appendingPathComponent("identity.plist")
    let unreadableBytes = try Data(contentsOf: unreadableURL)
    try FileManager.default.setAttributes([.posixPermissions: 0o000], ofItemAtPath: unreadableURL.path)
    defer { try? FileManager.default.setAttributes([.posixPermissions: 0o600], ofItemAtPath: unreadableURL.path) }
    try rejects("unreadable identity fails without replacement") {
        _ = try NodeIdentity.loadOrCreate(in: unreadableDirectory)
    }
    try FileManager.default.setAttributes([.posixPermissions: 0o600], ofItemAtPath: unreadableURL.path)
    try check(try Data(contentsOf: unreadableURL) == unreadableBytes, "unreadable identity remains unchanged")

    let original = InnerFrame(kind: .text, port: 7000, payload: Data("hello".utf8), sourceRecord: alice.record)
    let sealed = try alice.seal(original, to: bob.record)
    let opened = try bob.open(sealed, expectedSource: alice.record.address)
    try check(opened.payload == Data("hello".utf8), "end-to-end encryption round trip")
    var corrupted = sealed.combinedCiphertext; corrupted[corrupted.startIndex] ^= 1
    let tampered = SealedPayload(ephemeralPublicKey: sealed.ephemeralPublicKey, combinedCiphertext: corrupted)
    var rejected = false
    do { _ = try bob.open(tampered, expectedSource: alice.record.address) } catch { rejected = true }
    try check(rejected, "tampered ciphertext rejected")
    let replayID = UUID()
    let nativeFrame = InnerFrame(kind: .ipv6Packet, payload: Data([1, 2]), sourceRecord: alice.record, replayID: replayID)
    let nativeOpened = try bob.open(alice.seal(nativeFrame, to: bob.record), expectedSource: alice.record.address)
    try check(nativeOpened.replayID == replayID, "signed native packet replay UUID round trip")

    let router = LinkStateRouter(local: alice.record.address)
    try check(router.ingest(try alice.makeLinkState(sequence: 1, neighbors: [bob.record.address])), "local topology accepted")
    try check(router.ingest(try bob.makeLinkState(sequence: 1, neighbors: [alice.record.address, carol.record.address])), "relay topology accepted")
    try check(router.ingest(try carol.makeLinkState(sequence: 1, neighbors: [bob.record.address])), "remote topology accepted")
    let route = router.routes()[carol.record.address]
    try check(route?.nextHop == bob.record.address && route?.hopCount == 2, "three-node route uses relay")
    try check(!router.ingest(try carol.makeLinkState(sequence: 1, neighbors: [])), "stale topology rejected")
    let forged = LinkState(origin: carol.record, sequence: 2, neighbors: [], signature: Data(repeating: 0, count: 64))
    try check(!router.ingest(forged), "forged topology rejected")
    try check((try NodeIdentity.loadOrCreate(in: directory)).record.address == stored1.record.address,
              "route and peer topology changes do not alter persistent local address")

    let firewall = MeshFirewall()
    try check(!firewall.allows(port: 80, source: alice.record.address), "firewall defaults to deny")
    firewall.allow(port: 80, source: alice.record.address)
    try check(firewall.allows(port: 80, source: alice.record.address), "scoped firewall rule allows source")
    try check(!firewall.allows(port: 80, source: bob.record.address), "scoped firewall rule denies other source")
    firewall.allow(port: 443, source: nil)
    try check(firewall.allows(port: 443, source: bob.record.address), "any-source firewall rule")
    let migrated = try JSONDecoder().decode(FirewallRule.self, from: Data("{\"port\":7000}".utf8))
    try check(migrated.protocolKind == .overlay, "protocol-less firewall rule migrates to overlay")
    firewall.allow(protocol: .tcp, port: 8080, source: nil)
    try check(firewall.allows(protocol: .tcp, port: 8080, source: carol.record.address) &&
              !firewall.allows(protocol: .udp, port: 8080, source: carol.record.address), "firewall protocols remain distinct")

    let syn = ipv6Packet(source: alice.record.address, destination: bob.record.address, next: 6,
                         payload: tcpHeader(source: 50_000, destination: 8080, flags: 0x02))
    let synInfo = try IPv6PacketParser.parse(syn)
    try check(synInfo.transport == .tcp && synInfo.destinationPort == 8080, "TCP packet parsed")
    let udp = ipv6Packet(source: alice.record.address, destination: bob.record.address, next: 17,
                         payload: udpHeader(source: 50_001, destination: 5353, payload: Data([1, 2, 3])))
    try check(try IPv6PacketParser.parse(udp).transport == .udp, "UDP packet parsed")
    let echo = ipv6Packet(source: alice.record.address, destination: bob.record.address, next: 58,
                          payload: Data([128, 0, 0, 0, 0, 1, 0, 1]))
    try check(try IPv6PacketParser.parse(echo).icmpType == 128, "ICMPv6 echo parsed")
    let extensionHeader = Data([6, 0, 0, 0, 0, 0, 0, 0]) + tcpHeader(source: 1, destination: 2, flags: 0x02)
    let extended = ipv6Packet(source: alice.record.address, destination: bob.record.address, next: 0, payload: extensionHeader)
    try check(try IPv6PacketParser.parse(extended).transportOffset == 48, "IPv6 extension header parsed safely")
    var badOption = extensionHeader; badOption[2] = 5; badOption[3] = 10
    var badOptionRejected = false
    do { _ = try IPv6PacketParser.parse(ipv6Packet(source: alice.record.address, destination: bob.record.address,
                                                    next: 0, payload: badOption)) } catch { badOptionRejected = true }
    try check(badOptionRejected, "malformed IPv6 option rejected")
    var fragmentRejected = false
    do { _ = try IPv6PacketParser.parse(ipv6Packet(source: alice.record.address, destination: bob.record.address,
                                                    next: 44, payload: Data(repeating: 0, count: 8))) } catch { fragmentRejected = true }
    try check(fragmentRejected, "IPv6 fragments rejected")
    var malformedRejected = false; var malformedUDP = udp; malformedUDP[44] = 0; malformedUDP[45] = 8
    do { _ = try IPv6PacketParser.parse(malformedUDP) } catch { malformedRejected = true }
    try check(malformedRejected, "malformed UDP length rejected")
    try check(synInfo.source != carol.record.address, "packet source exposes spoof mismatch")

    let packetFrame = InnerFrame(kind: .ipv6Packet, payload: syn, sourceRecord: alice.record, replayID: UUID())
    let routedNative = RoutedPacket(source: alice.record.address, destination: carol.record.address,
                                    sealed: try alice.seal(packetFrame, to: carol.record), trafficClass: .nativeIPv6)
    try check(router.routes()[carol.record.address]?.nextHop == bob.record.address, "native packet selects three-node next hop")
    var relayCannotDecrypt = false
    do { _ = try bob.open(routedNative.sealed, expectedSource: alice.record.address) } catch { relayCannotDecrypt = true }
    let destinationFrame = try carol.open(routedNative.sealed, expectedSource: alice.record.address)
    try check(relayCannotDecrypt && destinationFrame.payload == syn, "native packet relays end-to-end encrypted")
    var replayCache = Set<UUID>(); let packetReplayID = destinationFrame.replayID!
    try check(replayCache.insert(packetReplayID).inserted && !replayCache.insert(packetReplayID).inserted,
              "native replay UUID rejects duplicate")
    var nativeBytes = 0
    for _ in 0..<900 {
        let large = ipv6Packet(source: alice.record.address, destination: carol.record.address, next: 17,
                               payload: udpHeader(source: 40_000, destination: 8080,
                                                  payload: Data(repeating: 0x5a, count: 1_160)))
        _ = try IPv6PacketParser.parse(large); nativeBytes += large.count
    }
    try check(nativeBytes > 1_048_576, "native packet path exceeds legacy 1 MiB stream cap")
    try check(router.ingest(try alice.makeLinkState(sequence: 2, neighbors: [carol.record.address])) &&
              router.ingest(try carol.makeLinkState(sequence: 2, neighbors: [alice.record.address])), "route-change topology accepted")
    try check(router.routes()[carol.record.address]?.nextHop == carol.record.address,
              "retransmitted native packet uses current route")

    let nativeRules = MeshFirewall(); let stateful = NativePacketFirewall(rules: nativeRules, maximumFlows: 4)
    try check(stateful.allowOutbound(synInfo), "outbound TCP allowed and tracked")
    let synAck = ipv6Packet(source: bob.record.address, destination: alice.record.address, next: 6,
                            payload: tcpHeader(source: 8080, destination: 50_000, flags: 0x12))
    try check(stateful.allowInbound(try IPv6PacketParser.parse(synAck)), "TCP return packet admitted")
    let unsolicited = ipv6Packet(source: bob.record.address, destination: alice.record.address, next: 6,
                                 payload: tcpHeader(source: 9000, destination: 9001, flags: 0x02))
    try check(!stateful.allowInbound(try IPv6PacketParser.parse(unsolicited)), "new inbound TCP denied by default")
    nativeRules.allow(protocol: .tcp, port: 9001, source: bob.record.address)
    try check(stateful.allowInbound(try IPv6PacketParser.parse(unsolicited)), "source-scoped inbound TCP admitted")

    let aaaaResponse = try MeshDNSCodec.response(to: dnsQuery("alice.mesh", type: 28), aliases: ["alice": alice.record.address])
    try check(aaaaResponse[7] == 1 && aaaaResponse.suffix(16) == alice.record.address.bytes, "DNS AAAA response")
    let noData = try MeshDNSCodec.response(to: dnsQuery("alice.mesh", type: 1), aliases: ["alice": alice.record.address])
    try check(noData[3] & 0x0f == 0 && noData[7] == 0, "DNS A query returns NODATA")
    let notFound = try MeshDNSCodec.response(to: dnsQuery("missing.mesh", type: 28), aliases: [:])
    try check(notFound[3] & 0x0f == 3, "DNS unknown alias returns NXDOMAIN")
    var compressionRejected = false
    var compressed = dnsQuery("x.mesh", type: 28); compressed.replaceSubrange(12..<20, with: Data([0xc0, 0x0c, 0, 28, 0, 1]))
    do { _ = try MeshDNSCodec.response(to: compressed, aliases: [:]) } catch { compressionRejected = true }
    try check(compressionRejected, "DNS compression loop rejected")

    let desired: Set<MeshAddress> = [bob.record.address, carol.record.address]
    try HelperRouteSet.validate(desired, local: alice.record.address)
    let routeDiff = HelperRouteSet.diff(current: [bob.record.address], desired: desired)
    try check(routeDiff.additions == [carol.record.address] && routeDiff.removals.isEmpty, "helper route-set diff is idempotent")
    var localRouteRejected = false
    do { try HelperRouteSet.validate([alice.record.address], local: alice.record.address) } catch { localRouteRejected = true }
    try check(localRouteRejected, "helper rejects a route to the local identity")

    var sockets = [Int32](repeating: -1, count: 2); var pipeFDs = [Int32](repeating: -1, count: 2)
    guard socketpair(AF_UNIX, SOCK_STREAM, 0, &sockets) == 0, pipe(&pipeFDs) == 0 else { throw TestFailure("socket setup") }
    defer { sockets.forEach { Darwin.close($0) }; pipeFDs.forEach { Darwin.close($0) } }
    var peerUID: uid_t = 0, peerGID: gid_t = 0
    try check(umn_get_peer_eid(sockets[0], &peerUID, &peerGID) == 0 && peerUID == getuid(), "helper peer UID authentication primitive")
    let marker = Data("FD\n".utf8)
    let fdSent = marker.withUnsafeBytes { umn_send_fd(sockets[0], pipeFDs[0], $0.baseAddress, marker.count) }
    var receivedFD: Int32 = -1; var fdBuffer = [UInt8](repeating: 0, count: 8)
    let fdCount = umn_recv_fd(sockets[1], &receivedFD, &fdBuffer, fdBuffer.count)
    defer { if receivedFD >= 0 { Darwin.close(receivedFD) } }
    try check(fdSent == 0 && fdCount == marker.count && receivedFD >= 0, "utun descriptor passing primitive")

    let framed = try FrameCodec.encode(WireEnvelope(message: .keepalive))
    try check(try FrameCodec.bodyLength(from: framed.prefix(4)) == framed.count - 4, "frame length prefix")
    let decoded = try FrameCodec.decode(WireEnvelope.self, body: framed.dropFirst(4))
    if decoded.version == 2, case .keepalive = decoded.message { try check(true, "wire v2 frame round trip") }
    else { throw TestFailure("wire frame round trip") }
    var zeroRejected = false
    do { _ = try FrameCodec.bodyLength(from: Data([0, 0, 0, 0])) } catch { zeroRejected = true }
    try check(zeroRejected, "zero-length frame rejected")

    print("1..\(passed)")
    print("all core self-tests passed")
} catch {
    fputs("not ok - \(error.localizedDescription)\n", stderr)
    exit(1)
}