import com.ionspin.kotlin.crypto.util.decodeFromUByteArray
import com.ionspin.kotlin.crypto.util.encodeToUByteArray
import kotlinx.coroutines.test.runTest
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import net.sergeych.crypto2.*
import net.sergeych.utools.now
import net.sergeych.utools.pack
import net.sergeych.utools.unpack
import kotlin.test.*

class KeysTest {
    @Test
    fun testSigningCreationAndMap() = runTest {
        initCrypto()
        val (stk,pbk) = SigningSecretKey.generatePair()

        val x = mapOf( stk to "STK!", pbk to "PBK!")
        assertEquals("STK!", x[stk])
        val s1 = SigningSecretKey(stk.keyBytes)
        assertEquals(stk, s1)
        assertEquals("STK!", x[s1])
        assertEquals("PBK!", x[pbk])

        val data = "8 rays dev!".encodeToUByteArray()
        val data1 = "8 rays dev!".encodeToUByteArray()
        val s = stk.seal(data)
        s.verify(data)

        data1[0] = 0x01u
        assertFalse(s.isValid(data1))
        val p2 = SigningSecretKey.generatePair()
        val p3 = SigningSecretKey.generatePair()

        val ms = SealedBox.create(data, s1) + p2.secretKey

        // non tampered:
        val ms1 = unpack<SealedBox>(pack(ms))
        assertContentEquals(data, ms1.message)
        assertTrue(pbk in ms1)
        assertTrue(p2.publicKey in ms1)
        assertTrue(p3.publicKey !in ms1)

        assertThrows<IllegalSignatureException> {
            unpack<SealedBox>(pack(ms).also { it[3] = 1u })
        }
    }

    @Test
    fun testNonDeterministicSeals() = runTest {
        initCrypto()
        val data = "Welcome to the crazy new world!".encodeToUByteArray()
        val (sk,_) = SigningSecretKey.generatePair()
        val t = now()
        val s1 = Seal.create(sk, data, createdAt = t)
        val s2 = Seal.create(sk, data, createdAt = t)
        val s2bad = Seal.create(sk, data + "!".encodeToUByteArray())
        val s3 = Seal.create(sk, data, createdAt = t, nonDeterministic = true)
        val s4 = Seal.create(sk, data, createdAt = t, nonDeterministic = true)

        for( seal in listOf(s1,s2,s3,s4)) {
            assertTrue { seal.isValid(data) }
            assertTrue { Seal.unpack(seal.packed).isValid(data) }
        }
        assertFalse { s2bad.isValid(data)}
        assertContentEquals(s1.packed, s2.packed)
        assertFalse { s1.packed contentEquals s3.packed }
        assertFalse { s4.packed contentEquals s3.packed }

    }

    @Test
    fun secretEncryptTest() = runTest {
        initCrypto()
        val key = SymmetricKey.new()
        val key1 = SymmetricKey.new()
        assertEquals("hello", key.decrypt(key.encrypt("hello".encodeToUByteArray())).decodeFromUByteArray())
        assertEquals("hello", key.decryptString(key.encrypt("hello")))
        assertEquals("hello", key.decryptObject(key.encryptObject("hello")))
        assertEquals("hello", key.decrypt(key.encrypt("hello".encodeToUByteArray(), 18..334)).decodeFromUByteArray())
        assertEquals("hello", key.decryptString(key.encrypt("hello", 18..334)))
        assertEquals("hello", key.decryptObject(key.encryptObject("hello", 18..334)))
        assertThrows<DecryptionFailedException> {
            key.decrypt(key1.encrypt("hello".encodeToUByteArray())).decodeFromUByteArray()
        }
    }

    @Test
    fun symmetricKeyTest() = runTest {
        initCrypto()
        val k1 = SymmetricKey.new()
        val src = "Buena Vista".encodeToUByteArray()
        val nonce = k1.randomNonce()

        assertContentEquals(src, k1.decryptWithNonce(k1.encryptWithNonce(src, nonce), nonce))
        assertThrows<DecryptionFailedException> {
            val n2 = nonce.copyOf()
            n2[4] =  n2[4].inv()
            k1.decryptWithNonce(k1.encryptWithNonce(src, nonce), n2)
        }

        assertContentEquals(src, k1.decrypt(k1.encrypt(src)))
        assertContentEquals(src, k1.decrypt(k1.encrypt(src, 0..117)))
        assertContentEquals(src, k1.decrypt(k1.encrypt(src, 7..117)))
    }

