package com.anypreta.bountyboard.data.ws

import com.anypreta.bountyboard.data.model.Message
import com.anypreta.bountyboard.data.model.Notification
import com.google.gson.Gson
import com.google.gson.JsonObject
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.SharedFlow
import kotlinx.coroutines.flow.asSharedFlow
import kotlinx.coroutines.isActive
import kotlinx.coroutines.launch
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.AtomicBoolean

/**
 * The single multiplexed live-update WebSocket (server §2.2): one connection per
 * session carrying `chat`, `board` and `notifications` channels. Auth rides on the
 * shared cookie jar; the server allows clients that send no Origin header (we don't),
 * so no server changes are required. Auto-reconnects with backoff while [connect]ed.
 */
class LiveSocket(
    baseClient: OkHttpClient,
    private val gson: Gson,
    private val baseUrlProvider: () -> String,
) {
    private val client = baseClient.newBuilder()
        .pingInterval(20, TimeUnit.SECONDS) // keep the socket alive through proxies
        .readTimeout(0, TimeUnit.SECONDS)
        .build()

    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
    private val running = AtomicBoolean(false)
    @Volatile private var socket: WebSocket? = null
    @Volatile private var reconnectJob: kotlinx.coroutines.Job? = null

    private val _chat = MutableSharedFlow<Message>(extraBufferCapacity = 64)
    val chatMessages: SharedFlow<Message> = _chat.asSharedFlow()

    private val _notifications = MutableSharedFlow<Notification>(extraBufferCapacity = 64)
    val notifications: SharedFlow<Notification> = _notifications.asSharedFlow()

    private val _board = MutableSharedFlow<BoardEvent>(extraBufferCapacity = 64)
    val boardEvents: SharedFlow<BoardEvent> = _board.asSharedFlow()

    private val _connected = MutableSharedFlow<Boolean>(replay = 1, extraBufferCapacity = 1)
    val connected: SharedFlow<Boolean> = _connected.asSharedFlow()

    fun connect() {
        if (!running.compareAndSet(false, true)) return
        open()
    }

    fun disconnect() {
        running.set(false)
        reconnectJob?.cancel()
        socket?.close(1000, "bye")
        socket = null
    }

    private fun open() {
        if (!running.get()) return
        val url = baseUrlProvider().trimEnd('/') + "/ws"
        val request = Request.Builder().url(url).build()
        socket = client.newWebSocket(request, listener)
    }

    private fun scheduleReconnect() {
        if (!running.get()) return
        if (reconnectJob?.isActive == true) return
        reconnectJob = scope.launch {
            var attempt = 0
            while (isActive && running.get() && socket == null) {
                val backoff = minOf(30_000L, 1_000L * (1L shl minOf(attempt, 4)))
                delay(backoff)
                attempt++
                open()
                // give the open a moment; if it failed, onFailure nulls socket again
                delay(2_000)
            }
        }
    }

    private val listener = object : WebSocketListener() {
        override fun onOpen(webSocket: WebSocket, response: Response) {
            _connected.tryEmit(true)
        }

        override fun onMessage(webSocket: WebSocket, text: String) {
            try {
                val env = gson.fromJson(text, JsonObject::class.java) ?: return
                val channel = env.get("channel")?.asString ?: return
                val event = env.get("event")?.asString ?: return
                val data = env.get("data")
                when (channel) {
                    "chat" -> if (event == "message" && data != null) {
                        runCatching { gson.fromJson(data, Message::class.java) }
                            .getOrNull()?.let { _chat.tryEmit(it) }
                    }
                    "notifications" -> if (event == "notification" && data != null) {
                        runCatching { gson.fromJson(data, Notification::class.java) }
                            .getOrNull()?.let { _notifications.tryEmit(it) }
                    }
                    "board" -> {
                        val taskId = data?.asJsonObject?.get("taskId")?.asString
                        _board.tryEmit(BoardEvent(event, taskId))
                    }
                }
            } catch (_: Exception) {
                // ignore malformed frames
            }
        }

        override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
            socket = null
            _connected.tryEmit(false)
            scheduleReconnect()
        }

        override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
            socket = null
            _connected.tryEmit(false)
            scheduleReconnect()
        }
    }
}

data class BoardEvent(val event: String, val taskId: String?)
