package com.peachpasswords.sync

import kotlinx.coroutines.async
import kotlinx.coroutines.awaitAll
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.withTimeout
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import okio.ByteString.Companion.toByteString
import java.io.IOException
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.coroutines.resume
import kotlin.coroutines.resumeWithException

const val DEFAULT_PROBE_TIMEOUT_MS: Long = 500
const val DEFAULT_PROBE_BATCH: Int = 16
const val DEFAULT_CONNECT_TIMEOUT_MS: Long = 4_000
const val SUBNET_PROBE_COOLDOWN_MS: Long = 60_000

data class ChirpSyncSession(
    val remoteStaticPublicKey: ByteArray,
    val send: (ByteArray) -> Unit,
    val receive: suspend () -> ByteArray,
    val close: suspend () -> Unit,
)

data class SyncEndpoint(
    val ip: String,
    val port: Int,
)

private val lastProbeAt = mutableMapOf<String, Long>()
private val webSocketClient = OkHttpClient.Builder().build()

suspend fun openChirpSyncSession(
    endpoint: SyncEndpoint,
    ourLongTermPrivateKey: ByteArray,
    ourLongTermPublicKey: ByteArray,
    expectedRemotePubkeyHex: String,
    connectTimeoutMs: Long = DEFAULT_CONNECT_TIMEOUT_MS,
): ChirpSyncSession {
    val connection = connectWebSocket(endpoint, connectTimeoutMs)

    val session = try {
        noiseXXInitiator(
            ourLongTermPrivateKey,
            ourLongTermPublicKey,
            send = connection::send,
            receive = connection::receive,
        )
    } catch (err: Throwable) {
        connection.close()
        throw err
    }

    val remoteHex = pubkeyToHex(session.remoteStaticPublicKey)
    if (remoteHex != expectedRemotePubkeyHex) {
        connection.close()
        throw IllegalStateException(
            "paired-device static key mismatch (got ${remoteHex.take(16)}…, expected ${expectedRemotePubkeyHex.take(16)}…)",
        )
    }
    voidHex(expectedRemotePubkeyHex)

    return ChirpSyncSession(
        remoteStaticPublicKey = session.remoteStaticPublicKey,
        send = { plaintext -> connection.send(session.encrypt(plaintext)) },
        receive = { session.decrypt(connection.receive()) },
        close = { connection.close() },
    )
}

fun subnetProbeOnCooldown(
    expectedRemotePubkeyHex: String,
    now: Long = System.currentTimeMillis(),
    cooldownMs: Long = SUBNET_PROBE_COOLDOWN_MS,
): Boolean {
    val last = lastProbeAt[expectedRemotePubkeyHex]
    return last != null && now - last < cooldownMs
}

fun resetSubnetProbeCooldown() {
    lastProbeAt.clear()
}

suspend fun probeSubnetForPeer(
    subnet: String,
    port: Int,
    ourLongTermPrivateKey: ByteArray,
    ourLongTermPublicKey: ByteArray,
    expectedRemotePubkeyHex: String,
    batch: Int = DEFAULT_PROBE_BATCH,
    timeoutMs: Long = DEFAULT_PROBE_TIMEOUT_MS,
    cooldownMs: Long = SUBNET_PROBE_COOLDOWN_MS,
    now: Long = System.currentTimeMillis(),
): SyncEndpoint? {
    if (subnetProbeOnCooldown(expectedRemotePubkeyHex, now, cooldownMs)) {
        return null
    }
    lastProbeAt[expectedRemotePubkeyHex] = now

    for (start in 1..254 step batch) {
        val end = minOf(start + batch - 1, 254)
        val results = coroutineScope {
            (start..end).map { i ->
                val ip = "$subnet.$i"
                async {
                    probeOne(
                        endpoint = SyncEndpoint(ip = ip, port = port),
                        ourLongTermPrivateKey = ourLongTermPrivateKey,
                        ourLongTermPublicKey = ourLongTermPublicKey,
                        expectedRemotePubkeyHex = expectedRemotePubkeyHex,
                        timeoutMs = timeoutMs,
                    )
                }
            }.awaitAll()
        }
        val hit = results.firstOrNull { it != null }
        if (hit != null) return hit
    }
    return null
}

