package com.openclaw.alfred.gateway

import android.content.Context
import android.net.ConnectivityManager
import android.net.NetworkCapabilities
import android.util.Log
import com.google.gson.Gson
import com.google.gson.JsonObject
import com.openclaw.alfred.BuildConfig
import okhttp3.*
import java.util.concurrent.TimeUnit

/**
 * WebSocket client for OpenClaw Gateway connection.
 * Handles protocol handshake, message framing, and reconnection.
 */
class GatewayClient(
    private val context: Context,
    private val accessToken: String,
    private val listener: GatewayListener
) {
    
    private val TAG = "GatewayClient"
    private val gson = Gson()
    
    // Get gateway URL from preferences, fallback to BuildConfig
    private fun getGatewayUrl(): String {
        val prefs = context.getSharedPreferences("alfred_settings", Context.MODE_PRIVATE)
        return prefs.getString("gateway_url", BuildConfig.GATEWAY_URL) ?: BuildConfig.GATEWAY_URL
    }
    
    private var webSocket: WebSocket? = null
    private var isConnected = false
    private var requestId = 0
    
    // Extract user ID from JWT for session key
    private val userId: String by lazy {
        try {
            // JWT format: header.payload.signature
            val parts = accessToken.split(".")
            if (parts.size >= 2) {
                val payload = String(android.util.Base64.decode(parts[1], android.util.Base64.URL_SAFE or android.util.Base64.NO_WRAP))
                val json = gson.fromJson(payload, JsonObject::class.java)
                
                // Prefer username over email over sub for consistent, readable session keys
                val username = json.get("preferred_username")?.asString
                val email = json.get("email")?.asString
                val sub = json.get("sub")?.asString
                
                when {
                    !username.isNullOrEmpty() -> username
                    !email.isNullOrEmpty() -> email
                    !sub.isNullOrEmpty() -> sub
                    else -> "mobile"
                }
            } else {
                "mobile"
            }
        } catch (e: Exception) {
            Log.e(TAG, "Failed to extract user ID from token", e)
            "mobile"
        }
    }
    
    // Reconnection state
    private var shouldReconnect = true
    private var reconnectAttempts = 0
    private val maxReconnectAttempts = 10
    private val baseReconnectDelayMs = 1000L // Start with 1 second
    private val maxReconnectDelayMs = 30000L // Max 30 seconds
    private var reconnectHandler: android.os.Handler? = null
    
    private val client = OkHttpClient.Builder()
        .connectTimeout(10, TimeUnit.SECONDS)
        .readTimeout(0, TimeUnit.MILLISECONDS)  // No timeout for WebSocket
        .build()
    
    /**
     * Connect to the OpenClaw Gateway.
     */
    fun connect() {
        if (isConnected) {
            Log.w(TAG, "Already connected")
            return
        }
        
        // Close any existing WebSocket before creating a new one
        webSocket?.let { existingWs ->
            Log.d(TAG, "Closing existing WebSocket before reconnect")
            try {
                existingWs.close(1000, "Reconnecting")
            } catch (e: Exception) {
                Log.e(TAG, "Error closing existing WebSocket: ${e.message}")
            }
            webSocket = null
        }
        
        // Enable reconnection and reset state
        shouldReconnect = true
        
        val gatewayUrl = getGatewayUrl()
        Log.d(TAG, "Connecting to $gatewayUrl")
        
        val request = Request.Builder()
            .url(gatewayUrl)
            .addHeader("Authorization", "Bearer $accessToken")
            .build()
        
        webSocket = client.newWebSocket(request, object : WebSocketListener() {
            override fun onOpen(webSocket: WebSocket, response: Response) {
                Log.d(TAG, "WebSocket opened")
                listener.onConnecting()
            }
            
            override fun onMessage(webSocket: WebSocket, text: String) {
                Log.d(TAG, "<<< Received TEXT: $text")
                handleMessage(text)
            }
            
            override fun onMessage(webSocket: WebSocket, bytes: okio.ByteString) {
                val text = bytes.utf8()
                Log.d(TAG, "<<< Received BINARY (converted): $text")
                handleMessage(text)
            }
            
            override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
                Log.d(TAG, "WebSocket closing: code=$code reason='$reason'")
                isConnected = false
                listener.onDisconnected()
                webSocket.close(1000, null)
                
                // Attempt reconnection unless explicitly disconnected
                if (shouldReconnect) {
                    scheduleReconnect()
                }
            }
            
            override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
                Log.d(TAG, "WebSocket closed: code=$code reason='$reason'")
                isConnected = false
                
                // Check for authentication failures (code 1008 = Policy Violation)
                if (code == 1008 || reason.contains("Authentication", ignoreCase = true) || 
                    reason.contains("Invalid token", ignoreCase = true)) {
                    Log.w(TAG, "Connection closed due to authentication failure")
                    listener.onError("Authentication failed: $reason")
                    // Don't auto-reconnect on auth failures - let app handle token refresh
                    shouldReconnect = false
                    return
                }
                
                // Attempt reconnection unless explicitly disconnected
                if (shouldReconnect) {
                    scheduleReconnect()
                }
            }
            
            override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
                Log.e(TAG, "WebSocket failure: ${t.message}, response code: ${response?.code}")
                isConnected = false
                
                // Check for authentication failures (401 Unauthorized)
                if (response?.code == 401 || response?.code == 403) {
                    Log.w(TAG, "Connection failed due to authentication (${response.code})")
                    listener.onError("Authentication failed (401 Unauthorized)")
                    // Don't auto-reconnect on auth failures - let app handle token refresh
                    shouldReconnect = false
                    return
                }
                
                listener.onError(t.message ?: "Connection failed")
                
                // Attempt reconnection unless explicitly disconnected
                if (shouldReconnect) {
                    scheduleReconnect()
                }
            }
        })
    }
    
    /**
     * Schedule a reconnection attempt with exponential backoff.
     */
    private fun scheduleReconnect() {
        // Check if we've exceeded max attempts
        if (reconnectAttempts >= maxReconnectAttempts) {
            Log.e(TAG, "Max reconnection attempts ($maxReconnectAttempts) reached. Giving up.")
            listener.onError("Connection lost - max retries exceeded")
            shouldReconnect = false
            return
        }
        
        // Check network availability
        if (!isNetworkAvailable()) {
            Log.w(TAG, "Network unavailable, waiting longer before retry...")
            // Use longer delay when network is unavailable (10 seconds)
            // Don't increment reconnectAttempts - we're not actually trying to connect
            val delay = 10000L
            
            Log.d(TAG, "Network unavailable - will check again in ${delay}ms (not counting as retry attempt)")
            listener.onReconnecting(reconnectAttempts, delay)
            
            // Cancel any pending reconnection
            reconnectHandler?.removeCallbacksAndMessages(null)
            
            // Schedule reconnection
            reconnectHandler = android.os.Handler(android.os.Looper.getMainLooper())
            reconnectHandler?.postDelayed({
                if (shouldReconnect && !isConnected) {
                    // Check network again before attempting
                    if (isNetworkAvailable()) {
                        Log.d(TAG, "Network restored, attempting reconnection")
                        connect()
                    } else {
                        Log.d(TAG, "Network still unavailable, rescheduling...")
                        scheduleReconnect()
                    }
                }
            }, delay)
            return
        }
        
        // Calculate delay with exponential backoff: 1s, 2s, 4s, 8s, 16s, 30s (max)
        val delay = minOf(
            baseReconnectDelayMs * (1 shl reconnectAttempts), // 2^n backoff
            maxReconnectDelayMs
        )
        
        reconnectAttempts++
        
        Log.d(TAG, "Scheduling reconnect attempt $reconnectAttempts in ${delay}ms")
        listener.onReconnecting(reconnectAttempts, delay)
        
        // Cancel any pending reconnection
        reconnectHandler?.removeCallbacksAndMessages(null)
        
        // Schedule reconnection
        reconnectHandler = android.os.Handler(android.os.Looper.getMainLooper())
        reconnectHandler?.postDelayed({
            if (shouldReconnect && !isConnected) {
                Log.d(TAG, "Attempting reconnection (attempt $reconnectAttempts)")
                connect()
            }
        }, delay)
    }
    
    /**
     * Reset reconnection state on successful connection.
     */
    private fun resetReconnectionState() {
        reconnectAttempts = 0
        reconnectHandler?.removeCallbacksAndMessages(null)
        reconnectHandler = null
    }
    
    /**
     * Check if device has network connectivity.
     */
    private fun isNetworkAvailable(): Boolean {
        val connectivityManager = context.getSystemService(Context.CONNECTIVITY_SERVICE) as? ConnectivityManager
        if (connectivityManager == null) {
            Log.w(TAG, "ConnectivityManager not available")
            return true // Assume available if we can't check
        }
        
        val network = connectivityManager.activeNetwork ?: return false
        val capabilities = connectivityManager.getNetworkCapabilities(network) ?: return false
        
        return capabilities.hasCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET) &&
               capabilities.hasCapability(NetworkCapabilities.NET_CAPABILITY_VALIDATED)
    }
    
    /**
     * Handle incoming WebSocket messages.
     */
    private fun handleMessage(text: String) {
        try {
            val json = gson.fromJson(text, JsonObject::class.java)
            val type = json.get("type")?.asString
            
            when (type) {
                "event" -> handleEvent(json)
                "res" -> handleResponse(json)
                "alarm_dismiss" -> {
                    // Handle alarm dismissal broadcast from proxy
                    val alarmId = json.get("alarmId")?.asString
                    if (alarmId != null) {
                        Log.d(TAG, "Received alarm dismiss broadcast: $alarmId")
                        listener.onAlarmDismissed(alarmId)
                    }
                }
                else -> Log.w(TAG, "Unknown message type: $type")
            }
        } catch (e: Exception) {
            Log.e(TAG, "Failed to parse message", e)
        }
    }
    
    /**
     * Handle gateway events.
     */
    private fun handleEvent(json: JsonObject) {
        val event = json.get("event")?.asString
        val payload = json.getAsJsonObject("payload")
        
        Log.d(TAG, "handleEvent: event=$event")
        
        when (event) {
            "connect.challenge" -> {
                // Gateway sent challenge, respond with connect request
                val nonce = payload?.get("nonce")?.asString
                Log.d(TAG, "Received connect challenge with nonce: $nonce")
                sendConnectRequest(nonce)
            }
            "chat" -> {
                // Handle chat message events
                handleChatEvent(payload)
            }
            "mobile.notification" -> {
                // Handle mobile notification events
                handleNotificationEvent(payload)
            }
            "mobile.alarm.dismissed" -> {
                // Handle alarm dismiss broadcast from other devices
                val alarmId = safeGetString(payload, "alarmId")
                if (alarmId != null) {
                    Log.d(TAG, "Received alarm dismiss broadcast: $alarmId")
                    listener.onAlarmDismissed(alarmId)
                }
            }
            "agent" -> {
                // Agent events can be logged but not shown to user
                Log.d(TAG, "Agent event received")
            }
            else -> {
                Log.d(TAG, "Received event: $event")
                listener.onEvent(event ?: "unknown", payload?.toString() ?: "{}")
            }
        }
    }
    
    /**
     * Safely get a string from JsonObject, handling JsonNull.
     */
    private fun safeGetString(obj: JsonObject, key: String): String? {
        val element = obj.get(key) ?: return null
        return if (element.isJsonNull) null else element.asString
    }
    
    /**
     * Handle notification-specific events.
     */
    private fun handleNotificationEvent(payload: JsonObject?) {
        if (payload == null) {
            Log.w(TAG, "Notification event with no payload")
            return
        }
        
        try {
            val notificationType = safeGetString(payload, "notificationType") ?: "alert"
            val title = safeGetString(payload, "title") ?: "AI Assistant"
            val message = safeGetString(payload, "message")
            val priority = safeGetString(payload, "priority") ?: "default"
            val sound = payload.get("sound")?.asBoolean ?: true
            val vibrate = payload.get("vibrate")?.asBoolean ?: true
            val timestamp = payload.get("timestamp")?.asLong ?: System.currentTimeMillis()
            val action = safeGetString(payload, "action")
            
            if (message != null && !message.isEmpty()) {
                Log.d(TAG, "Got notification: type=$notificationType title=$title message=$message")
                listener.onNotification(notificationType, title, message, priority, sound, vibrate, timestamp, action)
            } else {
                Log.w(TAG, "Notification event with no message")
            }
        } catch (e: Exception) {
            Log.e(TAG, "Failed to parse notification event", e)
        }
    }
    
    /**
     * Handle chat-specific events.
     */
    private fun handleChatEvent(payload: JsonObject?) {
        if (payload == null) return
        
        // Extract state and message object
        val state = payload.get("state")?.asString
        val messageObj = payload.getAsJsonObject("message")
        
        if (messageObj == null) {
            Log.w(TAG, "Chat event with no message object")
            return
        }
        
        // Extract role and content array
        val role = messageObj.get("role")?.asString ?: "assistant"
        val contentArray = messageObj.getAsJsonArray("content")
        
        if (contentArray == null || contentArray.size() == 0) {
            Log.w(TAG, "Chat event with empty content array")
            return
        }
        
        // Loop through all content blocks to find text (thinking blocks come first)
        var foundText: String? = null
        for (i in 0 until contentArray.size()) {
            val contentBlock = contentArray.get(i).asJsonObject
            val contentType = contentBlock.get("type")?.asString
            
            // Only extract text blocks, skip thinking/toolCall/etc
            if (contentType == "text") {
                foundText = contentBlock.get("text")?.asString
                if (foundText != null && foundText.isNotEmpty()) {
                    break
                }
            }
        }
        
        if (foundText != null && foundText.isNotEmpty()) {
            // Only show the final message to avoid duplicates with streaming
            if (state == "final") {
                Log.d(TAG, "Got final message: $foundText")
                listener.onMessage("Alfred", foundText)
            } else {
                Log.d(TAG, "Skipping delta state, waiting for final")
            }
        } else {
            Log.d(TAG, "Chat event with no text content blocks")
        }
    }
    
    /**
     * Handle gateway responses.
     */
    private fun handleResponse(json: JsonObject) {
        val id = json.get("id")?.asString
        val ok = json.get("ok")?.asBoolean ?: false
        val payload = json.getAsJsonObject("payload")
        
        if (ok) {
            val payloadType = payload?.get("type")?.asString
            
            if (payloadType == "hello-ok") {
                Log.d(TAG, "Connect successful!")
                isConnected = true
                resetReconnectionState() // Reset reconnection state on successful connection
                
                // Extract and save user preferences if present
                val userPrefs = payload?.getAsJsonObject("userPreferences")
                if (userPrefs != null) {
                    Log.d(TAG, "Received user preferences from server: $userPrefs")
                    val prefs = context.getSharedPreferences("alfred_settings", Context.MODE_PRIVATE)
                    val editor = prefs.edit()
                    
                    // Update assistant name if present
                    if (userPrefs.has("assistantName")) {
                        val assistantName = userPrefs.get("assistantName").asString
                        editor.putString("assistant_name", assistantName)
                        Log.d(TAG, "Updated assistant name from server: $assistantName")
                    }
                    
                    // Update voice ID if present
                    if (userPrefs.has("voiceId")) {
                        val voiceId = userPrefs.get("voiceId").asString
                        editor.putString("tts_voice_id", voiceId)
                        Log.d(TAG, "Updated voice ID from server: $voiceId")
                    }
                    
                    editor.apply()
                } else {
                    Log.d(TAG, "No user preferences in connect response")
                }
                
                listener.onConnected()
            } else {
                listener.onResponse(id ?: "unknown", payload?.toString() ?: "{}")
            }
        } else {
            val error = json.getAsJsonObject("error")
            val errorMsg = error?.get("message")?.asString ?: "Unknown error"
            Log.e(TAG, "Request failed: $errorMsg")
            listener.onError(errorMsg)
        }
    }
    
    /**
     * Send connect request to gateway.
     */
    private fun sendConnectRequest(nonce: String?) {
        val connectMsg = mapOf(
            "type" to "req",
            "id" to "connect-${requestId++}",
            "method" to "connect",
            "params" to mapOf(
                "minProtocol" to 3,
                "maxProtocol" to 3,
                "client" to mapOf(
                    "id" to "cli",
                    "version" to BuildConfig.VERSION_NAME,
                    "platform" to "android",
                    "mode" to "webchat"
                ),
                "role" to "operator",
                "scopes" to listOf("operator.read", "operator.write"),
                "caps" to emptyList<String>(),
                "commands" to emptyList<String>(),
                "permissions" to emptyMap<String, Boolean>(),
                "auth" to mapOf("token" to accessToken),
                "locale" to "en-US",
                "userAgent" to "alfred-mobile/${BuildConfig.VERSION_NAME}"
            )
        )
        
        val json = gson.toJson(connectMsg)
        Log.d(TAG, ">>> Sending connect request (sessionKey will be in chat.send): $json")
        val sent = webSocket?.send(json)
        Log.d(TAG, "Send result: $sent")
    }
    
    /**
     * Send a message to the gateway.
     */
    fun sendMessage(message: String) {
        if (!isConnected) {
            Log.w(TAG, "Not connected, cannot send message")
            listener.onError("Not connected")
            return
        }
        
        val idempotencyKey = "msg-${System.currentTimeMillis()}-${requestId++}"
        
        val msgObj = mapOf(
            "type" to "req",
            "id" to "chat-${requestId++}",
            "method" to "chat.send",
            "params" to mapOf(
                "sessionKey" to userId,
                "message" to message,
                "idempotencyKey" to idempotencyKey
            )
        )
        
        val json = gson.toJson(msgObj)
        Log.d(TAG, "Sending message: $message")
        webSocket?.send(json)
    }
    
    /**
     * Send alarm dismiss event to notify other devices.
     */
    fun dismissAlarm(alarmId: String) {
        if (!isConnected) {
            Log.w(TAG, "Not connected, cannot send alarm dismiss")
            return
        }
        
        val msgObj = mapOf(
            "type" to "alarm.dismiss",
            "alarmId" to alarmId,
            "timestamp" to System.currentTimeMillis()
        )
        
        val json = gson.toJson(msgObj)
        Log.d(TAG, "Sending alarm dismiss: $alarmId")
        webSocket?.send(json)
    }
    
    /**
     * Send FCM token to proxy for push notifications.
     */
    fun sendFCMToken(fcmToken: String) {
        if (!isConnected) {
            Log.w(TAG, "Not connected, cannot send FCM token")
            return
        }
        
        val msgObj = mapOf(
            "type" to "fcm.register",
            "token" to fcmToken,
            "timestamp" to System.currentTimeMillis()
        )
        
        val json = gson.toJson(msgObj)
        Log.d(TAG, "Sending FCM token: ${fcmToken.take(20)}...")
        webSocket?.send(json)
    }
    
    /**
     * Update user preferences on server.
     */
    fun updatePreferences(preferences: Map<String, Any>) {
        if (!isConnected) {
            Log.w(TAG, "Not connected, cannot update preferences")
            return
        }
        
        val msgObj = mapOf(
            "type" to "req",
            "id" to "prefs-update-${requestId++}",
            "method" to "user.preferences.update",
            "params" to preferences
        )
        
        val json = gson.toJson(msgObj)
        Log.d(TAG, "Updating preferences: $preferences")
        webSocket?.send(json)
    }
    
    /**
     * Get user preferences from server.
     */
    fun getPreferences() {
        if (!isConnected) {
            Log.w(TAG, "Not connected, cannot get preferences")
            return
        }
        
        val msgObj = mapOf(
            "type" to "req",
            "id" to "prefs-get-${requestId++}",
            "method" to "user.preferences.get",
            "params" to emptyMap<String, Any>()
        )
        
        val json = gson.toJson(msgObj)
        Log.d(TAG, "Requesting preferences")
        webSocket?.send(json)
    }
    
    /**
     * Disconnect from the gateway.
     */
    fun disconnect() {
        Log.d(TAG, "Disconnecting")
        
        // Disable automatic reconnection
        shouldReconnect = false
        
        // Cancel any pending reconnection attempts
        reconnectHandler?.removeCallbacksAndMessages(null)
        reconnectHandler = null
        
        // Close WebSocket
        isConnected = false
        webSocket?.close(1000, "Client disconnect")
        webSocket = null
        
        // Reset reconnection state
        resetReconnectionState()
    }
    
    /**
     * Check if currently connected.
     */
    fun isConnected(): Boolean = isConnected
}

/**
 * Listener for gateway events.
 */
interface GatewayListener {
    fun onConnecting()
    fun onConnected()
    fun onDisconnected()
    fun onReconnecting(attempt: Int, delayMs: Long)
    fun onError(error: String)
    fun onEvent(event: String, payload: String)
    fun onResponse(id: String, payload: String)
    fun onMessage(sender: String, text: String)
    fun onNotification(
        notificationType: String,
        title: String,
        message: String,
        priority: String,
        sound: Boolean,
        vibrate: Boolean,
        timestamp: Long,
        action: String?
    )
    fun onAlarmDismissed(alarmId: String)
    fun onWakeWordDetected()
}
