package com.tokilabs.proto_socket

import com.tokilabs.proto_socket.packets.TestData
import java.io.File
import java.net.ServerSocket
import java.security.KeyFactory
import java.security.KeyStore
import java.security.SecureRandom
import java.security.cert.CertificateFactory
import java.security.cert.X509Certificate
import java.security.spec.PKCS8EncodedKeySpec
import java.util.Base64
import java.util.Random
import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.atomic.AtomicLong
import javax.net.ssl.KeyManagerFactory
import javax.net.ssl.SSLContext
import javax.net.ssl.TrustManagerFactory
import kotlin.test.fail

fun testParserMap(): ParserMap =
    mapOf(
        typeNameOf<TestData>() to { TestData.parseFrom(it) },
    )

fun testData(): TestData =
    TestData.newBuilder()
        .setIndex(1)
        .setMessage("ping")
        .build()

private val nextPort = AtomicInteger(15000 + Random().nextInt(10000))
private val nextSecureRandomSeed = AtomicLong(0x50524f544f534f43L)

fun freePort(): Int {
    while (true) {
        val port = nextPort.getAndIncrement()
        try {
            ServerSocket(port).use { return port }
        } catch (_: java.io.IOException) {
            // port is in use, try next
        }
    }
}

fun createTestSslContexts(certPath: String, keyPath: String): Pair<SSLContext, SSLContext> {
    val certFactory = CertificateFactory.getInstance("X.509")
    val cert = File(certPath).inputStream().use {
        certFactory.generateCertificate(it) as X509Certificate
    }
    val keyPem = File(keyPath).readText()
        .replace("-----BEGIN PRIVATE KEY-----", "")
        .replace("-----END PRIVATE KEY-----", "")
        .replace("\\s+".toRegex(), "")
    val keyBytes = Base64.getDecoder().decode(keyPem)
    val privateKey = KeyFactory.getInstance("RSA")
        .generatePrivate(PKCS8EncodedKeySpec(keyBytes))

    val keyStore = KeyStore.getInstance("PKCS12")
    keyStore.load(null, null)
    keyStore.setKeyEntry("toki", privateKey, CharArray(0), arrayOf(cert))
    val keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm())
    keyManagerFactory.init(keyStore, CharArray(0))
    val serverContext = SSLContext.getInstance("TLSv1.2")
    serverContext.init(keyManagerFactory.keyManagers, null, testSecureRandom())

    val trustStore = KeyStore.getInstance("PKCS12")
    trustStore.load(null, null)
    trustStore.setCertificateEntry("toki", cert)
    val trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm())
    trustManagerFactory.init(trustStore)
    val clientContext = SSLContext.getInstance("TLSv1.2")
    clientContext.init(null, trustManagerFactory.trustManagers, testSecureRandom())
    return Pair(serverContext, clientContext)
}

private fun testSecureRandom(): SecureRandom =
    SecureRandom.getInstance("SHA1PRNG").apply {
        setSeed(nextSecureRandomSeed.getAndIncrement())
    }

suspend fun waitForCondition(
    timeoutMs: Long = 15_000L,
    intervalMs: Long = 10L,
    message: String,
    predicate: () -> Boolean,
) {
    val deadline = System.nanoTime() + timeoutMs * 1_000_000L
    while (System.nanoTime() < deadline) {
        if (predicate()) return
        kotlinx.coroutines.delay(intervalMs)
    }
    fail(message)
}