    @Test
    fun keyExchangeTest() = runTest {
        initCrypto()
        val ske = SafeKeyExchange()
        val cke = SafeKeyExchange()

        val clientSessionKey = cke.clientSessionKey(ske.publicKey)
        val serverSessionKey = ske.serverSessionKey(cke.publicKey)

        val src = "Hello, Dolly!"
        assertEquals(src, serverSessionKey.decryptString(clientSessionKey.encrypt(src)))
        assertEquals(src, serverSessionKey.decryptString(clientSessionKey.encrypt(src)))
        assertEquals(src, serverSessionKey.decryptString(clientSessionKey.encrypt(src)))
        assertEquals(src, serverSessionKey.decryptString(clientSessionKey.encrypt(src)))

        assertEquals(src, clientSessionKey.decryptString(serverSessionKey.encrypt(src)))
        assertEquals(src, clientSessionKey.decryptString(serverSessionKey.encrypt(src)))
        assertEquals(src, clientSessionKey.decryptString(serverSessionKey.encrypt(src)))
        assertEquals(src, clientSessionKey.decryptString(serverSessionKey.encrypt(src)))

        assertContentEquals(clientSessionKey.sessionTag, serverSessionKey.sessionTag)

    }

    @Test
    fun asymmetricKeyTest() = runTest {
        initCrypto()
        val (sk0, pk0) = Asymmetric.generateKeys()
        assertEquals(pk0, sk0.publicKey)

        val (sk1, pk1) = Asymmetric.generateKeys()
        val (sk2, pk2) = Asymmetric.generateKeys()

        val plain = "The fake vaccine kills".encodeToUByteArray()

        var m = pk1.encryptMessage(plain, sk1)
        assertContentEquals(plain, m.decrypt(sk1))
        assertThrows<DecryptionFailedException> {
            assertContentEquals(plain, m.decrypt(sk2))
        }

        m = pk2.encryptAnonymousMessage(plain)
        assertContentEquals(plain, m.decrypt(sk2))
        assertContentEquals(plain, sk2.decrypt(m))
//        assertContentEquals(plain, sk2.decrypt(sk1.encrypt(plain, pk2)))
        assertThrows<DecryptionFailedException> {
            assertContentEquals(plain, m.decrypt(sk1))
        }

        val x1 = pk1.encryptMessage(plain, sk1).encoded
        val x2 = pk1.encryptMessage(plain, sk1).encoded

        assertFalse { x1 contentEquals x2 }
    }
    @Test
    fun asymmetricKeySerializationTest() = runTest {
        initCrypto()
        val (sk0, pk0) = Asymmetric.generateKeys()

//        println(sk0.publicKey)
        val j = Json { prettyPrint = true}

        val sk1 = j.decodeFromString<Asymmetric.SecretKey>(j.encodeToString(sk0))
        assertEquals(sk0, sk1)
        assertEquals(pk0, sk1.publicKey)
//        println(j.encodeToString(sk1))
    }


    @Test
    fun testUniKeys() = runTest {
        initCrypto()
        val sy1 = SymmetricKey.new()
        val sy2 = SymmetricKey(sy1.keyBytes)
        val sy3 = SymmetricKey.new()

        assertEquals(sy1, sy2)
        assertEquals(sy2, sy1)
        assertFalse { sy1 == sy3 }

        assertEquals(sy1, deepCopy(sy1), "symmetric key should be equal to the restored copy")
        assertEquals(sy1.hashCode(), deepCopy(sy1).hashCode(), "hashcode of the restored symmetric key is wrong")

        val usy1 = UniversalKey.from(sy1)
        val usy2 = UniversalKey.from(sy2)
        val usy3 = UniversalKey.from(sy3)

        assertEquals(usy1, usy2)
        assertEquals(usy2, usy1)
        assertFalse { usy1 == usy3 }

        val sk1 = Asymmetric.newSecretKey()
        val sk2 = Asymmetric.SecretKey(sk1.keyBytes)
        val sk3 = Asymmetric.newSecretKey()

        assertEquals(sk1, sk2)
        assertEquals(sk2, sk1)
        assertFalse { sk1 == sk3 }

        var usk1 = UniversalKey.from(sk1)
        var usk2 = UniversalKey.from(sk2)
        var usk3 = UniversalKey.from(sk3)
        val usk4 = UniversalKey.from(sy3)

        assertEquals(usk1, usk2)
        assertEquals(usk2, usk1)
        assertFalse { usk1 == usk3 }

        var a = setOf(usk1, usk2, usk3, usk4)
        var b = setOf(usk1, usk2, usk3, usk4)

        assertEquals(a, b)
        // usk1 and usk2 are equal so set with only one of should be the same
        assertEquals(a, setOf(usk1, usk3, usk4))

        usk1 = deepCopy(usk1)
        usk2 = deepCopy(usk2)
        usk3 = deepCopy(usk3)

        assertEquals(usk1.hashCode(),usk2.hashCode())

        assertEquals(usk1, usk2)
        assertEquals(usk2, usk1)
        assertFalse { usk1 == usk3 }

        a = setOf(usk1, usk2, usk3, usk4)
        b = setOf(usk4, usk3, usk2, usk1)

        assertEquals(3, a.size)
        assertEquals(3, b.size)

        assertEquals(a, b)
        // usk1 and usk2 are equal so set with only one of should be the same
        assertEquals(a, setOf(usk1, usk3, usk4))
    }
}