package com.tokilabs.proto_socket

import com.tokilabs.proto_socket.packets.TestData
import java.io.File
import java.security.KeyFactory
import java.security.KeyStore
import java.security.cert.CertificateFactory
import java.security.cert.X509Certificate
import java.security.spec.PKCS8EncodedKeySpec
import java.net.ServerSocket
import java.util.Base64
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()

fun freePort(): Int =
    ServerSocket(0).use { it.localPort }

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("TLS")
    serverContext.init(keyManagerFactory.keyManagers, null, null)

    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("TLS")
    clientContext.init(null, trustManagerFactory.trustManagers, null)
    return Pair(serverContext, clientContext)
}

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