package com.peachpasswords.sync

import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import javax.inject.Inject
import javax.inject.Singleton

/** Whether this device will relay for others. Settings-backed; see [RelaySettingsStore]. */
interface RelayPolicy {
    fun isRelayEnabled(): Boolean
}

/**
 * One authenticated peer connection, as seen by the relay.
 *
 * A link is in exactly one mode at a time, which is what keeps the single
 * Noise session per connection safe to write to: while SYNCING the connection's
 * own coroutine is the only writer, and while RELAYING the partner's read loop
 * is. The mutex covers the overlap at the edges (RELAY_OPEN/RELAY_CLOSE).
 */
class RelayLink internal constructor(
    val peerKeyHex: String,
    private val sendPlaintext: suspend (ByteArray) -> Unit,
) {
    internal enum class Mode { IDLE, SYNCING, RELAYING }

    private val writeLock = Mutex()
    internal var mode: Mode = Mode.IDLE
    internal var partner: RelayLink? = null

    internal suspend fun send(message: RelayMessage) = sendRaw(message.encodeToBytes())

    /** Sends already-encoded plaintext — the forwarding path never re-encodes. */
    internal suspend fun sendRaw(frame: ByteArray) = writeLock.withLock { sendPlaintext(frame) }
}

/**
 * Phone-side relay dispatch (`PEACH_SYNC_RELAY_DESIGN.md` §2 and §4).
 *
 * Two laptops can never connect to each other — neither can listen — so a phone
 * paired with both forwards between them. Frames pass through verbatim: the
 * relay decrypts only the per-hop link encryption it is a party to, sees a
 * RELAY_FRAME wrapper, and re-encrypts the same bytes onto the other link. The
 * inner Noise KK session belongs to the two laptops and terminates there.
 *
 * The hub holds every live authenticated connection keyed by the peer's static
 * public key, which is the addressing layer the design calls for — nothing else
 * on the wire identifies a destination.
 *
 * Multiple laptop pairs can be relayed at once (design §4); what is *not*
 * allowed is one link serving two channels, since a RELAY_FRAME carries no
 * channel id and would be ambiguous.
 */
@Singleton
class RelayHub @Inject constructor(
    private val policy: RelayPolicy,
) {
    private val lock = Mutex()
    private val links = mutableMapOf<String, RelayLink>()

    /** Registers an authenticated connection. Call [unregister] when it ends. */
    suspend fun register(peerKeyHex: String, sendPlaintext: suspend (ByteArray) -> Unit): RelayLink {
        val link = RelayLink(peerKeyHex.lowercase(), sendPlaintext)
        lock.withLock { links[link.peerKeyHex] = link }
        return link
    }

    suspend fun unregister(link: RelayLink) {
        val partner = lock.withLock {
            if (links[link.peerKeyHex] === link) links.remove(link.peerKeyHex)
            detachLocked(link)
        }
        // A relayed peer whose partner vanished would otherwise sit waiting on a
        // channel that can never deliver another frame.
        partner?.let { runCatching { it.send(RelayMessage.Close) } }
    }

    /** Marks a link busy with its own vault's sync, so it is not offered as a relay target. */
    suspend fun beginSync(link: RelayLink): Boolean = lock.withLock {
        if (link.mode != RelayLink.Mode.IDLE) return@withLock false
        link.mode = RelayLink.Mode.SYNCING
        true
    }

    suspend fun endSync(link: RelayLink) = lock.withLock {
        if (link.mode == RelayLink.Mode.SYNCING) link.mode = RelayLink.Mode.IDLE
    }

    /**
     * Answers a RELAY_REQUEST. On success both links enter RELAYING and every
     * later RELAY_FRAME on either is forwarded to the other.
     */
    suspend fun onRequest(from: RelayLink, targetKeyHex: String) {
        if (!policy.isRelayEnabled()) {
            from.send(RelayMessage.Disabled)
            return
        }
        val target = lock.withLock {
            val candidate = links[targetKeyHex.lowercase()]
            if (candidate == null ||
                candidate === from ||
                candidate.mode != RelayLink.Mode.IDLE ||
                from.mode != RelayLink.Mode.IDLE
            ) {
                null
            } else {
                from.partner = candidate
                candidate.partner = from
                from.mode = RelayLink.Mode.RELAYING
                candidate.mode = RelayLink.Mode.RELAYING
                candidate
            }
        }
        if (target == null) {
            from.send(RelayMessage.Unavailable)
            return
        }
        // The target learns first: an ACCEPT that beat the OPEN would invite the
        // requester to send frames into a peer that has not been told to expect
        // them.
        try {
            target.send(RelayMessage.Open(from.peerKeyHex))
        } catch (_: Exception) {
            lock.withLock { detachLocked(from) }
            from.send(RelayMessage.Unavailable)
            return
        }
        from.send(RelayMessage.Accept)
    }

    /**
     * Forwards one frame to the link's partner, byte for byte. Returns false if
     * there is no channel, which the caller answers with RELAY_CLOSE.
     */
    suspend fun forward(from: RelayLink, frame: ByteArray): Boolean {
        val partner = lock.withLock { from.partner } ?: return false
        return try {
            partner.sendRaw(frame)
            true
        } catch (_: Exception) {
            closeChannel(from)
            false
        }
    }

    /** Tears the channel down and tells the other end, if there is one. */
    suspend fun closeChannel(link: RelayLink) {
        val partner = lock.withLock { detachLocked(link) }
        partner?.let { runCatching { it.send(RelayMessage.Close) } }
    }

    /** Must be called with [lock] held. Returns the detached partner, if any. */
    private fun detachLocked(link: RelayLink): RelayLink? {
        val partner = link.partner
        link.partner = null
        if (link.mode == RelayLink.Mode.RELAYING) link.mode = RelayLink.Mode.IDLE
        if (partner != null) {
            partner.partner = null
            if (partner.mode == RelayLink.Mode.RELAYING) partner.mode = RelayLink.Mode.IDLE
        }
        return partner
    }
}