package com.peachpasswords.sync

import java.util.concurrent.atomic.AtomicBoolean
import kotlinx.coroutines.async
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.runBlocking
import org.json.JSONObject
import org.junit.Assert.assertEquals
import org.junit.Assert.assertFalse
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertNull
import org.junit.Assert.assertTrue
import org.junit.Test

class RelayAdoptionHostTest {
    private class MemoryRelaySocket(
        private val inbound: Channel<ByteArray>,
        private val outbound: Channel<ByteArray>,
    ) : RelayBinarySocket {
        override fun send(frame: ByteArray) {
            check(outbound.trySend(frame.copyOf()).isSuccess)
        }

        override suspend fun receive(): ByteArray = inbound.receive()
        override fun close(code: Int, reason: String) {
            outbound.close()
        }
    }

    /**
     * Models OkHttp's asynchronous send queue: each receive flushes the prior
     * outbound frame, but ADOPT_APPLIED has no following receive and therefore
     * reaches the peer only through a graceful close.
     */
    private class QueuedFinalFrameSocket(
        private val delegate: RelayBinarySocket,
    ) : RelayBinarySocket {
        private val queued = mutableListOf<ByteArray>()
        var gracefullyClosed = false
            private set
        var aborted = false
            private set

        override fun send(frame: ByteArray) {
            queued += frame.copyOf()
        }

        override suspend fun receive(): ByteArray {
            flush()
            return delegate.receive()
        }

        override fun close(code: Int, reason: String) {
            aborted = true
            queued.clear()
            delegate.close(code, reason)
        }

        override fun closeAfterQueuedFrames(code: Int, reason: String) {
            gracefullyClosed = true
            flush()
            delegate.closeAfterQueuedFrames(code, reason)
        }

        private fun flush() {
            queued.forEach(delegate::send)
            queued.clear()
        }
    }

    private fun socketPair(): Pair<RelayBinarySocket, RelayBinarySocket> {
        val leftToRight = Channel<ByteArray>(Channel.UNLIMITED)
        val rightToLeft = Channel<ByteArray>(Channel.UNLIMITED)
        return MemoryRelaySocket(rightToLeft, leftToRight) to
            MemoryRelaySocket(leftToRight, rightToLeft)
    }

    private fun identity(name: String, publicKey: ByteArray) = LocalDeviceIdentity(
        deviceId = "$name-id",
        deviceName = name,
        deviceType = if (name == "source") DeviceType.ANDROID else DeviceType.EXTENSION,
        longTermPublicKeyHex = Rendezvous.bytesToHex(publicKey),
        signingPublicKeyHex = "cd".repeat(32),
        wsPort = 49142,
        relayUrl = "wss://relay.example.test",
    )

