宜博 / phoneuse

package ai.lemon.phoneuse.eventbus

import ai.lemon.phoneuse.core.event.Event
import ai.lemon.phoneuse.core.event.EventSerializer
import android.net.LocalServerSocket
import android.net.LocalSocket
import android.util.Log
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancel
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import java.io.IOException
import java.io.PrintWriter
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicLong

/**
 * Unix Domain Socket server that streams [Event] objects to external subscribers.
 *
 * Each event is serialized as a single-line JSON string followed by a newline
 * character (newline-delimited JSON / NDJSON format), making it trivial to
 * consume from Python or any other language with a socket reader.
 *
 * ### Android root note
 * The abstract socket namespace (`\0agent_events`) is used by default and
 * requires no filesystem permissions.  If a concrete path is preferred
 * (`/data/local/tmp/agent_events.sock`), set [useAbstractNamespace] to
 * `false`; the caller is then responsible for ensuring the path is writable
 * (e.g. via `su chmod 666 <path>`).
 *
 * ### Usage
 * ```kotlin
 * val server = EventStreamServer()
 * server.start()
 *
 * // later, from EventPublisher:
 * server.broadcast(event)
 *
 * // on teardown:
 * server.stop()
 * ```
 */
class EventStreamServer(
    /**
     * Abstract socket name (no leading NUL — the Android API adds it automatically).
     * Clients connect using the same name.
     */
    private val socketName: String = SOCKET_NAME,
) {

    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)

    /** Tracks active client writers keyed by a monotonic connection ID. */
    private val clients = ConcurrentHashMap<Long, ClientWriter>()
    private val nextId = AtomicLong(0)

    private var serverSocket: LocalServerSocket? = null
    private var acceptJob: Job? = null

    // ------------------------------------------------------------------ //
    // Public API                                                           //
    // ------------------------------------------------------------------ //

    /** Starts the server; idempotent if already running. */
    fun start() {
        if (acceptJob?.isActive == true) return
        try {
            serverSocket = LocalServerSocket(socketName)
            Log.i(TAG, "Listening on abstract socket @$socketName")
        } catch (e: IOException) {
            Log.e(TAG, "Failed to bind socket: ${e.message}", e)
            return
        }
        acceptJob = scope.launch { acceptLoop() }
    }

    /** Stops the server and disconnects all clients. */
    fun stop() {
        acceptJob?.cancel()
        acceptJob = null
        clients.values.forEach { it.close() }
        clients.clear()
        try {
            serverSocket?.close()
        } catch (_: IOException) {}
        serverSocket = null
        scope.cancel()
        Log.i(TAG, "Server stopped")
    }

    /**
     * Serializes [event] to NDJSON and enqueues it for delivery to every
     * currently connected client.  Non-blocking: slow or disconnected clients
     * are dropped rather than blocking the publisher.
     */
    fun broadcast(event: Event) {
        val line = EventSerializer.toJson(event)
        val deadClients = mutableListOf<Long>()
        clients.forEach { (id, writer) ->
            val sent = writer.offer(line)
            if (!sent) {
                Log.w(TAG, "Client $id dropped (channel full or closed)")
                deadClients += id
            }
        }
        deadClients.forEach { id ->
            clients.remove(id)?.close()
        }
    }

    // ------------------------------------------------------------------ //
    // Internal                                                             //
    // ------------------------------------------------------------------ //

    private suspend fun acceptLoop() {
        val srv = serverSocket ?: return
        while (scope.isActive) {
            val socket = try {
                withContext(Dispatchers.IO) { srv.accept() }
            } catch (e: IOException) {
                if (scope.isActive) Log.e(TAG, "Accept error: ${e.message}", e)
                break
            }
            val id = nextId.getAndIncrement()
            Log.i(TAG, "Client $id connected")
            val writer = ClientWriter(id, socket)
            clients[id] = writer
            writer.startIn(scope) { clients.remove(id) }
        }
    }

    // ------------------------------------------------------------------ //
    // ClientWriter                                                         //
    // ------------------------------------------------------------------ //

    /**
     * Owns a single client [LocalSocket] and a dedicated coroutine that
     * drains a [Channel] into the socket's output stream.
     *
     * Using a channel per client ensures that a slow consumer never blocks
     * [broadcast] — once the channel is full the client is evicted.
     */
    private class ClientWriter(
        private val id: Long,
        private val socket: LocalSocket,
    ) {
        private val channel = Channel<String>(capacity = CLIENT_BUFFER_LINES)
        private var writeJob: Job? = null

        fun startIn(scope: CoroutineScope, onDisconnect: () -> Unit) {
            writeJob = scope.launch(Dispatchers.IO) {
                try {
                    val writer = PrintWriter(socket.outputStream.bufferedWriter())
                    for (line in channel) {
                        writer.println(line)
                        writer.flush()
                        if (writer.checkError()) {
                            Log.i(TAG, "Client $id write error — disconnecting")
                            break
                        }
                    }
                } catch (e: IOException) {
                    Log.i(TAG, "Client $id IO error: ${e.message}")
                } finally {
                    onDisconnect()
                    close()
                }
            }
        }

        /**
         * Enqueues a line for delivery.
         * @return false if the channel is full or already closed.
         */
        fun offer(line: String): Boolean = channel.trySend(line).isSuccess

        fun close() {
            channel.close()
            writeJob?.cancel()
            try { socket.close() } catch (_: IOException) {}
        }
    }

    companion object {
        private const val TAG = "EventStreamServer"

        /** Abstract socket name used by both server and client. */
        const val SOCKET_NAME = "agent_events"

        /** Per-client in-memory buffer (lines).  Clients lagging beyond this are dropped. */
        private const val CLIENT_BUFFER_LINES = 512
    }
}