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.SigningSecretKey
import net.sergeych.crypto2.VerifyingPublicKey
import net.sergeych.crypto2.initCrypto
import net.sergeych.kiloparsec.*
import net.sergeych.mp_logger.Log
import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals
import kotlin.test.fail

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

    class TestException(text: String) : Exception(text)

    @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, VerifyingPublicKey?>()
        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 = SigningSecretKey.generatePair()
        val clientId = SigningSecretKey.generatePair()

        val cmdException by command<Unit, Unit>()
        val cmdRemoteExceptionTest by command<Unit, String>()
        val cmdBreak by command<Unit, Unit>()

        val cmdPushServer by push<String>()
        val pushedFromServer = CompletableDeferred<String>()

        val serverInterface = KiloInterface<String>().apply {
            on(cmdPing) {
                "pong! [$it]"
            }
            on(cmdPushServer) {
                pushedFromServer.complete(it)
            }
            on(cmdGetToken) {
                sessionToken
            }
            on(cmdGetClientId) {
                remoteIdentity
            }
            on(cmdChainCallServer1) {
                remote.call(cmdChainCallClient1, it + "-s1")
            }
            on(cmdChainCallServer2) {
                remote.call(cmdChainCallClient2, "$it-s2")
            }
            on(cmdException) { throw TestException("te1") }
            on(cmdRemoteExceptionTest) {
                try {
                    remote.call(cmdException)
                    "error!"
                } catch (e: TestException) {
                    "ok: ${e.message}"
                }
            }
            on(cmdBreak) { throw LocalInterface.BreakConnectionException() }
            registerError { TestException(it) }
        }
        val kiloServerConnection = KiloServerConnection(
            serverInterface, d1, "server session", serverId.secretKey
        )
        launch { kiloServerConnection.run() }

        var cnt = 0
        val client = KiloClient {
            addErrors(serverInterface)
            session { "client session!" }
            secretIdKey = clientId.secretKey
            local {
                on(cmdPush) {
                    "server push: $it"
                }
                on(cmdPing) {
                    "client pong: $it"
                }
                on(cmdChainCallClient1) {
                    remote.call(cmdChainCallServer2, "$it-c1")
                }
                on(cmdChainCallClient2) { "$it-c2" }
                on(cmdException) { throw TestException("te-local") }
            }
            connect {
                println("Called connect: $cnt")
                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"))

        client.push(cmdPushServer, "42")

        assertEquals("**-s1-c1-s2-c2", client.call(cmdChainCallServer1, "**"))


        assertThrows<TestException> { client.call(cmdException) }
        assertEquals("ok: te-local", client.call(cmdRemoteExceptionTest))

        // wait for push to be received and check
        assertEquals("42", pushedFromServer.await())

        assertThrows<RemoteInterface.ClosedException> {
            client.call(cmdBreak)
        }

        // Note that current transport test is too simple,
        // therefore, we can't test reconnecting, also we need server and client instances
        // not connections, so that's all

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