    @Test fun `relay invitation runs Noise and persists session-derived pair state`() = runBlocking {
        val (sourceSocket, receiverSocket) = socketPair()
        val token = "AAECAwQFBgcICQoLDA0ODw"
        val relay = "wss://relay.example.test"
        val sourceRoot = ByteArray(32) { (it + 1).toByte() }
        val sourceKeys = localSyncKeypair(randomLocalIdentityCredentials("source-id"))
        val receiverKeys = localSyncKeypair(randomLocalIdentityCredentials("receiver-id"))
        val sourceIdentity = identity("source", sourceKeys.publicKey)
        val receiverIdentity = identity("receiver", receiverKeys.publicKey)
        var sourceStored: PairedDevice? = null
        var receiverStored: PairedDevice? = null

        val sourceJob = async {
            runRelayAdoptionSourceSession(
                socket = sourceSocket,
                expectedToken = token,
                relayUrl = relay,
                ourStaticPrivateKey = sourceKeys.secretKey,
                ourStaticPublicKey = sourceKeys.publicKey,
                localIdentity = sourceIdentity,
                sourceDeps = object : AdoptionSourceV2Deps {
        override suspend fun rotateSourceIdentity() = sourceIdentity
        override fun hasDeviceIdentityCollision(device: PairedDevice) = false
        override suspend fun persistPendingIdentity(device: PairedDevice, nowMs: Long) = true
        override suspend fun completePendingIdentity(device: PairedDevice) = true
        override fun snapshotCapable() = true

                    override val sourceIdentity = sourceIdentity
                    override suspend fun rootKeyHex() = Rendezvous.bytesToHex(sourceRoot)
                    override suspend fun vaultId() = computeVaultId(sourceRoot)
                    override suspend fun license() = "PPK-test"
                    override suspend fun onConsent(request: AdoptionFrame.Request) = true
                    override suspend fun onConsent(request: AdoptionFrame.RequestV2) = true
                    override suspend fun upsertPairedDevice(device: PairedDevice) {
                        sourceStored = device
                    }
                },
                confirm = { sas, request ->
                    assertEquals("receiver", request.deviceName)
                    assertEquals(8, sas.length)
                    true
                },
            )
        }

        receiverSocket.send(
            JSONObject().put("type", "PEACH_HELLO").put("token", token)
                .toString().toByteArray(),
        )
        val ack = JSONObject(String(receiverSocket.receive()))
        val sourcePublic = hexToPubkey(ack.getString("ephemeralPublicKeyHex"))
        val noise = noiseNKInitiator(
            sourcePublic,
            receiverSocket::send,
            receiverSocket::receive,
        )
        val channel = encryptedChannel(receiverSocket, noise)
        val received = requestAdoption(
            channel,
            object : AdoptionReceiverV2Deps {
        override fun localIdentity() = receiverIdentity
        override suspend fun rotateLocalIdentity() = localIdentity()
        override suspend fun validateSource(prepare: AdoptionFrame.Prepare) {}
        override suspend fun hasSourceIdentityCollision(prepare: AdoptionFrame.Prepare) = false

                override val deviceName = "receiver"
                override val deviceType = DeviceType.EXTENSION
                override val pairState = AdoptionPairState(
                    Rendezvous.bytesToHex(Rendezvous.deriveSecret(noise.handshakeHash)),
                    relay,
                )
                override fun generateDeviceId() = "session-request"
                override suspend fun confirmSas() = true
                override suspend fun createVaultFromRootKey(rootKeyHex: String) =
                    computeVaultId(hexToPubkey(rootKeyHex))
                override suspend fun deriveLocalIdentity() = receiverIdentity
                override suspend fun upsertPairedDevice(device: PairedDevice) {
                    receiverStored = device
                }
                override fun onError(error: String) = throw AssertionError(error)
            },
        )

        assertNotNull(received)
        assertNotNull(sourceJob.await())
        assertEquals(relay, receiverStored?.relayUrl)
        assertEquals(relay, sourceStored?.relayUrl)
        assertEquals(receiverStored?.rendezvousSecretHex, sourceStored?.rendezvousSecretHex)
        assertEquals(64, receiverStored?.rendezvousSecretHex?.length)
    }

