import com.ionspin.kotlin.crypto.keyexchange.KeyExchange
import kotlinx.coroutines.*
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ReceiveChannel
import kotlinx.coroutines.channels.SendChannel
import kotlinx.coroutines.test.runTest
import net.sergeych.crypto2.SigningKey
import net.sergeych.crypto2.initCrypto
import net.sergeych.kiloparsec.*
import net.sergeych.mp_logger.Log
import kotlin.test.*

private var dcnt = 0
fun createTestDevice(): Pair<Transport.Device, Transport.Device> {
    val p1 = Channel<UByteArray>(256)
    val p2 = Channel<UByteArray>(256)
    val id = ++dcnt
    val d1 = object : Transport.Device {
        override val input: ReceiveChannel<UByteArray> = p1
        override val output: SendChannel<UByteArray> = p2
        override suspend fun close() {
            p2.close()
        }

        override fun toString(): String {
            return "D1:$id"
        }
    }
    val d2 = object : Transport.Device {
        override val input: ReceiveChannel<UByteArray> = p2
        override val output: SendChannel<UByteArray> = p1
        override suspend fun close() {
            p1.close()
        }

        override fun toString(): String {
            return "D2:$id"
        }
    }
    return d1 to d2
}

class TransportTest {


    @Test
    fun testTransportL0AndEncryption() = runTest {
        initCrypto()
        val cmdPing by command<String, String>()
        val cmdSlow by command<String, String>()

        Log.connectConsole()
        Log.defaultLevel = Log.Level.DEBUG
        val (d1, d2) = createTestDevice()
        val l1 = LocalInterface<Unit>().apply {
            on(cmdPing) {
                "p1: $it"
            }
        }
        val l2 = LocalInterface<Unit>().apply {
            on(cmdPing) {
                "p2: $it"
            }
            on(cmdSlow) {
                delay(100)
                "done"
            }
        }
        val t1 = Transport(d1, l1, Unit)
        val t2 = Transport(d2, l2, Unit)

        val clip = KeyExchange.keypair()
        val serp = KeyExchange.keypair()
        val clisk = KeyExchange.clientSessionKeys(clip.publicKey, clip.secretKey, serp.publicKey)
        val sersk = KeyExchange.serverSessionKeys(serp.publicKey, serp.secretKey, clip.publicKey)
        val pser = KiloParams(true, t1, sersk, Unit, null, t1)
        val pcli = KiloParams(false, t2, clisk, Unit, null, t2)

        assertContentEquals(pcli.token, pser.token)
        assertEquals(pser.decryptString(pcli.encrypt("hello!")), "hello!")
        assertEquals(pser.decryptString(pcli.encrypt("hello!")), "hello!")
        assertEquals(pser.decryptString(pcli.encrypt("hello2!")), "hello2!")
        assertEquals(pser.decryptString(pcli.encrypt("hello3!")), "hello3!")
        assertEquals(pser.decryptString(pcli.encrypt("hello!")), "hello!")

        // test nonce increment
        assertFalse { pcli.encrypt("once") contentEquals pcli.encrypt("once") }

        assertEquals(pcli.decryptString(pser.encrypt("hello!")), "hello!")
        assertEquals(pcli.decryptString(pser.encrypt("hello!")), "hello!")
        assertEquals(pcli.decryptString(pser.encrypt("hello!")), "hello!")
        assertEquals(pcli.decryptString(pser.encrypt("hello!")), "hello!")
        assertEquals(pcli.decryptString(pser.encrypt("hello!")), "hello!")
        assertEquals(pcli.decryptString(pser.encrypt("hello!")), "hello!")


        coroutineScope {
            val j1 = launch { t1.run() }
            val j2 = launch { t2.run() }
            launch {
                assertThrows<RemoteInterface.ClosedException> {
                    t1.call(cmdSlow, "foo1")
                }
            }
            assertEquals("p2: foo", t1.call(cmdPing, "foo"))
            assertEquals("p1: bar", t2.call(cmdPing, "bar"))
            assertEquals("p2: foo", t1.call(cmdPing, "foo"))
            j1.cancelAndJoin()
            j2.cancelAndJoin()
        }
        d1.close()
        d2.close()
    }

