/*
 * Copyright (c) 2025. Sergey S. Chernov - All Rights Reserved
 *
 *  You may use, distribute and modify this code under the
 * terms of the private license, which you must obtain from the author
 *
 * To obtain the license, contact the author: https://t.me/real_sergeych or email to
 * real dot sergeych at gmail.
 */

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.bipack.BipackDecoder
import net.sergeych.bipack.BipackEncoder
import net.sergeych.crypto2.*
import net.sergeych.tools.bipack
import net.sergeych.tools.biunpack
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 }

        // public key ID should use key bytes instead
        assertContentEquals(pk1.keyBytes, pk1.id.binaryTag.take(32).toUByteArray())
        assertContentEquals(pk1.keyBytes, sk1.id.binaryTag.take(32).toUByteArray())

        assertEquals(pk1, pk1.id.id.asPublicKey)
    }

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

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

        val sk1 = j.decodeFromString<DecryptingSecretKey>(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 = sy1 as UniversalKey
        val usy2 = sy2 as UniversalKey
        val usy3 = sy3 as UniversalKey

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

        val sk1 = DecryptingSecretKey.new()
        val sk2 = DecryptingSecretKey(sk1.keyBytes)
        val sk3 = DecryptingSecretKey.new()

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

        var usk1 = sk1 as UniversalKey
        var usk2 = sk2 as UniversalKey
        var usk3 = sk3 as UniversalKey
        val usk4 = sy3 as UniversalKey

        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))
    }

    @Test
    fun testKiSerialization() = runTest {
        initCrypto()
        val k: UniversalKey = SymmetricKey.new()
        val d = bipack(k)
        val k1: UniversalKey = biunpack<UniversalKey>(d) as SymmetricKey
        assertEquals(k, k1)
    }

    @Test
    fun verifyingKeySerializationTest() = runTest {
        initCrypto()
        val k = SigningSecretKey.new()
        val s1 = pack(k)
        val s2 = pack(k.verifyingKey)
        val dk1 = unpack<SigningSecretKey>(s1)
        val dk2 = unpack<VerifyingPublicKey>(s2)
        assertEquals(k, dk1)
        assertEquals(k.verifyingKey, dk2)

        // id for public/shared keys should be of the key itself to increase safety.
        // not hashed!
        assertContentEquals(k.verifyingKey.keyBytes, dk2.id.binaryTag.take(32).toUByteArray())
        assertContentEquals(k.verifyingKey.keyBytes, dk1.id.binaryTag.take(32).toUByteArray())
        // and restored from id should be the same:
        assertEquals(k.verifyingKey, dk2.id.id.asVerifyingKey)
    }

    @Test
    fun multiKeyTestSome() = runTest {
        initCrypto()
        val k1 = SigningSecretKey.new()
        val k2 = SigningSecretKey.new()
        val k3 = SigningSecretKey.new()
        val k4 = SigningSecretKey.new()
        val k5 = SigningSecretKey.new()
//        val k6 = SigningSecretKey.new()
        val mk: Multikey = Multikey.Keys(1, setOf(k1.verifyingKey))
        val mk23: Multikey = Multikey.Keys(2, setOf(k1.verifyingKey, k2.verifyingKey, k3.verifyingKey))
        val mk13: Multikey = Multikey.Keys(1, setOf(k1.verifyingKey, k2.verifyingKey, k3.verifyingKey))

        assertTrue { mk.check(k1.verifyingKey) }
        assertFalse { mk.check(k2.verifyingKey) }

        assertTrue { mk23.check(k1.verifyingKey, k2.verifyingKey, k4.verifyingKey) }
        assertTrue { mk23.check(k3.verifyingKey, k2.verifyingKey, k4.verifyingKey) }
        assertFalse { mk23.check(k4.verifyingKey, k2.verifyingKey, k5.verifyingKey) }
        assertTrue { mk13.check(k4.verifyingKey, k2.verifyingKey, k5.verifyingKey) }

        println(pack(mk23).toDump())
        println(pack(mk23).size)

        val smk23: Multikey = Multikey.someOf(2, k1.verifyingKey, k2.verifyingKey, k3.verifyingKey)
//        val smk13: Multikey = Multikey.Keys(1, setOf(k1.verifyingKey, k2.verifyingKey, k3.verifyingKey))

        assertTrue { smk23.check(k1.verifyingKey, k2.verifyingKey, k4.verifyingKey) }
        assertTrue { smk23.check(k3.verifyingKey, k2.verifyingKey, k4.verifyingKey) }
        assertFalse { smk23.check(k4.verifyingKey, k2.verifyingKey, k5.verifyingKey) }
//        assertTrue { smk13.check(k4.verifyingKey, k2.verifyingKey, k5.verifyingKey) }

        println(pack(smk23).toDump())
        println(pack(smk23).size)

        val s1 = k1 or k2 or k3
        println(pack(s1).toDump())
        println(pack(s1).size)
        assertTrue { s1.check(k1.verifyingKey, k2.verifyingKey, k3.verifyingKey) }
        assertTrue { s1.check(k1.verifyingKey) }
        assertTrue { s1.check(k2.verifyingKey) }
        assertTrue { s1.check(k3.verifyingKey) }
        assertFalse { s1.check(k4.verifyingKey) }

        val s2 = (k1 or k2) and k3
        println(pack(s2).toDump())
        println(pack(s2).size)
        assertTrue { s2.check(k1.verifyingKey, k3.verifyingKey) }
        assertTrue { s2.check(k2.verifyingKey, k3.verifyingKey) }
        assertTrue { s2.check(k1.verifyingKey, k2.verifyingKey, k3.verifyingKey) }
        assertFalse { s2.check(k4.verifyingKey) }
        assertFalse { s2.check(k1.verifyingKey) }
        assertFalse { s2.check(k2.verifyingKey) }
        assertFalse { s2.check(k3.verifyingKey) }
        assertFalse { s2.check(k1.verifyingKey, k2.verifyingKey) }

        val s3 = (k1 and k2) or k3
        println(pack(s3).toDump())
        println(pack(s3).size)
        assertTrue { s3.check(k1.verifyingKey, k3.verifyingKey) }
        assertTrue { s3.check(k3.verifyingKey) }
        assertTrue { s3.check(k2.verifyingKey, k1.verifyingKey) }
        assertFalse { s3.check(k1.verifyingKey) }
        assertFalse { s3.check(k2.verifyingKey) }
        assertFalse { s3.check(k1.verifyingKey, k4.verifyingKey) }
    }

    @Test
    fun multiKeyTestAny() = runTest {
        initCrypto()
        val k1 = SigningSecretKey.new()
        val k2 = SigningSecretKey.new()
        val k3 = SigningSecretKey.new()
        val k4 = SigningSecretKey.new()
        val k5 = SigningSecretKey.new()
//        val k6 = SigningSecretKey.new()
        val mk: Multikey = Multikey.AnyKey
        assertTrue { mk.check(k1.verifyingKey) }
        assertTrue { mk.check(k2.verifyingKey) }
        assertTrue { mk.check(k3.verifyingKey) }
        assertTrue { mk.check(k4.verifyingKey) }
        assertTrue { mk.check(k5.verifyingKey) }
    }

    @Test
    fun testCombinedKeys() = runTest {
        initCrypto()
        val k1 = UniversalPrivateKey.new()
        val k2 = UniversalPrivateKey.new()
        val k3: UniversalPrivateKey = unpack(pack(k1))
        assertEquals(k1, k3)
        assertEquals(k1.publicKey, k3.publicKey)
        assertEquals(k1.signingKey, k3.signingKey)
        assertEquals(k1.verifyingKey, k3.verifyingKey)

        val k4: UniversalPublicKey = unpack(pack(k1.publicKey))
        assertEquals(k1.publicKey, k4)
        assertEquals(k1.publicKey.encryptingKey, k4.encryptingKey)
        assertEquals(k1.publicKey.verifyingKey, k4.verifyingKey)

        val data =
            """We hold these truths to be self-evident, that all men are created equal, 
                |that they are endowed by their Creator with certain unalienable Rights, 
                |that among these are Life, Liberty and the pursuit of Happiness."""
                .trimMargin()

        val kr1 = UniversalRing.from(k1)
        val kr2: UniversalRing = UniversalRing.from(k2)

        val bytes = Container.encrypt(data, k2.publicKey)
        assertNull(Container.decrypt<String>(bytes, kr1))
        assertEquals(data, Container.decrypt<String>(bytes, kr2))
    }

    @Test
    fun testEncodedSizes() = runTest {
        initCrypto()
        val x = SigningSecretKey.new()
//        println("key bytes: ${x.keyBytes.size}:\n${x.keyBytes.toDump()}")
        val y = BipackEncoder.encode(x)
//        println("packed: ${y.size}: ${y.toDump()}")
        assertTrue { x.keyBytes.size + 5 > y.size }
        assertEquals(x, BipackDecoder.decode<SigningSecretKey>(y))
        assertContentEquals(x.keyBytes, BipackDecoder.decode<SigningSecretKey>(y).keyBytes)
    }

    @Test
    fun testEncodedSizes2() = runTest {
        initCrypto()
        val x = DecryptingSecretKey.new()
//        println("key bytes: ${x.keyBytes.size}:\n${x.keyBytes.toDump()}")
        val y = BipackEncoder.encode(x)
//        println("packed: ${y.size}: ${y.toDump()}")
        assertTrue { x.keyBytes.size + 5 > y.size }
        assertEquals(x, BipackDecoder.decode<DecryptingSecretKey>(y))
        assertContentEquals(x.keyBytes, BipackDecoder.decode<DecryptingSecretKey>(y).keyBytes)
    }

    @Test
    fun testStringRepresentationAndParse() = runTest {
        initCrypto()
        val k1 = SigningSecretKey.new()
        val k2 = k1.verifyingKey
        val k3 = DecryptingSecretKey.new()
        val k4 = k3.publicKey

        val k5 = UniversalPrivateKey.new()
        val k6 = k5.publicKey

        assertEquals(32, k2.keyBytes.size)
        assertContentEquals(k2.keyBytes, k2.id.id.body)

        val k7 = SymmetricKey.new()
        val k8 = KDF.Complexity.Interactive.derive("super", KDF.Argon.randomSalt())

        fun testToString(k: UniversalKey) {
            val s = k.toString()
            val kx = UniversalKey.parseString(s)
            assertEquals(kx::class, k::class)
            assertContentEquals(k.keyBytes, kx.keyBytes)
            assertEquals(k.id, kx.id)
            assertEquals(k, kx)
        }

        fun testAsString(k: UniversalKey) {
            val s = k.asString()
            val kx = UniversalKey.parseString(s)
            assertEquals(kx::class, k::class)
            assertContentEquals(k.keyBytes, kx.keyBytes)
            assertEquals(k.id, kx.id)
            assertEquals(k, kx)
        }

        testToString(k2)
        testToString(k4)

        for( i in listOf(k1, k2, k3, k4, k5, k6, k7, k8)) testAsString(i)

        val x = VerifyingPublicKey.parse("I1po9Y2I7p2aOxeh4nFyGPm3e0YunBEu1Mo-PmIqP84Evg")
        println(x)
    }
}