package com.jihe.punchnet.sdlan.utils

import com.jihe.punchnet.sdlan.config.AESConfig
import com.jihe.punchnet.sdlan.logs.TerminalLogger
import org.bouncycastle.crypto.engines.ChaCha7539Engine
import org.bouncycastle.crypto.params.KeyParameter
import org.bouncycastle.crypto.params.ParametersWithIV
import java.nio.ByteBuffer
import java.nio.ByteOrder
import javax.crypto.Cipher
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec

class AES private constructor () {
    private var _secret: ByteArray = ByteArray(0)
    private var _iv: ByteArray = ByteArray(0)
    private var _algorithm: String = "aes"
    private var _regionId: Long = 0
    companion object {
        private var instance: AES? = null
            get() {
                if (field == null) {
                    field = AES()
                }
                return field
            }

        @Synchronized
        fun getAES(): AES {
            return instance!!
        }
    }

    fun isAuthorized(): Boolean {
        return _secret.size != 0
    }

    fun setSecret(secret: ByteArray?, algorithm: String = "aes", regionId: Long = 0) {
        if (secret == null) {
            _secret = ByteArray(0)
            _iv = ByteArray(0)
            _algorithm = "aes"
            _regionId = 0
        } else {
            _secret = secret
            _iv = _secret.sliceArray(0..<16)
            _algorithm = algorithm.lowercase()
            _regionId = regionId
        }
    }

    fun encrypt(content: ByteBuffer): ByteBuffer? {
        if (!isAuthorized()) {
            return null
        }
        try {
            val input = ByteArray(content.remaining())
            content.get(input)
            val encrypted = encrypt(input) ?: return null
            return ByteBuffer.wrap(encrypted)
        } catch (e: Exception) {
            println("failed to encrypt: ${e.toString()}")
            return null
        }
    }

    fun encrypt(content: ByteArray): ByteArray? {
        if (!isAuthorized()) {
            TerminalLogger.errorf { "not authed, so not encrypting" }
            return null
        }
        try {
            return when (_algorithm) {
                "aes" -> aesCrypt(Cipher.ENCRYPT_MODE, content)
                "chacha20" -> chacha20Crypt(content)
                else -> {
                    TerminalLogger.errorf { "unsupported encryption algorithm: $_algorithm" }
                    null
                }
            }
        } catch (e: Exception) {
            TerminalLogger.errorf {"encrypt failed: $e"}
            return null
        }
    }

    fun decrypt(ciphered: ByteBuffer): ByteBuffer? {
        if (!isAuthorized()) {
            return null
        }
        try {
            val input = ByteArray(ciphered.remaining())
            ciphered.get(input)
            val decrypted = decrypt(input) ?: return null
            return ByteBuffer.wrap(decrypted)
        } catch (e: Exception) {
            return null
        }
    }

    fun decrypt(ciphered: ByteArray): ByteArray? {
        if (!isAuthorized()) {
            return null
        }
        try {
            return when (_algorithm) {
                "aes" -> aesCrypt(Cipher.DECRYPT_MODE, ciphered)
                "chacha20" -> chacha20Crypt(ciphered)
                else -> null
            }
        } catch (e: Exception) {
            return null
        }
    }

    private fun aesCrypt(mode: Int, content: ByteArray): ByteArray {
        val cipher = Cipher.getInstance(AESConfig.CIPHER_ALGORITHM)
        val keyspec = SecretKeySpec(_secret, AESConfig.KEY_SPEC)
        cipher.init(mode, keyspec, IvParameterSpec(_iv))
        return cipher.doFinal(content)
    }

    private fun chacha20Crypt(content: ByteArray): ByteArray {
        val engine = ChaCha7539Engine()
        val nonce = chacha20Nonce()
        engine.init(true, ParametersWithIV(KeyParameter(_secret), nonce))
        val output = ByteArray(content.size)
        engine.processBytes(content, 0, content.size, output, 0)
        return output
    }

    private fun chacha20Nonce(): ByteArray {
        return ByteBuffer.allocate(12)
            .order(ByteOrder.BIG_ENDIAN)
            .putLong(0L)
            .putInt(_regionId.toInt())
            .array()
    }
}

fun byteArray2Hex(array: ByteArray): String {
    var result = mutableListOf<String>()
    for (item in array) {
        result.add(String.format("0x%02x", item.toInt() and 0xff))
    }
    return result.joinToString(" ")
}
