package com.peachpasswords.sync

import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.io.DataInputStream
import java.io.EOFException
import java.io.InputStream
import java.io.OutputStream

/**
 * Length-delimited frame IO for a pairing/sync connection.
 *
 * Originally declared alongside the Chirp audio pairing that has since been
 * removed, but it was never audio-specific — it is just "send me a frame, give
 * me a frame", which the Noise NK pipeline and the adoption flow both build on.
 * It lives here now, with the frame helpers that use it.
 */
interface PairingTransport {
    suspend fun send(msg: ByteArray)
    suspend fun receive(): ByteArray
    suspend fun close()
}

/**
 * Length-prefixed framing for the C7 TCP transport, shared by [PairingServer]
 * and [PairingTransport] so neither duplicates the wire format.
 *
 * Frame = [4-byte big-endian length][payload]. This is the framing the existing
 * ChirpPairing.runPairingInitiator/Responder already assume implicitly when they
 * call `transport.send(bytes)` / `transport.receive()` — the WebSocket transport
 * gets framing for free from the protocol; a raw TCP socket does not, so we add
 * it here.
 *
 * A frame length cap guards against a hostile/garbled peer asking us to allocate
 * gigabytes before the Noise layer has authenticated anything.
 */
const val MAX_PAIRING_FRAME_BYTES: Int = 1 shl 20 // 1 MiB

/** Blocking write of a single length-prefixed frame. */
fun writeFrame(out: OutputStream, payload: ByteArray) {
    require(payload.size <= MAX_PAIRING_FRAME_BYTES) { "frame too large: ${payload.size}" }
    val header = ByteArray(4)
    header[0] = (payload.size ushr 24).toByte()
    header[1] = (payload.size ushr 16).toByte()
    header[2] = (payload.size ushr 8).toByte()
    header[3] = payload.size.toByte()
    out.write(header)
    out.write(payload)
    out.flush()
}

/** Blocking read of a single length-prefixed frame. Throws on EOF / oversize. */
fun readFrame(input: InputStream): ByteArray {
    val data = if (input is DataInputStream) input else DataInputStream(input)
    val length = try {
        data.readInt()
    } catch (e: EOFException) {
        throw java.io.IOException("peer closed before frame header", e)
    }
    require(length in 0..MAX_PAIRING_FRAME_BYTES) { "invalid frame length: $length" }
    val buf = ByteArray(length)
    data.readFully(buf)
    return buf
}

/**
 * Adapts a blocking TCP socket pair to the suspend send/receive shape the REAL
 * crypto stack expects:
 *   - noiseNKInitiator / noiseNKResponder take `send: suspend (ByteArray)->Unit`
 *     and `receive: suspend ()->ByteArray` (see NoiseNK.kt).
 *   - ChirpPairing's [PairingTransport] interface is `send/receive/close`.
 *
 * Both the Noise handshake messages and the post-handshake encrypted payloads
 * travel as length-prefixed frames over the same socket, so one [SocketFrameIO]
 * serves the whole session. All IO hops to Dispatchers.IO.
 *
 * This implements the EXISTING [PairingTransport] interface (ChirpPairing.kt),
 * which is the contract `runPairingInitiator` / `runPairingResponder` consume —
 * so a TCP socket can drive the already-written handshake unchanged.
 */
class SocketFrameIO(
    private val input: InputStream,
    private val output: OutputStream,
    private val onClose: suspend () -> Unit = {},
) : PairingTransport {
    private val buffered = DataInputStream(input)

    /**
     * Synchronous frame write. [ChirpSyncSession.send] is non-suspending (see
     * ChirpSyncTransport.kt), so the sync-exchange bridge needs a blocking send
     * that does NOT wrap a coroutine. Socket writes are bounded by soTimeout.
     */
    fun sendBlocking(msg: ByteArray) {
        writeFrame(output, msg)
    }

    override suspend fun send(msg: ByteArray) = withContext(Dispatchers.IO) {
        writeFrame(output, msg)
    }

    override suspend fun receive(): ByteArray = withContext(Dispatchers.IO) {
        readFrame(buffered)
    }

    override suspend fun close() {
        onClose()
    }
}

/**
 * Convenience wrappers for the post-handshake encrypted exchange. The Noise
 * [NoiseNKSession] already owns the cipher (encrypt/decrypt); these just splice
 * it onto the frame layer so call sites read clearly.
 *
 * NOTE: the existing ChirpPairing handshake calls `session.encrypt` /
 * `session.decrypt` directly around `transport.send` / `transport.receive`; the
 * post-pairing SYNC exchange (runChirpSyncInitiator/Responder) does the same via
 * a [ChirpSyncSession]. These helpers exist for the raw-frame paths in
 * PairingServer/PairingTransport that don't go through a ChirpSyncSession.
 */
suspend fun PairingTransport.sendEncrypted(session: NoiseNKSession, plaintext: ByteArray) {
    send(session.encrypt(plaintext))
}

suspend fun PairingTransport.receiveEncrypted(session: NoiseNKSession): ByteArray {
    return session.decrypt(receive())
}