private suspend fun probeOne(
    endpoint: SyncEndpoint,
    ourLongTermPrivateKey: ByteArray,
    ourLongTermPublicKey: ByteArray,
    expectedRemotePubkeyHex: String,
    timeoutMs: Long,
): SyncEndpoint? {
    return try {
        val session = withTimeout(timeoutMs) {
            openChirpSyncSession(
                endpoint = endpoint,
                ourLongTermPrivateKey = ourLongTermPrivateKey,
                ourLongTermPublicKey = ourLongTermPublicKey,
                expectedRemotePubkeyHex = expectedRemotePubkeyHex,
                connectTimeoutMs = timeoutMs,
            )
        }
        session.close()
        endpoint
    } catch (_: Throwable) {
        null
    }
}

private suspend fun connectWebSocket(
    endpoint: SyncEndpoint,
    timeoutMs: Long,
): RawWebSocketConnection {
    return withTimeout(timeoutMs) {
        suspendCancellableCoroutine { cont ->
            val request = Request.Builder()
                .url("ws://${endpoint.ip}:${endpoint.port}/")
                .build()
            val connection = RawWebSocketConnection()
            val listener = object : WebSocketListener() {
                override fun onMessage(webSocket: WebSocket, bytes: okio.ByteString) {
                    connection.onMessage(bytes)
                }

                override fun onOpen(webSocket: WebSocket, response: Response) {
                    connection.install(webSocket)
                    if (cont.isActive) cont.resume(connection) else webSocket.cancel()
                }

                override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
                    connection.onClosed(IOException("WebSocket closed"))
                }

                override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
                    connection.onFailure(IOException("WebSocket connect to ${endpoint.ip}:${endpoint.port} errored", t))
                    if (cont.isActive) {
                        cont.resumeWithException(IOException("WebSocket connect to ${endpoint.ip}:${endpoint.port} errored", t))
                    }
                }
            }
            val ws = webSocketClient.newWebSocket(request, listener)
            connection.install(ws)
            cont.invokeOnCancellation {
                try {
                    ws.cancel()
                } catch (_: Throwable) {
                    // ignore
                }
            }
        }
    }
}

private class RawWebSocketConnection {
    private val inbox = Channel<ByteArray>(Channel.UNLIMITED)
    private val closed = AtomicBoolean(false)
    private var ws: WebSocket? = null

    fun install(webSocket: WebSocket) {
        ws = webSocket
    }

    fun onMessage(bytes: okio.ByteString) {
        inbox.trySend(bytes.toByteArray())
    }

    @OptIn(ExperimentalCoroutinesApi::class)
    fun onClosed(cause: Throwable) {
        if (closed.compareAndSet(false, true)) {
            inbox.close(cause)
        }
    }

    fun onFailure(cause: Throwable) {
        onClosed(cause)
    }

    fun send(bytes: ByteArray) {
        val socket = ws ?: throw IOException("WebSocket not open")
        if (!socket.send(bytes.toByteString())) {
            throw IOException("WebSocket send failed")
        }
    }

    suspend fun receive(): ByteArray {
        if (closed.get() && inbox.isEmpty) {
            throw IOException("WebSocket closed before frame")
        }
        return try {
            inbox.receive()
        } catch (err: kotlinx.coroutines.channels.ClosedReceiveChannelException) {
            throw IOException("WebSocket closed before frame", err)
        }
    }

    fun close() {
        try {
            onClosed(IOException("WebSocket closed"))
            ws?.close(1000, null)
        } catch (_: Throwable) {
            // ignore
        }
    }
}

private fun voidHex(value: String) {
    // Keep the expected pubkey source pinned in the call graph.
    hexToPubkey(value)
}