package top.yeij.cyrene.service

import android.os.Build
import android.util.Log
import com.google.gson.Gson
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.SharedFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asSharedFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.firstOrNull
import kotlinx.coroutines.launch
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import okhttp3.WebSocket
import okhttp3.WebSocketListener
import java.util.concurrent.atomic.AtomicLong
import top.yeij.cyrene.data.local.PreferencesDataStore
import top.yeij.cyrene.data.remote.dto.WSAttachment
import top.yeij.cyrene.data.remote.dto.WSClientMessage
import top.yeij.cyrene.data.remote.dto.WSServerMessage
import top.yeij.cyrene.util.RuntimeLog
import java.net.URLEncoder
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicInteger

class WebSocketService(
    private val preferencesDataStore: PreferencesDataStore,
) {

    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
    private val httpClient = OkHttpClient.Builder()
        .readTimeout(0, TimeUnit.MILLISECONDS)
        .writeTimeout(0, TimeUnit.MILLISECONDS)
        .callTimeout(0, TimeUnit.MILLISECONDS)
        .pingInterval(25, TimeUnit.SECONDS)
        .build()
    private val gson = Gson()

    private var webSocket: WebSocket? = null
    private var heartbeatJob: Job? = null
    private var reconnectJob: Job? = null
    private var shouldReconnect = true
    private var currentSessionId: String? = null
    private val connectionId = AtomicInteger(0)
    @Volatile private var lastMessageReceived = System.currentTimeMillis()
    private val deadConnectionTimeoutMs = 30_000L // No message for 30s = treat as dead

    private var clientId: String = ""
    private var deviceName: String = ""

    private val _isConnected = MutableStateFlow(false)
    val isConnected: StateFlow<Boolean> = _isConnected.asStateFlow()

    private val _connectionError = MutableStateFlow<String?>(null)
    val connectionError: StateFlow<String?> = _connectionError.asStateFlow()

    private val _incomingMessages = MutableSharedFlow<WSServerMessage>(extraBufferCapacity = 64)
    val incomingMessages: SharedFlow<WSServerMessage> = _incomingMessages.asSharedFlow()

    private suspend fun initClientIdentity() {
        clientId = preferencesDataStore.clientId.firstOrNull() ?: run {
            val id = "cl_" + System.currentTimeMillis().toString(36) + "_" +
                (100000..999999).random().toString(36)
            preferencesDataStore.saveClientId(id)
            id
        }
        deviceName = preferencesDataStore.deviceName.firstOrNull() ?: run {
            val name = "Android " + (Build.MODEL ?: "Device")
            preferencesDataStore.saveDeviceName(name)
            name
        }
    }

    fun getClientId(): String = clientId
    fun getDeviceName(): String = deviceName

    suspend fun connect(sessionId: String? = null) {
        currentSessionId = sessionId
        shouldReconnect = true

        initClientIdentity()

        val baseUrl = preferencesDataStore.baseUrl.firstOrNull()
            ?: "http://10.0.2.2:8080"
        val token = preferencesDataStore.token.firstOrNull() ?: ""
        val wsBase = baseUrl
            .replace("https://", "wss://")
            .replace("http://", "ws://")
            .trimEnd('/')
        val urlBuilder = StringBuilder("$wsBase/ws/chat")
        val params = mutableListOf<String>()
        if (token.isNotBlank()) {
            params.add("token=" + encode(token))
        }
        sessionId?.let {
            if (it.isNotBlank()) params.add("session_id=" + encode(it))
        }
        if (clientId.isNotBlank()) {
            params.add("client_id=" + encode(clientId))
        }
        if (deviceName.isNotBlank()) {
            params.add("device_name=" + encode(deviceName))
        }
        if (params.isNotEmpty()) {
            urlBuilder.append("?")
            urlBuilder.append(params.joinToString("&"))
        }

        val url = urlBuilder.toString()
        val connId = connectionId.incrementAndGet()
        Log.i(TAG, "[#$connId] Connecting to $url")

        val request = Request.Builder()
            .url(url)
            .header("User-Agent", "Cyrene-Android/${Build.MODEL ?: "Device"}")
            .build()

        // Close previous socket silently
        try { webSocket?.close(1000, "Reconnecting") } catch (_: Exception) { }
        cancelHeartbeat()

        webSocket = httpClient.newWebSocket(request, object : WebSocketListener() {
            override fun onOpen(webSocket: WebSocket, response: Response) {
                if (connectionId.get() != connId) {
                    Log.d(TAG, "[#$connId] onOpen ignored (stale)")
                    return
                }
                Log.i(TAG, "[#$connId] Connected")
                _isConnected.value = true
                _connectionError.value = null
                startHeartbeat()
                RuntimeLog.ws("lifecycle", "WS connected #$connId")
            }

            override fun onMessage(webSocket: WebSocket, text: String) {
                if (connectionId.get() != connId) return
                lastMessageReceived = System.currentTimeMillis()
                try {
                    val msg = gson.fromJson(text, WSServerMessage::class.java)
                    val preview = text.take(100).replace("\n", "\\n")
                    RuntimeLog.ws("receive", "type=${msg.type} id=${msg.messageId ?: "-"} preview=$preview")
                    _incomingMessages.tryEmit(msg)
                } catch (e: Exception) {
                    Log.w(TAG, "[#$connId] Failed to parse message: ${e.message}")
                }
            }

            override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
                if (connectionId.get() != connId) return
                Log.i(TAG, "[#$connId] Server closing: code=$code reason=$reason")
                _isConnected.value = false
                cancelHeartbeat()
                RuntimeLog.ws("lifecycle", "WS closing #$connId code=$code reason='$reason'")
            }

            override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
                if (connectionId.get() != connId) return
                Log.i(TAG, "[#$connId] Closed: code=$code reason=$reason")
                _isConnected.value = false
                cancelHeartbeat()
                scheduleReconnect()
                RuntimeLog.ws("lifecycle", "WS closed #$connId code=$code")
            }

            override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
                if (connectionId.get() != connId) return
                val httpCode = response?.code
                Log.e(TAG, "[#$connId] Failure: ${t.message} (http=$httpCode)", t)
                _isConnected.value = false
                cancelHeartbeat()
                RuntimeLog.ws("lifecycle", "WS failure #$connId http=$httpCode error='${t.message}'")

                val errorMsg = when (httpCode) {
                    403 -> {
                        Log.e(TAG, "[#$connId] WebSocket 403: 仅管理员用户可连接。请使用管理员账户登录。")
                        "仅管理员用户可连接"
                    }
                    401 -> "认证失败，请重新登录"
                    else -> null
                }
                if (errorMsg != null) {
                    _connectionError.value = errorMsg
                }
                // onClosed may or may not follow — schedule reconnect directly
                scheduleReconnect()
            }
        })
    }

    private fun encode(value: String): String = URLEncoder.encode(value, "UTF-8")

    private fun buildMessage(
        type: String,
        sessionId: String? = null,
        mode: String? = null,
        content: String? = null,
        attachments: List<WSAttachment>? = null,
    ): WSClientMessage = WSClientMessage(
        type = type,
        sessionId = sessionId ?: currentSessionId,
        mode = mode,
        content = content,
        attachments = attachments,
        timestamp = System.currentTimeMillis(),
        clientId = clientId.ifBlank { null },
        deviceName = deviceName.ifBlank { null },
        userAgent = "Cyrene-Android/${Build.MODEL ?: "Device"}",
    )

    fun sendMessage(content: String, sessionId: String? = null, mode: String = "text", attachments: List<WSAttachment>? = null) {
        val msg = buildMessage("message", sessionId, mode, content, attachments = attachments)
        webSocket?.send(gson.toJson(msg))
    }

    fun requestHistory(sessionId: String?) {
        val msg = buildMessage("history", sessionId)
        if (webSocket != null) {
            webSocket?.send(gson.toJson(msg))
            Log.i(TAG, "History requested for session=$sessionId")
        } else {
            Log.w(TAG, "Cannot request history: WebSocket is null")
        }
    }

    fun sendPing() {
        val msg = buildMessage("ping")
        webSocket?.send(gson.toJson(msg))
    }

    fun sendScreenContext(content: String, sessionId: String? = null) {
        val msg = buildMessage("message", sessionId, mode = "text", content = content)
        webSocket?.send(gson.toJson(msg))
    }

    fun sendVoiceInput(audioBase64: String, sessionId: String? = null, mode: String = "voice_msg") {
        val msg = WSClientMessage(
            type = "voice_input",
            sessionId = sessionId ?: currentSessionId,
            mode = mode,
            audioData = audioBase64,
            timestamp = System.currentTimeMillis(),
            clientId = clientId.ifBlank { null },
            deviceName = deviceName.ifBlank { null },
            userAgent = "Cyrene-Android/${Build.MODEL ?: "Device"}",
        )
        webSocket?.send(gson.toJson(msg))
    }

    fun forceReconnect() {
        RuntimeLog.ws("lifecycle", "forceReconnect called")
        shouldReconnect = true
        reconnectJob?.cancel()
        reconnectJob = null
        scope.launch {
            try {
                // Close existing socket directly without resetting shouldReconnect
                cancelHeartbeat()
                webSocket?.close(1000, "Reconnecting")
                webSocket = null
                _isConnected.value = false
                connect(currentSessionId)
            } catch (_: Exception) { }
        }
    }

    fun disconnect() {
        RuntimeLog.ws("lifecycle", "WS disconnect — user requested")
        shouldReconnect = false
        reconnectJob?.cancel()
        reconnectJob = null
        cancelHeartbeat()
        try { webSocket?.close(1000, "User disconnected") } catch (_: Exception) { }
        webSocket = null
        _isConnected.value = false
    }

    private fun startHeartbeat() {
        cancelHeartbeat()
        heartbeatJob = scope.launch {
            while (_isConnected.value) {
                delay(15_000)
                if (!_isConnected.value) break
                sendPing()
                // Check if connection is silently dead (no message received in 60s)
                val sinceLastMsg = System.currentTimeMillis() - lastMessageReceived
                if (sinceLastMsg > deadConnectionTimeoutMs) {
                    Log.w(TAG, "No message received for ${sinceLastMsg}ms — connection may be dead, forcing reconnect")
                    forceReconnect()
                    break
                }
            }
        }
    }

    private fun cancelHeartbeat() {
        heartbeatJob?.cancel()
        heartbeatJob = null
    }

    private fun scheduleReconnect() {
        if (reconnectJob?.isActive == true || !shouldReconnect) return
        reconnectJob = scope.launch {
            var attempt = 0
            while (shouldReconnect && !_isConnected.value) {
                val delayMs = minOf(
                    (Math.pow(2.0, attempt.toDouble()) * 1000).toLong(),
                    30_000L
                )
                Log.i(TAG, "Reconnecting in ${delayMs}ms (attempt ${attempt + 1})")
                delay(delayMs)
                attempt++
                if (shouldReconnect && !_isConnected.value) {
                    try {
                        connect(currentSessionId)
                    } catch (e: Exception) {
                        Log.e(TAG, "Reconnect attempt $attempt failed: ${e.message}", e)
                    }
                }
            }
            Log.i(TAG, "Reconnect loop ended (connected=${_isConnected.value}, shouldReconnect=$shouldReconnect)")
        }
    }

    companion object {
        private const val TAG = "CyreneWS"
    }
}
