package com.peachpasswords.sync

import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.withContext
import java.io.DataInputStream
import java.io.EOFException
import java.io.IOException
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
 * PairingFrames.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

/**
 * PB-26-205: relay pairing/adoption inbox backstop, mirrored from the
 * browser (`pairing-dial.ts`). The 1 MiB frame ceiling clears the 768 KiB
 * adoption frame ceiling with headroom; the 12 MiB / 256-message inbox
 * clears the 10 MiB aggregate adoption budget plus handshake/control
 * overhead. Both clients agree on accept/reject at these boundaries.
 */
const val MAX_PAIRING_INBOX_BYTES: Int = 12 * 1024 * 1024
const val MAX_PAIRING_INBOX_MESSAGES: Int = 256

/** Terminal pairing error: an IOException so it lands in the ERRORED outcome like a transport failure. */
class PairingResourceLimitException(
    message: String = "pairing_resource_limit",
) : IOException(message)

/**
 * Bounded relay inbox for one pairing/adoption session. The size check runs
 * before [read] performs the allocate/copy, so an oversized frame never
 * materializes. On any violation the channel closes with
 * [PairingResourceLimitException] — every parked and future reader gets that
 * one terminal error — [onLimitExceeded] fires once (the owner closes the
 * socket with status 1009), and later frames are ignored.
 */
internal class RelayPairingInbox(private val onLimitExceeded: () -> Unit) {
    private val channel = Channel<ByteArray>(Channel.UNLIMITED)
    private var queuedBytes = 0
    private var queuedMessages = 0
    private var terminal: Throwable? = null

    @Synchronized
    fun onFrame(size: Int, read: () -> ByteArray): Boolean {
        if (terminal != null) return false
        if (
            size > MAX_PAIRING_FRAME_BYTES ||
            queuedMessages + 1 > MAX_PAIRING_INBOX_MESSAGES ||
            queuedBytes + size > MAX_PAIRING_INBOX_BYTES
        ) {
            tripLimit()
            return false
        }
        val frame = read()
        queuedBytes += frame.size
        queuedMessages += 1
        channel.trySend(frame)
        return true
    }

    /** Trips the limit for a non-frame violation (e.g. a text frame). */
    @Synchronized
    fun reject(): Boolean {
        if (terminal != null) return false
        tripLimit()
        return false
    }

    /** Benign socket close/failure: fails parked readers without a limit close. */
    @Synchronized
    fun close(cause: Throwable) {
        if (terminal != null) return
        terminal = cause
        channel.close(cause)
    }

    private fun tripLimit() {
        terminal = PairingResourceLimitException()
        // A closed channel still delivers buffered elements first, so drop
        // the queued refs for parity with the browser's cleared inbox.
        while (channel.tryReceive().isSuccess) { /* drop */ }
        queuedBytes = 0
        queuedMessages = 0
        channel.close(terminal)
        onLimitExceeded()
    }

    suspend fun receive(): ByteArray {
        val frame = channel.receive()
        synchronized(this) {
            queuedMessages -= 1
            queuedBytes -= frame.size
        }
        return frame
    }
}

/** 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).
 *   - PairingFrames'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 (PairingFrames.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. [SyncSession.send] is non-suspending (see
     * SyncTransport.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 PairingFrames handshake calls `session.encrypt` /
 * `session.decrypt` directly around `transport.send` / `transport.receive`; the
 * post-pairing SYNC exchange (runSyncInitiator/Responder) does the same via
 * a [SyncSession]. These helpers exist for the raw-frame paths in
 * PairingServer/PairingTransport that don't go through a SyncSession.
 */
suspend fun PairingTransport.sendEncrypted(session: NoiseNKSession, plaintext: ByteArray) {
    send(session.encrypt(plaintext))
}

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