import Foundation
import CryptoKit

public enum UMNError: Error, LocalizedError {
    case invalidAddress, invalidFrame, invalidIdentity, cryptoFailure
    case noRoute, denied, unbound, transferLimit, timeout
    case message(String)

    public var errorDescription: String? {
        switch self {
        case .invalidAddress: "invalid mesh address"
        case .invalidFrame: "invalid protocol frame"
        case .invalidIdentity: "identity verification failed"
        case .cryptoFailure: "cryptographic operation failed"
        case .noRoute: "no route to destination"
        case .denied: "destination firewall denied the connection"
        case .unbound: "destination port has no bound service"
        case .transferLimit: "the 1 MiB stream limit was exceeded"
        case .timeout: "operation timed out"
        case .message(let value): value
        }
    }
}

public struct MeshAddress: Hashable, Codable, Comparable, CustomStringConvertible, Sendable {
    public let bytes: Data

    public init(bytes: Data) throws {
        guard bytes.count == 16, bytes.first == 0xfd else { throw UMNError.invalidAddress }
        self.bytes = bytes
    }

    public init(_ text: String) throws {
        var raw = in6_addr()
        guard inet_pton(AF_INET6, text, &raw) == 1 else { throw UMNError.invalidAddress }
        let data = withUnsafeBytes(of: &raw) { Data($0) }
        try self.init(bytes: data)
    }

    public static func derive(from signingPublicKey: Data) throws -> MeshAddress {
        let digest = Data(SHA256.hash(data: signingPublicKey))
        return try MeshAddress(bytes: Data([0xfd]) + digest.prefix(15))
    }

    public var description: String {
        var raw = in6_addr()
        _ = withUnsafeMutableBytes(of: &raw) { bytes.copyBytes(to: $0) }
        var output = [CChar](repeating: 0, count: Int(INET6_ADDRSTRLEN))
        guard inet_ntop(AF_INET6, &raw, &output, socklen_t(INET6_ADDRSTRLEN)) != nil else {
            return bytes.map { String(format: "%02x", $0) }.joined()
        }
        return String(cString: output)
    }

    public static func < (lhs: MeshAddress, rhs: MeshAddress) -> Bool {
        lhs.bytes.lexicographicallyPrecedes(rhs.bytes)
    }
}

public struct NodeRecord: Codable, Hashable, Sendable {
    public let address: MeshAddress
    public let signingPublicKey: Data
    public let agreementPublicKey: Data

    public init(address: MeshAddress, signingPublicKey: Data, agreementPublicKey: Data) {
        self.address = address
        self.signingPublicKey = signingPublicKey
        self.agreementPublicKey = agreementPublicKey
    }

    public func validate() -> Bool {
        guard signingPublicKey.count == 32, agreementPublicKey.count == 32,
              let derived = try? MeshAddress.derive(from: signingPublicKey) else { return false }
        return derived == address
    }
}

public struct LinkState: Codable, Hashable, Sendable {
    public let origin: NodeRecord
    public let sequence: UInt64
    public let neighbors: [MeshAddress]
    public let signature: Data

    public init(origin: NodeRecord, sequence: UInt64, neighbors: [MeshAddress], signature: Data) {
        self.origin = origin
        self.sequence = sequence
        self.neighbors = neighbors.sorted()
        self.signature = signature
    }

    public func signingBytes() throws -> Data {
        struct Unsigned: Codable { let origin: NodeRecord; let sequence: UInt64; let neighbors: [MeshAddress] }
        let encoder = PropertyListEncoder(); encoder.outputFormat = .binary
        return try encoder.encode(Unsigned(origin: origin, sequence: sequence, neighbors: neighbors.sorted()))
    }

    public func validate() -> Bool {
        guard origin.validate(), let key = try? Curve25519.Signing.PublicKey(rawRepresentation: origin.signingPublicKey),
              let bytes = try? signingBytes() else { return false }
        return key.isValidSignature(signature, for: bytes)
    }
}

public enum PacketKind: String, Codable, Sendable {
    case pingRequest, pingReply, text, streamOpen, streamData, streamAck, streamClose, streamReset, ipv6Packet, error
}

public struct InnerFrame: Codable, Sendable {
    public let kind: PacketKind
    public let streamID: UInt64
    public let port: UInt16?
    public let sequence: UInt64
    public let acknowledgment: UInt64
    public let payload: Data
    public let sourceRecord: NodeRecord
    /// Required for ipv6Packet frames. It is covered by the end-to-end signature.
    public let replayID: UUID?

    public init(kind: PacketKind, streamID: UInt64 = 0, port: UInt16? = nil,
                sequence: UInt64 = 0, acknowledgment: UInt64 = 0,
                payload: Data = Data(), sourceRecord: NodeRecord, replayID: UUID? = nil) {
        self.kind = kind; self.streamID = streamID; self.port = port
        self.sequence = sequence; self.acknowledgment = acknowledgment
        self.payload = payload; self.sourceRecord = sourceRecord; self.replayID = replayID
    }
}

public struct SealedPayload: Codable, Sendable {
    public let ephemeralPublicKey: Data
    public let combinedCiphertext: Data

    public init(ephemeralPublicKey: Data, combinedCiphertext: Data) {
        self.ephemeralPublicKey = ephemeralPublicKey
        self.combinedCiphertext = combinedCiphertext
    }
}

public struct RoutedPacket: Codable, Sendable {
    public enum TrafficClass: String, Codable, Sendable { case overlay, nativeIPv6 }
    public let id: UUID
    public let source: MeshAddress
    public let destination: MeshAddress
    public var hopLimit: UInt8
    public let sealed: SealedPayload
    public let trafficClass: TrafficClass

    public init(id: UUID = UUID(), source: MeshAddress, destination: MeshAddress,
                hopLimit: UInt8 = 16, sealed: SealedPayload, trafficClass: TrafficClass = .overlay) {
        self.id = id; self.source = source; self.destination = destination
        self.hopLimit = hopLimit; self.sealed = sealed; self.trafficClass = trafficClass
    }
}

public enum WireMessage: Codable, Sendable {
    case hello(NodeRecord)
    case linkState(LinkState)
    case packet(RoutedPacket)
    case keepalive

    private enum CodingKeys: String, CodingKey { case type, payload }
    private enum Kind: String, Codable { case hello, linkState, packet, keepalive }

    public init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        switch try c.decode(Kind.self, forKey: .type) {
        case .hello: self = .hello(try c.decode(NodeRecord.self, forKey: .payload))
        case .linkState: self = .linkState(try c.decode(LinkState.self, forKey: .payload))
        case .packet: self = .packet(try c.decode(RoutedPacket.self, forKey: .payload))
        case .keepalive: self = .keepalive
        }
    }

    public func encode(to encoder: Encoder) throws {
        var c = encoder.container(keyedBy: CodingKeys.self)
        switch self {
        case .hello(let value): try c.encode(Kind.hello, forKey: .type); try c.encode(value, forKey: .payload)
        case .linkState(let value): try c.encode(Kind.linkState, forKey: .type); try c.encode(value, forKey: .payload)
        case .packet(let value): try c.encode(Kind.packet, forKey: .type); try c.encode(value, forKey: .payload)
        case .keepalive: try c.encode(Kind.keepalive, forKey: .type)
        }
    }
}

public struct WireEnvelope: Codable, Sendable {
    public static let currentVersion: UInt16 = 2
    public let version: UInt16
    public let message: WireMessage

    public init(message: WireMessage, version: UInt16 = currentVersion) {
        self.version = version; self.message = message
    }
}