import Foundation
import NIOCore

#if canImport(CryptoKit)
import CryptoKit
#else
import Crypto
#endif

/**
 An authenticated message to or from the device.
 */
struct Message: Equatable, Hashable {

    /// The message authentication code for the message (32 bytes)
    let mac: Data

    /// The message content
    let content: Content

    /**
     Create an authenticated message
     - Parameter mac: The message authentication code
     - Parameter content: The message content
     */
    init(mac: Data, content: Content) {
        self.mac = mac
        self.content = content
    }
}

extension Message: Codable {
    
    enum CodingKeys: Int, CodingKey {
        case mac = 1
        case content = 2
    }
}

extension Message {

    /**
     The message content without authentication.
     */
    struct Content: Equatable, Hashable {

        /// The time of message creation, in UNIX time (seconds since 1970)
        let time: UInt32

        /// The counter of the message (for freshness)
        let id: UInt32

        /**
         Create new message content.
         - Parameter time: The time of message creation,
         - Parameter id: The counter of the message
         */
        init(time: UInt32, id: UInt32) {
            self.time = time
            self.id = id
        }

        /**
         Decode message content from data.

         The data consists of two `UInt32` encoded in little endian format
         - Warning: The sequence must contain at least 8 bytes, or the function will crash.
         - Parameter data: The sequence containing the bytes.
         */
        init<T: Sequence>(decodeFrom data: T) where T.Element == UInt8 {
            self.time = UInt32(data: Data(data.prefix(MemoryLayout<UInt32>.size)))
            self.id = UInt32(data: Data(data.dropFirst(MemoryLayout<UInt32>.size)))
        }

        /// The byte length of an encoded message content
        static var length: Int {
            MemoryLayout<UInt32>.size * 2
        }

        /// The message content encoded to data
        var encoded: Data {
            time.encoded + id.encoded
        }
    }
}

extension Message.Content: Codable {
    
    enum CodingKeys: Int, CodingKey {
        case time = 1
        case id = 2
    }
}

extension Message {

    /// The length of a message in bytes
    static var length: Int {
        SHA256.byteCount + Content.length
    }

    /**
     Decode a message from a byte buffer.
     The buffer must contain at least `Message.length` bytes, or it will return `nil`.
     - Parameter buffer: The buffer containing the bytes.
     */
    init?(decodeFrom buffer: ByteBuffer) {
        guard let data = buffer.getBytes(at: 0, length: Message.length) else {
            return nil
        }
        self.init(decodeFrom: data)
    }

    init?(decodeFrom data: Data, index: inout Int) {
        guard index + Message.length <= data.count else {
            return nil
        }
        self.init(decodeFrom: data.advanced(by: index))
        index += Message.length
    }

    /// The message encoded to data
    var encoded: Data {
        mac + content.encoded
    }

    var bytes: [UInt8] {
        Array(encoded)
    }

    /**
     Create a message from received bytes.
     - Parameter data: The sequence of bytes
     - Note: The sequence must contain at least `Message.length` bytes, or the function will crash.
     */
    init<T: Sequence>(decodeFrom data: T) where T.Element == UInt8 {
        let count = SHA256.byteCount
        self.mac = Data(data.prefix(count))
        self.content = .init(decodeFrom: Array(data.dropFirst(count)))
    }

    /**
     Check if the message contains a valid authentication code
     - Parameter key: The key used to sign the message.
     - Returns: `true`, if the message is valid.
     */
    func isValid(using key: SymmetricKey) -> Bool {
         HMAC<SHA256>.isValidAuthenticationCode(mac, authenticating: content.encoded, using: key)
    }
}

extension Message.Content {

    /**
     Calculate an authentication code for the message content.
     - Parameter key: The key to use to sign the content.
     - Returns: The new message signed with the key.
     */
    func authenticate(using key: SymmetricKey) -> Message {
        let mac = HMAC<SHA256>.authenticationCode(for: encoded, using: key)
        return .init(mac: Data(mac.map { $0 }), content: self)
    }

    /**
     Calculate an authentication code for the message content and convert everything to data.
     - Parameter key: The key to use to sign the content.
     - Returns: The new message signed with the key, serialized to bytes.
     */
    func authenticateAndSerialize(using key: SymmetricKey) -> Data {
        let encoded = self.encoded
        let mac = HMAC<SHA256>.authenticationCode(for: encoded, using: key)
        return Data(mac.map { $0 }) + encoded
    }
}