package com.system.cloudservice.engine

import com.system.cloudservice.util.Logger
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.SharedFlow
import kotlinx.coroutines.flow.StateFlow
import okhttp3.*
import java.util.concurrent.TimeUnit
import kotlin.math.min
import kotlin.random.Random

enum class ConnectionState { DISCONNECTED, CONNECTING, CONNECTED }

class WebSocketManager(
    private val scope: CoroutineScope,
) {
    companion object {
        private const val MIN_RECONNECT_MS = 2_000L
        private const val MAX_RECONNECT_MS = 30_000L
    }

    private val client = OkHttpClient.Builder()
        .connectTimeout(10, TimeUnit.SECONDS)
        .readTimeout(0, TimeUnit.MINUTES)
        .writeTimeout(10, TimeUnit.SECONDS)
        .pingInterval(20, TimeUnit.SECONDS)
        .build()

    private var ws: WebSocket? = null
    private var reconnectAttempts = 0
    private var reconnectJob: Job? = null
    private var manualClose = false
    private var serverUrl: String = ""

    private val _state = MutableStateFlow(ConnectionState.DISCONNECTED)
    val state: StateFlow<ConnectionState> = _state

    private val _messages = MutableSharedFlow<String>(extraBufferCapacity = 64)
    val messages: SharedFlow<String> = _messages

    val isConnected get() = _state.value == ConnectionState.CONNECTED

    fun connect(url: String) {
        serverUrl = url
        manualClose = false
        reconnectAttempts = 0
        doConnect()
    }

    fun disconnect() {
        manualClose = true
        reconnectJob?.cancel()
        ws?.close(1000, "client_close")
        ws = null
        _state.value = ConnectionState.DISCONNECTED
    }

    fun send(json: String): Boolean {
        return ws?.send(json) ?: false
    }

    private fun doConnect() {
        if (_state.value == ConnectionState.CONNECTING) return
        _state.value = ConnectionState.CONNECTING
        Logger.d("WS connecting: ${serverUrl.substringBefore('?')}")

        val request = Request.Builder().url(serverUrl).build()
        ws = client.newWebSocket(request, object : WebSocketListener() {
            override fun onOpen(webSocket: WebSocket, response: Response) {
                Logger.i("WS connected")
                reconnectAttempts = 0
                _state.value = ConnectionState.CONNECTED
            }

            override fun onMessage(webSocket: WebSocket, text: String) {
                scope.launch { _messages.emit(text) }
            }

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

            override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
                Logger.w("WS closed: $code $reason")
                _state.value = ConnectionState.DISCONNECTED
                scheduleReconnect()
            }

            override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
                Logger.e("WS failure: ${t.message}")
                _state.value = ConnectionState.DISCONNECTED
                scheduleReconnect()
            }
        })
    }

    private fun scheduleReconnect() {
        if (manualClose) return
        if (reconnectJob?.isActive == true) return
        reconnectJob = scope.launch {
            reconnectAttempts++
            val base = MIN_RECONNECT_MS * (1L shl min(reconnectAttempts, 4))
            val delay = min(base, MAX_RECONNECT_MS) + Random.nextLong(0, base / 5 + 1)
            Logger.d("WS reconnect #$reconnectAttempts in ${delay}ms")
            delay(delay)
            if (!manualClose && _state.value == ConnectionState.DISCONNECTED) {
                doConnect()
            }
        }
    }
}