    @Test fun `source rejection never requests the root`() = runBlocking {
        val (sourceSocket, receiverSocket) = socketPair()
        val token = "AAECAwQFBgcICQoLDA0ODw"
        val root = ByteArray(32) { (it + 7).toByte() }
        val sourceKeys = localSyncKeypair(randomLocalIdentityCredentials("source-id"))
        val rootRequested = AtomicBoolean(false)

        val sourceJob = async {
            runRelayAdoptionSourceSession(
                sourceSocket,
                token,
                "wss://relay.example.test",
                sourceKeys.secretKey,
                sourceKeys.publicKey,
                identity("source", sourceKeys.publicKey),
                object : AdoptionSourceV2Deps {
        override suspend fun rotateSourceIdentity() = sourceIdentity
        override fun hasDeviceIdentityCollision(device: PairedDevice) = false
        override suspend fun persistPendingIdentity(device: PairedDevice, nowMs: Long) = true
        override suspend fun completePendingIdentity(device: PairedDevice) = true
        override fun snapshotCapable() = true

                    override val sourceIdentity = identity("source", sourceKeys.publicKey)
                    override suspend fun rootKeyHex(): String {
                        rootRequested.set(true)
                        return Rendezvous.bytesToHex(root)
                    }
                    override suspend fun vaultId() = computeVaultId(root)
                    override suspend fun onConsent(request: AdoptionFrame.Request) = true
                    override suspend fun onConsent(request: AdoptionFrame.RequestV2) = true
                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit
                },
                confirm = { _, _ -> false },
            )
        }
        receiverSocket.send(
            JSONObject().put("type", "PEACH_HELLO").put("token", token)
                .toString().toByteArray(),
        )
        val ack = JSONObject(String(receiverSocket.receive()))
        val noise = noiseNKInitiator(
            hexToPubkey(ack.getString("ephemeralPublicKeyHex")),
            receiverSocket::send,
            receiverSocket::receive,
        )
        val channel = encryptedChannel(receiverSocket, noise)
        sendFrame(channel, AdoptionFrame.Request("session-request", "receiver", "extension"))

        assertNull(sourceJob.await())
        assertFalse(rootRequested.get())
    }

    @Test fun `a pre-reset source still runs the v1 transcript`() = runBlocking {
        val (sourceSocket, receiverSocket) = socketPair()
        val token = "AAECAwQFBgcICQoLDA0ODw"
        val root = ByteArray(32) { (it + 9).toByte() }
        val sourceKeys = localSyncKeypair(randomLocalIdentityCredentials("source-id"))
        val rootRequested = AtomicBoolean(false)
        val prompted = AtomicBoolean(false)
        val outdated = AtomicBoolean(false)

        val sourceJob = async {
            sourceSession(
                socket = sourceSocket,
                token = token,
                root = root,
                sourceKeys = sourceKeys,
                rootRequested = rootRequested,
                confirm = { _, _ -> prompted.set(true); false },
                postReset = { false },
                onPeerOutdated = { outdated.set(true) },
            )
        }
        receiverSocket.send(
            JSONObject().put("type", "PEACH_HELLO").put("token", token)
                .toString().toByteArray(),
        )
        val ack = JSONObject(String(receiverSocket.receive()))
        val noise = noiseNKInitiator(
            hexToPubkey(ack.getString("ephemeralPublicKeyHex")),
            receiverSocket::send,
            receiverSocket::receive,
        )
        val channel = encryptedChannel(receiverSocket, noise)
        sendFrame(channel, AdoptionFrame.Request("session-request", "receiver", "extension"))

        assertNull(sourceJob.await())
        assertTrue("v1 is still offered before a reset", prompted.get())
        assertFalse("v1 is not an outdated peer before a reset", outdated.get())
        assertFalse("a declined prompt never reads the root", rootRequested.get())
    }

    @Test fun `a post-reset source refuses v1 before any prompt or grant`() = runBlocking {
        val (sourceSocket, receiverSocket) = socketPair()
        val token = "AAECAwQFBgcICQoLDA0ODw"
        val root = ByteArray(32) { (it + 13).toByte() }
        val sourceKeys = localSyncKeypair(randomLocalIdentityCredentials("source-id"))
        val rootRequested = AtomicBoolean(false)
        val prompted = AtomicBoolean(false)
        val outdated = AtomicBoolean(false)

        val sourceJob = async {
            sourceSession(
                socket = sourceSocket,
                token = token,
                root = root,
                sourceKeys = sourceKeys,
                rootRequested = rootRequested,
                confirm = { _, _ -> prompted.set(true); true },
                postReset = { true },
                onPeerOutdated = { outdated.set(true) },
            )
        }
        receiverSocket.send(
            JSONObject().put("type", "PEACH_HELLO").put("token", token)
                .toString().toByteArray(),
        )
        val ack = JSONObject(String(receiverSocket.receive()))
        val noise = noiseNKInitiator(
            hexToPubkey(ack.getString("ephemeralPublicKeyHex")),
            receiverSocket::send,
            receiverSocket::receive,
        )
        val channel = encryptedChannel(receiverSocket, noise)
        sendFrame(channel, AdoptionFrame.Request("session-request", "receiver", "extension"))

        assertNull(sourceJob.await())
        assertTrue("the refusal must be reported as an outdated peer", outdated.get())
        assertFalse("no SAS prompt for a peer that is refused", prompted.get())
        assertFalse("the root must never be read for a refused peer", rootRequested.get())
    }

