package dev.kaizen.app.live

import android.util.Log
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.SharedFlow
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.atomic.AtomicReference

/**
 * WebSocket client for the Kaizen Live Call protocol.
 *
 * Connects to `wss://{baseUrl}/ws/live` with Bearer auth, speaks the
 * provider-agnostic JSON protocol defined in `lib/live/protocol.ts`.
 * Handles auto-reconnect with exponential backoff + resume.
 *
 * Thread-safe: OkHttp dispatches callbacks on its own threads,
 * all state mutations go through atomic refs or the coroutine scope.
 */
class LiveClient(
    private val baseUrl: String,
    private val token: String,
    private val scope: CoroutineScope,
) {
    companion object {
        private const val TAG = "LiveClient"
        private const val INITIAL_BACKOFF_MS = 500L
        private const val MAX_BACKOFF_MS = 8_000L
        private const val MAX_RECONNECT_ATTEMPTS = 5
    }

    // ─── Public event stream ──────────────────────────────────────────────

    private val _events = MutableSharedFlow<ServerEvent>(extraBufferCapacity = 64)
    val events: SharedFlow<ServerEvent> = _events

    // ─── State ────────────────────────────────────────────────────────────

    private val ws = AtomicReference<WebSocket?>(null)
    private val sessionId = AtomicReference<String?>(null)
    private val lastSeqId = AtomicInteger(0)
    private val reconnectAttempts = AtomicInteger(0)
    private var chunkIntervalMs = 100L

    @Volatile private var active = false
    @Volatile private var pendingSetup: SetupParams? = null

    private data class SetupParams(
        val model: String,
        val voice: String,
        val conversationId: String?,
    )

    private val client = OkHttpClient.Builder()
        .readTimeout(0, TimeUnit.MILLISECONDS)
        .pingInterval(15, TimeUnit.SECONDS)
        .connectTimeout(10, TimeUnit.SECONDS)
        .build()

    // ─── Connect ──────────────────────────────────────────────────────────

    fun connect(model: String, voice: String, conversationId: String?) {
        active = true
        pendingSetup = SetupParams(model, voice, conversationId)
        reconnectAttempts.set(0)
        openSocket()
    }

    private fun openSocket() {
        if (!active) return

        val wsUrl = baseUrl
            .replace("https://", "wss://")
            .replace("http://", "ws://") + "/ws/live"

        val request = Request.Builder()
            .url(wsUrl)
            .header("Authorization", "Bearer $token")
            .build()

        val socket = client.newWebSocket(request, Listener())
        ws.set(socket)
    }

    // ─── Send ─────────────────────────────────────────────────────────────

    fun sendAudio(base64Pcm: String) {
        ws.get()?.send(encodeAudio(base64Pcm))
    }

    fun sendText(text: String) {
        ws.get()?.send(encodeText(text))
    }

    fun interrupt() {
        ws.get()?.send(encodeInterrupt())
    }

    fun end() {
        active = false
        ws.get()?.send(encodeEnd())
        ws.getAndSet(null)?.close(1000, "end")
    }

    fun dispose() {
        active = false
        ws.getAndSet(null)?.cancel()
        client.dispatcher.executorService.shutdown()
    }

    // ─── Reconnect ────────────────────────────────────────────────────────

    private fun scheduleReconnect() {
        if (!active) return
        val attempt = reconnectAttempts.incrementAndGet()
        if (attempt > MAX_RECONNECT_ATTEMPTS) {
            Log.w(TAG, "Max reconnect attempts reached")
            scope.launch { _events.emit(ServerEvent.Error("reconnect_failed", "Max reconnect attempts")) }
            active = false
            return
        }

        val backoff = (INITIAL_BACKOFF_MS * (1 shl (attempt - 1))).coerceAtMost(MAX_BACKOFF_MS)
        Log.d(TAG, "Reconnecting in ${backoff}ms (attempt $attempt)")

        scope.launch {
            delay(backoff)
            if (active) openSocket()
        }
    }

    // ─── WebSocket Listener ───────────────────────────────────────────────

    private inner class Listener : WebSocketListener() {

        override fun onOpen(webSocket: WebSocket, response: Response) {
            Log.d(TAG, "WebSocket opened")
            reconnectAttempts.set(0)

            val sid = sessionId.get()
            if (sid != null) {
                webSocket.send(encodeResume(sid, lastSeqId.get()))
            } else {
                val setup = pendingSetup ?: return
                webSocket.send(encodeSetup(setup.model, setup.voice, setup.conversationId))
            }
        }

        override fun onMessage(webSocket: WebSocket, text: String) {
            val event = parseServerEvent(text) ?: return

            when (event) {
                is ServerEvent.Ready -> {
                    sessionId.set(event.sessionId)
                }
                is ServerEvent.Resumed -> {
                    sessionId.set(event.sessionId)
                }
                is ServerEvent.Audio -> {
                    lastSeqId.set(maxOf(lastSeqId.get(), event.seq))
                }
                is ServerEvent.Text -> {
                    lastSeqId.set(maxOf(lastSeqId.get(), event.seq))
                }
                is ServerEvent.Backpressure -> {
                    chunkIntervalMs = if (event.throttleMs > 0) event.throttleMs.toLong() else 100L
                }
                is ServerEvent.GoAway -> {
                    if (event.reason == "shutdown") {
                        scope.launch {
                            delay(500)
                            if (active) openSocket()
                        }
                    }
                }
                is ServerEvent.Error -> {
                    val terminal = event.code in setOf(
                        "auth_failed", "restricted_model", "restricted_feature",
                        "budget_exceeded", "concurrent_limit",
                    )
                    if (terminal) active = false
                }
                else -> {}
            }

            scope.launch { _events.emit(event) }
        }

        override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
            webSocket.close(code, reason)
        }

        override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
            Log.d(TAG, "WebSocket closed: $code $reason")
            ws.compareAndSet(webSocket, null)
            if (active && code != 1000) scheduleReconnect()
        }

        override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
            Log.w(TAG, "WebSocket failure: ${t.message}")
            ws.compareAndSet(webSocket, null)
            if (active) scheduleReconnect()
        }
    }
}
