//
//  CryptoChatModelMsg.swift
//  Wallet
//
//  Created by Saveliy Stavitsky on 8/17/20.
//  Copyright © 2020 List. All rights reserved.
//

import Foundation
import CryptoSwift
import microecc

extension CryptoChat.Model {

    struct Msg: Codable {
        let from: String
        let to: String
        var iv: String
        let ephemPublicKey: String
        var cipherText: String
        var mac: String
        
        var ephemPublicKeyTemp: [UInt8]?
        var ephemPrivateKey: [UInt8]?
        var msg: String?
        
        private enum CodingKeys: String, CodingKey {
            case from
            case to
            case iv
            case ephemPublicKey = "ephem_key"
            case cipherText = "cipher_text"
            case mac
        }

        let curve = uECC_secp256k1()

        init(from: String, to: String, iv: String, ephemPublicKey: String, cipherText: String, mac: String) {
            self.from = from
            self.to = to
            self.iv = iv
            self.ephemPublicKey = ephemPublicKey
            self.cipherText = cipherText
            self.mac = mac
        }

        // Encrypt
        init(from: String, to: String, msg: String) {
            self.from = from
            self.to = to

            self.msg = msg

            // Generate ephem public and private key pair into bytes array
            var ephemPublicKey = Array(repeating: UInt8(0), count: Int(uECC_curve_public_key_size(curve)) * 2)
            var ephemPrivateKey = Array(repeating: UInt8(0), count: Int(uECC_curve_private_key_size(curve)))
            uECC_make_key(&ephemPublicKey, &ephemPrivateKey, curve) 
            ephemPublicKey = Data(bytes: ephemPublicKey, count: Int(uECC_curve_public_key_size(curve))).bytes

            self.ephemPrivateKey = ephemPrivateKey
            self.ephemPublicKeyTemp = ephemPublicKey

            // Calculate compressed EphemPublicKey
            let compressedEphemPublicKey = ([4] + ephemPublicKey).toData().toCompressedPublicKey ?? Data()
            self.ephemPublicKey = compressedEphemPublicKey.toEosioK1PublicKey

            self.iv = ""
            self.cipherText = ""
            self.mac = ""
        }
        
        // Init and Encrypt
        init(from: String, to: String, msg: String, receiverPublicKey: String) throws {
            self.init(from: from, to: to, msg: msg)
            try self.encrypt(receiverPublicKey: receiverPublicKey)
        }
        
        // Encrypt
        mutating func encrypt(receiverPublicKey: String) throws {
            // Read end decompress receiver public key into bytes array for future use in alg
            let receiverPublicKeyCompressed = try Data(eosioPublicKey: receiverPublicKey).bytes
            var receiverPublicKeyDecompressed = Array(repeating: UInt8(0), count: 64)
            uECC_decompress(receiverPublicKeyCompressed, &receiverPublicKeyDecompressed, curve)
            
            // Calculate symetric shared key for encryption
            var sharedKey = Array(repeating: UInt8(0), count: 32)
            uECC_shared_secret(receiverPublicKeyDecompressed, ephemPrivateKey ?? [], &sharedKey, curve)

            // Calculate sharedKey sha512 (64 bytes) and split it to encryptionKey - first 32 bytes and macKey - second 32 bytes
            let hash = sharedKey.sha512()
            let encryptionKey = [UInt8](hash[..<32])
            let macKey = [UInt8](hash[32...])

            // Generate iv random 16 bytes
            let iv = (0..<16).map({ _ in UInt8.random(in: 0..<UInt8.max) })

            // Encrypt msg with encryptionKey and iv
            let ciphertext = try AES(key: encryptionKey, blockMode: CBC(iv: iv), padding: .pkcs7).encrypt((msg ?? "").bytes)
            
            // Calculate HMAC for msg validation
            let compressedEphemPublicKey = ([4] + (ephemPublicKeyTemp ?? [])).toData().toCompressedPublicKey ?? Data()
            let dataToMac = iv + compressedEphemPublicKey.bytes + ciphertext
            let mac = try HMAC(key: macKey, variant: .sha256).authenticate(dataToMac)
            
            self.iv = iv.toHexString()
            self.cipherText = ciphertext.toHexString()
            self.mac = mac.toHexString()
        }

        func decrypt(privateKey: String) throws -> String {
            let privateKeyBytes = try Data(eosioPrivateKey: privateKey).bytes

            // Read end decompress ephem public key into bytes array for future use in alg
            let compressedEphemPublicKey = try Data(eosioPublicKey: self.ephemPublicKey).bytes
            var decompressedEphemPublicKey = Array(repeating: UInt8(0), count: 64)
            uECC_decompress(compressedEphemPublicKey, &decompressedEphemPublicKey, curve)

            // Calculate symetric shared key for encryption
            var sharedKey = Array(repeating: UInt8(0), count: 32)
            uECC_shared_secret(decompressedEphemPublicKey, privateKeyBytes, &sharedKey, curve)

            // Calculate sharedKey sha512 (64 bytes) and split it to encryptionKey - first 32 bytes and macKey - second 32 bytes
            let hash = sharedKey.sha512()
            let encryptionKey = [UInt8](hash[..<32])
            let macKey = [UInt8](hash[32...])

            // Calculate HMAC for msg validation
            if let ivBytes = Data(hexString: iv)?.bytes,
                let cipherTextBytes = Data(hexString: cipherText)?.bytes {
                let dataToMac = ivBytes + compressedEphemPublicKey + cipherTextBytes
                let realMac = try HMAC(key: macKey, variant: .sha256).authenticate(dataToMac)

                if realMac.toHexString().uppercased() == mac.uppercased() {
                    let plaintext = try AES(key: encryptionKey, blockMode: CBC(iv: ivBytes), padding: .pkcs7).decrypt(cipherTextBytes)
                    return String(bytes: plaintext, encoding: .utf8)
                        ?? "<plaintext>*** Wrong message encryption format ***</plaintext>"
                } else {
                    return "<plaintext>\(L10n.CryptoChat.Chats.decodeError)</plaintext>"
                }
            }

            return "<plaintext>*** Wrong iv or cipherText hex format ***</plaintext>"
        }
    }
}