    /** Runs the production source session with only the deps these two cases need. */
    private suspend fun sourceSession(
        socket: RelayBinarySocket,
        token: String,
        root: ByteArray,
        sourceKeys: SyncIdentityKeypair,
        rootRequested: AtomicBoolean,
        confirm: suspend (String, AdoptionFrame.Request) -> Boolean,
        postReset: () -> Boolean,
        onPeerOutdated: () -> Unit,
    ): PairedDevice? = runRelayAdoptionSourceSession(
        socket,
        token,
        "wss://relay.example.test",
        sourceKeys.secretKey,
        sourceKeys.publicKey,
        identity("source", sourceKeys.publicKey),
        object : AdoptionSourceV2Deps {
            override suspend fun rotateSourceIdentity() = sourceIdentity
            override fun hasDeviceIdentityCollision(device: PairedDevice) = false
            override suspend fun persistPendingIdentity(device: PairedDevice, nowMs: Long) = true
            override suspend fun completePendingIdentity(device: PairedDevice) = true
            override fun snapshotCapable() = true

            override val sourceIdentity = identity("source", sourceKeys.publicKey)
            override suspend fun rootKeyHex(): String {
                rootRequested.set(true)
                return Rendezvous.bytesToHex(root)
            }

            override suspend fun vaultId() = computeVaultId(root)
            override suspend fun onConsent(request: AdoptionFrame.Request) = true
            override suspend fun onConsent(request: AdoptionFrame.RequestV2) = true
            override suspend fun upsertPairedDevice(device: PairedDevice) = Unit
        },
        confirm = confirm,
        postReset = postReset,
        onPeerOutdated = onPeerOutdated,
    )

