import Foundation
import CryptoKit

extension SymmetricKey {
    
    var data: Data {
        withUnsafeBytes { Data(Array($0)) }
    }
    
    var base64: String {
        data.base64EncodedString()
    }

    var displayString: String {
        data.hexEncoded.uppercased().split(by: 4).joined(separator: " ")
    }
    
    var codeString: String {
        "    {" +
        withUnsafeBytes {
            return Data(Array($0))
        }.map(String.init).joined(separator: ", ") +
        "},"
    }
}

extension SHA256.Digest {

    var hexEncoded: String {
        Data(map { $0 }).hexEncoded
    }
}

extension String {

    func split(by length: Int) -> [String] {
        var startIndex = self.startIndex
        var results = [Substring]()

        while startIndex < self.endIndex {
            let endIndex = self.index(startIndex, offsetBy: length, limitedBy: self.endIndex) ?? self.endIndex
            results.append(self[startIndex..<endIndex])
            startIndex = endIndex
        }

        return results.map { String($0) }
    }
}

let protocolSalt = "CryptoKit Playgrounds Putting It Together".data(using: .utf8)!

/// Generates an ephemeral key agreement key and performs key agreement to get the shared secret and derive the symmetric encryption key.
func encrypt(_ data: Data, to theirEncryptionKey: Curve25519.KeyAgreement.PublicKey, signedBy ourSigningKey: Curve25519.Signing.PrivateKey) throws ->
    (ephemeralPublicKeyData: Data, ciphertext: Data, signature: Data) {
        let ephemeralKey = Curve25519.KeyAgreement.PrivateKey()
        let ephemeralPublicKey = ephemeralKey.publicKey.rawRepresentation
        
        let sharedSecret = try ephemeralKey.sharedSecretFromKeyAgreement(with: theirEncryptionKey)
        
        let symmetricKey = sharedSecret.hkdfDerivedSymmetricKey(using: SHA256.self,
                                                                salt: protocolSalt,
                                                                sharedInfo: ephemeralPublicKey +
                                                                    theirEncryptionKey.rawRepresentation +
                                                                    ourSigningKey.publicKey.rawRepresentation,
                                                                outputByteCount: 32)
        
        let ciphertext = try ChaChaPoly.seal(data, using: symmetricKey).combined
        let signature = try ourSigningKey.signature(for: ciphertext + ephemeralPublicKey + theirEncryptionKey.rawRepresentation)
        
        return (ephemeralPublicKey, ciphertext, signature)
}