    @Test
    fun testConnection() = runTest {
        initCrypto()

        val cmdPing by command<String, String>()
        val cmdPush by command<String, String>()
        val cmdGetToken by command<Unit, UByteArray>()
        Log.connectConsole()
//        Log.defaultLevel = Log.Level.DEBUG
        val (d1, d2) = createTestDevice()
        val serverInterface = KiloInterface<String>().apply {
            on(cmdPing) {
                "pong! [$it]"
            }
            on(cmdGetToken) {
                sessionToken
            }
            registerError { IllegalStateException() }
            registerError { IllegalArgumentException(it) }
        }
        val kiloServerConnection = KiloServerConnection(serverInterface, d1, "server session")
        launch { kiloServerConnection.run() }

        val clientInterface = KiloInterface<String>().apply {
            on(cmdPush) {
                "server push: $it"
            }
            on(cmdPing) {
                "client pong: $it"
            }
        }
        val client = KiloClientConnection(clientInterface, d2, "client session")
        launch { client.run() }
        assertEquals("pong! [hello]", client.call(cmdPing, "hello"))
        assertEquals("pong! [foo]", client.call(cmdPing, "foo"))
        assertEquals("client pong: foo", kiloServerConnection.call(cmdPing, "foo"))
        assertEquals("server push: bar", kiloServerConnection.call(cmdPush, "bar"))

        assertContentEquals(client.token(), client.call(cmdGetToken))
        d1.close()
        d2.close()
    }

    @Test
    fun testClient() = runTest {
        initCrypto()

        val cmdPing by command<String, String>()
        val cmdPush by command<String, String>()
        val cmdGetToken by command<Unit, UByteArray>()
        val cmdGetClientId by command<Unit, SigningKey.Public?>()
        val cmdChainCallServer1 by command<String, String>()
        val cmdChainCallClient1 by command<String, String>()
        val cmdChainCallServer2 by command<String, String>()
        val cmdChainCallClient2 by command<String, String>()
        Log.connectConsole()
//        Log.defaultLevel = Log.Level.DEBUG
        val (d1, d2) = createTestDevice()

        val serverId = SigningKey.Secret.pair()
        val clientId = SigningKey.Secret.pair()

        val serverInterface = KiloInterface<String>().apply {
            on(cmdPing) {
                "pong! [$it]"
            }
            on(cmdGetToken) {
                sessionToken
            }
            on(cmdGetClientId) {
                remoteIdentity
            }
            on(cmdChainCallServer1) {
                remote.call(cmdChainCallClient1, it + "-s1")
            }
            on(cmdChainCallServer2) {
                remote.call(cmdChainCallClient2, "$it-s2")
            }
            registerError { IllegalStateException() }
            registerError { IllegalArgumentException(it) }
        }
        val kiloServerConnection = KiloServerConnection(serverInterface, d1, "server session", serverId.signing)
        launch { kiloServerConnection.run() }

        var cnt = 0
        val client = KiloClient {
            session { "client session!" }
            secretIdKey = clientId.signing
            local {
                on(cmdPush) {
                    "server push: $it"
                }
                on(cmdPing) {
                    "client pong: $it"
                }
                on(cmdChainCallClient1) {
                    remote.call(cmdChainCallServer2,"$it-c1")
                }
                on(cmdChainCallClient2) { "$it-c2" }
            }
            connect {
                if (cnt++ > 0) {
                    cancel()
                    fail("connect called once again")
                }
                d2
            }
        }
        assertEquals("pong! [hello]", client.call(cmdPing, "hello"))
        assertEquals("pong! [foo]", client.call(cmdPing, "foo"))
        assertEquals("client pong: foo", kiloServerConnection.call(cmdPing, "foo"))
        assertEquals("server push: bar", kiloServerConnection.call(cmdPush, "bar"))

        assertEquals("**-s1-c1-s2-c2", client.call(cmdChainCallServer1, "**"))
        d1.close()
        d2.close()
        client.close()
//        assertEquals(1, connectionCounter)
    }

    @Test
    fun testAuthentication() = runTest {
        initCrypto()

        val cmdPing by command<String, String>()
        val cmdPush by command<String, String>()
        val cmdGetToken by command<Unit, UByteArray>()
        Log.connectConsole()
        Log.defaultLevel = Log.Level.DEBUG
        val (d1, d2) = createTestDevice()
        val serverInterface = KiloInterface<String>().apply {
            on(cmdPing) {
                "pong! [$it]"
            }
            on(cmdGetToken) {
                sessionToken
            }
            registerError { IllegalStateException() }
            registerError { IllegalArgumentException(it) }
        }
        val kiloServerConnection = KiloServerConnection(serverInterface, d1, "server session")
        launch { kiloServerConnection.run() }

        var cnt = 0
        val client = KiloClient {
            session { "client session!" }
            local {
                on(cmdPush) {
                    "server push: $it"
                }
                on(cmdPing) {
                    "client pong: $it"
                }
            }
            connect {
                if (cnt++ > 0) {
                    cancel()
                    fail("connect called once again")
                }
                d2
            }
        }
        assertEquals("pong! [hello]", client.call(cmdPing, "hello"))
        assertEquals("pong! [foo]", client.call(cmdPing, "foo"))
        assertEquals("client pong: foo", kiloServerConnection.call(cmdPing, "foo"))
        assertEquals("server push: bar", kiloServerConnection.call(cmdPush, "bar"))

        assertContentEquals(client.call(cmdGetToken), client.token())
        client.close()
        d1.close()
        d2.close()
    }
}