    @Test fun `receiver session completes real Noise adoption with matching pair state`() = runBlocking {
        val (sourceSocket, receiverSocket) = socketPair()
        val token = "AAECAwQFBgcICQoLDA0ODw"
        val relay = "wss://relay.example.test"
        val root = ByteArray(32) { (it + 3).toByte() }
        val sourceKeys = localSyncKeypair(randomLocalIdentityCredentials("source-id"))
        val receiverKeys = localSyncKeypair(randomLocalIdentityCredentials("receiver-id"))
        val sourceIdentity = identity("source", sourceKeys.publicKey)
        val receiverIdentity = identity("receiver", receiverKeys.publicKey)
        var sourceStored: PairedDevice? = null
        var receiverStored: PairedDevice? = null
        var sourceSas = ""
        var receiverSas = ""

        val sourceJob = async {
            runRelayAdoptionSourceSession(
                sourceSocket,
                token,
                relay,
                sourceKeys.secretKey,
                sourceKeys.publicKey,
                sourceIdentity,
                object : AdoptionSourceV2Deps {
        override suspend fun rotateSourceIdentity() = sourceIdentity
        override fun hasDeviceIdentityCollision(device: PairedDevice) = false
        override suspend fun persistPendingIdentity(device: PairedDevice, nowMs: Long) = true
        override suspend fun completePendingIdentity(device: PairedDevice) = true
        override fun snapshotCapable() = false

                    override val sourceIdentity = sourceIdentity
                    override suspend fun rootKeyHex() = Rendezvous.bytesToHex(root)
                    override suspend fun vaultId() = computeVaultId(root)
                    override suspend fun license() = "PPK-test"
                    override suspend fun onConsent(request: AdoptionFrame.Request) = true
                    override suspend fun onConsent(request: AdoptionFrame.RequestV2) = true
                    override suspend fun upsertPairedDevice(device: PairedDevice) {
                        sourceStored = device
                    }
                },
                confirm = { sas, _ ->
                    sourceSas = sas
                    true
                },
            )
        }

        val queuedReceiverSocket = QueuedFinalFrameSocket(receiverSocket)
        val received = adoptFromRelaySocket(
            socket = queuedReceiverSocket,
            invitation = PairingInvitation(
                relay,
                token,
                noiseStaticPublicKeyHex = Rendezvous.bytesToHex(sourceKeys.publicKey),
            ),
            deps = object : AdoptionReceiverV2Deps {
        override fun localIdentity() = receiverIdentity
        override suspend fun rotateLocalIdentity() = localIdentity()
        override suspend fun validateSource(prepare: AdoptionFrame.Prepare) {}
        override suspend fun hasSourceIdentityCollision(prepare: AdoptionFrame.Prepare) = false

                override val deviceName = "receiver"
                override val deviceType = DeviceType.ANDROID
                override fun generateDeviceId() = "session-only-request"
                override suspend fun confirmSas() = error("session wrapper owns SAS")
                override suspend fun createVaultFromRootKey(rootKeyHex: String): String {
                    assertEquals(Rendezvous.bytesToHex(root), rootKeyHex)
                    return computeVaultId(hexToPubkey(rootKeyHex))
                }
                override suspend fun deriveLocalIdentity() = receiverIdentity
                override suspend fun upsertPairedDevice(device: PairedDevice) {
                    receiverStored = device
                }
                override fun onError(error: String) = throw AssertionError(error)
            },
            confirm = { sas ->
                receiverSas = sas
                true
            },
        )

        assertNotNull(received)
        assertNotNull(sourceJob.await())
        assertTrue(queuedReceiverSocket.gracefullyClosed)
        assertFalse(queuedReceiverSocket.aborted)
        assertEquals(sourceSas, receiverSas)
        assertEquals(relay, sourceStored?.relayUrl)
        assertEquals(relay, receiverStored?.relayUrl)
        assertEquals(sourceStored?.rendezvousSecretHex, receiverStored?.rendezvousSecretHex)
    }

