package net.hivetechnology.o20010_06.vpn

import kotlinx.coroutines.CompletableDeferred
import net.hivetechnology.o20010_06.model.RuntimeTarget
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import org.json.JSONObject
import java.io.ByteArrayInputStream
import java.security.KeyStore
import java.security.MessageDigest
import java.security.cert.CertificateFactory
import java.security.cert.X509Certificate
import java.util.Locale
import java.util.concurrent.TimeUnit
import javax.net.ssl.SSLContext
import javax.net.ssl.TrustManagerFactory
import javax.net.ssl.X509TrustManager

class RuntimeSocket(private val target: RuntimeTarget) : AutoCloseable {
    private var socket: WebSocket? = null
    private var client: OkHttpClient? = null
    private val disconnected = CompletableDeferred<Unit>()
    @Volatile var lastError: String? = null
        private set
    @Volatile private var observedPath = "unknown"
    @Volatile private var observedPathAtMs = 0L

    suspend fun connect(): Boolean {
        val opened = CompletableDeferred<Boolean>()
        lastError = null
        val builder = OkHttpClient.Builder().pingInterval(30, TimeUnit.SECONDS).connectTimeout(8, TimeUnit.SECONDS)
        configureTls(builder)
        client = builder.build()
        val scheme = if (target.secure) "wss" else "ws"
        val request = Request.Builder().url("$scheme://${target.host}:${target.port}${target.path}").build()
        socket = client!!.newWebSocket(request, object : WebSocketListener() {
            override fun onOpen(webSocket: WebSocket, response: Response) {
                opened.complete(true)
            }
            override fun onFailure(webSocket: WebSocket, throwable: Throwable, response: Response?) {
                lastError = throwable.message ?: response?.message
                opened.complete(false)
                disconnected.complete(Unit)
            }
            override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
                lastError = reason.takeIf { it.isNotBlank() }
                disconnected.complete(Unit)
            }
            override fun onMessage(webSocket: WebSocket, text: String) {
                val message = runCatching { JSONObject(text) }.getOrNull() ?: return
                if (message.optString("type") != "vpn-path-status" ||
                    message.optInt("version") != 1
                ) return
                val path = message.optString("observedPath")
                if (path !in setOf("p2p", "relay", "unknown")) return
                observedPath = path
                observedPathAtMs = android.os.SystemClock.elapsedRealtime()
            }
        })
        return opened.await()
    }

    suspend fun awaitDisconnect() = disconnected.await()

    fun pathObservation(nowMs: Long = android.os.SystemClock.elapsedRealtime()): String =
        if (observedPathAtMs > 0L && nowMs - observedPathAtMs <= 5_000L) observedPath else "unknown"

    private fun configureTls(builder: OkHttpClient.Builder) {
        if (!target.secure) return
        val pem = target.tls.caCertificatePem
        val expectedFingerprint = target.tls.certificateSha256?.normalizedFingerprint()
        if (pem.isNullOrBlank() && expectedFingerprint.isNullOrBlank()) return

        val delegate = if (!pem.isNullOrBlank()) trustManagerForPem(pem) else platformTrustManager()
        val manager = object : X509TrustManager {
            override fun getAcceptedIssuers(): Array<X509Certificate> = delegate.acceptedIssuers
            override fun checkClientTrusted(chain: Array<X509Certificate>, authType: String) =
                delegate.checkClientTrusted(chain, authType)
            override fun checkServerTrusted(chain: Array<X509Certificate>, authType: String) {
                if (!expectedFingerprint.isNullOrBlank()) {
                    val actual = MessageDigest.getInstance("SHA-256").digest(chain.first().encoded)
                        .joinToString("") { "%02x".format(it) }
                    require(actual == expectedFingerprint) { "Runtime certificate fingerprint mismatch" }
                } else {
                    delegate.checkServerTrusted(chain, authType)
                }
            }
        }
        val context = SSLContext.getInstance("TLS")
        context.init(null, arrayOf(manager), null)
        builder.sslSocketFactory(context.socketFactory, manager)
        if (!expectedFingerprint.isNullOrBlank()) builder.hostnameVerifier { _, _ -> true }
    }

    private fun trustManagerForPem(pem: String): X509TrustManager {
        val certificate = CertificateFactory.getInstance("X.509")
            .generateCertificate(ByteArrayInputStream(pem.toByteArray()))
        val store = KeyStore.getInstance(KeyStore.getDefaultType()).apply {
            load(null)
            setCertificateEntry("runtime", certificate)
        }
        return trustManager(store)
    }

    private fun platformTrustManager(): X509TrustManager = trustManager(null)

    private fun trustManager(store: KeyStore?): X509TrustManager {
        val factory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm())
        factory.init(store)
        return factory.trustManagers.filterIsInstance<X509TrustManager>().first()
    }

    override fun close() {
        disconnected.complete(Unit)
        socket?.close(1000, "VPN stopped")
        socket = null
        client?.dispatcher?.executorService?.shutdown()
        client = null
    }
}

private fun String.normalizedFingerprint(): String =
    lowercase(Locale.ROOT).replace(Regex("[^0-9a-f]"), "")