    @Test fun `receiver rejection is encrypted and source never reads root`() = runBlocking {
        val token = "AAECAwQFBgcICQoLDA0ODw"
        val relay = "wss://relay.example.test"
        val root = ByteArray(32) { (it + 9).toByte() }
        val sourceKeys = localSyncKeypair(randomLocalIdentityCredentials("source-id"))
        val receiverKeys = localSyncKeypair(randomLocalIdentityCredentials("receiver-id"))
        val receiverIdentity = identity("receiver", receiverKeys.publicKey)
        val (sourceSocket, receiverSocket) = socketPair()
        val rootRequested = AtomicBoolean(false)

        val sourceJob = async {
            runRelayAdoptionSourceSession(
                sourceSocket,
                token,
                relay,
                sourceKeys.secretKey,
                sourceKeys.publicKey,
                identity("source", sourceKeys.publicKey),
                object : AdoptionSourceV2Deps {
        override suspend fun rotateSourceIdentity() = sourceIdentity
        override fun hasDeviceIdentityCollision(device: PairedDevice) = false
        override suspend fun persistPendingIdentity(device: PairedDevice, nowMs: Long) = true
        override suspend fun completePendingIdentity(device: PairedDevice) = true
        override fun snapshotCapable() = true

                    override val sourceIdentity = identity("source", sourceKeys.publicKey)
                    override suspend fun rootKeyHex(): String {
                        rootRequested.set(true)
                        return Rendezvous.bytesToHex(root)
                    }
                    override suspend fun vaultId() = computeVaultId(root)
                    override suspend fun onConsent(request: AdoptionFrame.Request) = true
                    override suspend fun onConsent(request: AdoptionFrame.RequestV2) = true
                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit
                },
                confirm = { _, _ -> true },
            )
        }
        var created = false
        val received = runRelayAdoptionReceiverSession(
            receiverSocket,
            PairingInvitation(
                relay,
                token,
                noiseStaticPublicKeyHex = Rendezvous.bytesToHex(sourceKeys.publicKey),
            ),
            object : AdoptionReceiverV2Deps {
        override fun localIdentity() = receiverIdentity
        override suspend fun rotateLocalIdentity() = localIdentity()
        override suspend fun validateSource(prepare: AdoptionFrame.Prepare) {}
        override suspend fun hasSourceIdentityCollision(prepare: AdoptionFrame.Prepare) = false

                override val deviceName = "receiver"
                override val deviceType = DeviceType.ANDROID
                override fun generateDeviceId() = "session-only-request"
                override suspend fun confirmSas() = error("session wrapper owns SAS")
                override suspend fun createVaultFromRootKey(rootKeyHex: String): String {
                    created = true
                    return computeVaultId(hexToPubkey(rootKeyHex))
                }
                override suspend fun deriveLocalIdentity() =
                    identity("receiver", sourceKeys.publicKey)
                override suspend fun upsertPairedDevice(device: PairedDevice) = Unit
                override fun onError(error: String) {
                    assertEquals("ADOPT_SAS_REJECTED", error)
                }
            },
            confirm = { false },
        )

        assertNull(received)
        assertNull(sourceJob.await())
        assertFalse(rootRequested.get())
        assertFalse(created)
    }

    @Test fun `malformed receiver ACK fails before any vault mutation`() = runBlocking {
        val (unusedSource, receiverSocket) = socketPair()
        unusedSource.send(
            JSONObject()
                .put("type", "PEACH_HELLO_ACK")
                .put("ephemeralPublicKeyHex", "ab")
                .toString()
                .toByteArray(),
        )
        var created = false
        val result = runCatching {
            runRelayAdoptionReceiverSession(
                receiverSocket,
                PairingInvitation(
                    "wss://relay.example.test",
                    "AAECAwQFBgcICQoLDA0ODw",
                ),
                object : AdoptionReceiverV2Deps {
        override fun localIdentity() = LocalDeviceIdentity(generateDeviceId(), deviceName, deviceType, "", "")
        override suspend fun rotateLocalIdentity() = localIdentity()
        override suspend fun validateSource(prepare: AdoptionFrame.Prepare) {}
        override suspend fun hasSourceIdentityCollision(prepare: AdoptionFrame.Prepare) = false

                    override val deviceName = "receiver"
                    override val deviceType = DeviceType.ANDROID
                    override fun generateDeviceId() = "request"
                    override suspend fun confirmSas() = false
                    override suspend fun createVaultFromRootKey(rootKeyHex: String): String {
                        created = true
                        return ""
                    }
                    override suspend fun deriveLocalIdentity() =
                        error("must not derive identity")
                    override suspend fun upsertPairedDevice(device: PairedDevice) = Unit
                    override fun onError(error: String) = Unit
                },
                confirm = { true },
            )
        }
        assertTrue(result.isFailure)
        assertFalse(created)
    }

    private fun encryptedChannel(
        socket: RelayBinarySocket,
        noise: NoiseNKSession,
    ): PairingTransport = object : PairingTransport {
        override suspend fun send(msg: ByteArray) = socket.send(noise.encrypt(msg))
        override suspend fun receive(): ByteArray = noise.decrypt(socket.receive())
        override suspend fun close() = socket.close()
    }
}