package com.nexlab.openclaw.node.data.remote

import android.content.Context
import android.content.SharedPreferences
import android.util.Log
import com.nexlab.openclaw.node.domain.model.NodeCapabilities
import com.nexlab.openclaw.node.domain.model.NodeCommand
import com.nexlab.openclaw.node.domain.model.NodeEvent
import com.nexlab.openclaw.node.domain.model.NodeResponse
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.receiveAsFlow
import okhttp3.*
import org.json.JSONArray
import org.json.JSONObject
import java.security.KeyFactory
import java.security.KeyPairGenerator
import java.security.MessageDigest
import java.security.PrivateKey
import java.security.Signature
import java.security.spec.PKCS8EncodedKeySpec
import java.security.spec.X509EncodedKeySpec
import java.util.Base64
import java.util.concurrent.TimeUnit
import javax.inject.Inject
import javax.inject.Singleton

private const val TAG = "NodeClient"

@Singleton
class NodeClient() {

    private var webSocket: WebSocket? = null
    private val client = OkHttpClient.Builder()
        .readTimeout(30, TimeUnit.SECONDS)
        .pingInterval(20, TimeUnit.SECONDS)
        .build()

    private val _connectionStatus = Channel<ConnectionStatus>(Channel.BUFFERED)
    val connectionStatus: Flow<ConnectionStatus> = _connectionStatus.receiveAsFlow()

    private val _incomingCommands = Channel<NodeCommand>(Channel.BUFFERED)
    val incomingCommands: Flow<NodeCommand> = _incomingCommands.receiveAsFlow()

    private var isConnected = false
    private var currentNodeId: String = ""
    private var currentToken: String = ""
    private var pendingRequestId: String? = null
    private var isPairingMode = false

    data class ConnectionStatus(val connected: Boolean, val nodeId: String)

    // Global gateway token for initial pairing attempt
    private val GLOBAL_GATEWAY_TOKEN = "415fa3c21b7ef06f22aff571697d88c59c2dc67737681267"

    // Device identity
    private var deviceId: String = ""
    private var devicePublicKey: String = ""
    private var devicePrivateKey: String = ""
    private var devicePrivateKeyObject: PrivateKey? = null
    
    // Challenge tracking
    private var challengeNonce: String = ""
    private var challengeTimestamp: Long = 0
    private var hasReceivedChallenge = false
    private var pendingPairingAfterConnect = false  // Flag to send pairing request after reconnect

    // Public method to check if pairing request should be sent
    fun shouldSendPairingRequest(): Boolean = pendingPairingAfterConnect

    // SharedPreferences for device identity persistence
    private var prefs: SharedPreferences? = null
    
    fun setContext(context: Context) {
        prefs = context.getSharedPreferences("openclaw_device_identity", Context.MODE_PRIVATE)
        // Regenerate identity if needed (will load existing if available)
        generateDeviceIdentity()
    }

    init {
        // Generate device identity on initialization
        // Note: setContext should be called after creation to enable persistence
        generateDeviceIdentity()
    }

    private fun generateDeviceIdentity() {
        try {
            // Try to load existing keys from SharedPreferences
            val savedPublicKey = prefs?.getString("device_public_key", null)
            val savedPrivateKey = prefs?.getString("device_private_key", null)
            
            if (savedPublicKey != null && savedPrivateKey != null) {
                // Load existing keys
                devicePublicKey = savedPublicKey
                devicePrivateKey = savedPrivateKey
                
                // Reconstruct PrivateKey object from PEM
                val keyFactory = KeyFactory.getInstance("RSA")
                val privateKeyBytes = Base64.getDecoder().decode(
                    savedPrivateKey
                        .replace("-----BEGIN PRIVATE KEY-----", "")
                        .replace("-----END PRIVATE KEY-----", "")
                        .replace("\n", "")
                )
                val privateKeySpec = PKCS8EncodedKeySpec(privateKeyBytes)
                devicePrivateKeyObject = keyFactory.generatePrivate(privateKeySpec)
                
                // Derive deviceId from public key
                val publicKeyBytes = Base64.getDecoder().decode(
                    savedPublicKey
                        .replace("-----BEGIN PUBLIC KEY-----", "")
                        .replace("-----END PUBLIC KEY-----", "")
                        .replace("\n", "")
                )
                val publicKey = keyFactory.generatePublic(X509EncodedKeySpec(publicKeyBytes))
                val digest = MessageDigest.getInstance("SHA-256")
                val hash = digest.digest(publicKey.encoded)
                deviceId = hash.joinToString("") { "%02x".format(it) }
                
                Log.d(TAG, "Loaded existing device identity: deviceId=$deviceId")
            } else {
                // Generate new key pair
                val keyGen = KeyPairGenerator.getInstance("RSA")
                keyGen.initialize(2048)
                val keyPair = keyGen.generateKeyPair()

                val publicKey = keyPair.public
                val privateKey = keyPair.private
                devicePrivateKeyObject = privateKey

                // Convert to PEM format
                devicePublicKey = "-----BEGIN PUBLIC KEY-----\n" +
                    Base64.getEncoder().encodeToString(publicKey.encoded).chunked(64).joinToString("\n") +
                    "\n-----END PUBLIC KEY-----"

                devicePrivateKey = "-----BEGIN PRIVATE KEY-----\n" +
                    Base64.getEncoder().encodeToString(privateKey.encoded).chunked(64).joinToString("\n") +
                    "\n-----END PRIVATE KEY-----"

                // Generate device ID from public key hash
                val digest = MessageDigest.getInstance("SHA-256")
                val hash = digest.digest(publicKey.encoded)
                deviceId = hash.joinToString("") { "%02x".format(it) }
                
                // Persist keys for future use
                prefs?.edit()
                    ?.putString("device_public_key", devicePublicKey)
                    ?.putString("device_private_key", devicePrivateKey)
                    ?.apply()

                Log.d(TAG, "Generated new device identity: deviceId=$deviceId")
            }
        } catch (e: Exception) {
            Log.e(TAG, "Failed to generate/load device identity: ${e.message}")
            // Generate a random device ID as fallback
            deviceId = "device-${System.currentTimeMillis()}"
        }
    }

    private fun signData(data: String): String {
        return try {
            val privateKey = devicePrivateKeyObject ?: return ""
            val signature = Signature.getInstance("SHA256withRSA")
            signature.initSign(privateKey)
            signature.update(data.toByteArray())
            val sigBytes = signature.sign()
            Base64.getEncoder().encodeToString(sigBytes)
        } catch (e: Exception) {
            Log.e(TAG, "Failed to sign data: ${e.message}")
            ""
        }
    }

    // Sign the v3 payload for device authentication
    // Format: v3|{nonce}|{clientId}|{role}|{scopes}|{token}|{platform}|{deviceFamily}
    private fun signV3Payload(nonce: String, clientId: String, role: String, scopes: JSONArray, token: String, platform: String, deviceFamily: String): String {
        return try {
            val privateKey = devicePrivateKeyObject ?: return ""
            
            // Build v3 payload
            val scopesStr = scopes.join(",")
            val payload = "v3|$nonce|$clientId|$role|$scopesStr|$token|$platform|$deviceFamily"
            
            val signature = Signature.getInstance("SHA256withRSA")
            signature.initSign(privateKey)
            signature.update(payload.toByteArray())
            val sigBytes = signature.sign()
            Base64.getEncoder().encodeToString(sigBytes)
        } catch (e: Exception) {
            Log.e(TAG, "Failed to sign v3 payload: ${e.message}")
            ""
        }
    }

    fun connect(gatewayUrl: String, nodeId: String, token: String, capabilities: NodeCapabilities) {
        Log.d(TAG, "Connecting to $gatewayUrl with nodeId=$nodeId")
        currentNodeId = nodeId
        // Use global token for initial connection attempt (for pairing flow)
        // The user-provided token will be used after pairing
        currentToken = GLOBAL_GATEWAY_TOKEN
        isPairingMode = true

        val request = Request.Builder()
            .url(gatewayUrl)
            .build()

        Log.e(TAG, "Opening WebSocket connection to: $gatewayUrl")
        webSocket = client.newWebSocket(request, object : WebSocketListener() {
            override fun onOpen(webSocket: WebSocket, response: Response) {
                Log.e(TAG, "WebSocket opened!")
                // Wait for challenge, don't send anything yet
            }

            override fun onMessage(webSocket: WebSocket, text: String) {
                try {
                    val json = JSONObject(text)
                    val type = json.getString("type")
                    
                    when (type) {
                        "event" -> {
                            val event = json.optString("event", "")
                            if (event == "connect.challenge") {
                                // Capture the challenge nonce
                                val payload = json.optJSONObject("payload")
                                challengeNonce = payload?.optString("nonce", "") ?: ""
                                challengeTimestamp = payload?.optLong("ts", System.currentTimeMillis()) ?: System.currentTimeMillis()
                                hasReceivedChallenge = true
                                Log.d(TAG, "Received challenge: nonce=$challengeNonce, ts=$challengeTimestamp")
                                
                                // BUG FIX: Always send connect first - the gateway requires "connect" as the first request
                                // If connect is rejected with "device-required", we handle that in the response handler
                                // This was causing "invalid handshake: first request must be connect" error
                                pendingPairingAfterConnect = false
                                Log.d(TAG, "Always sending connect request first (protocol requirement)")
                                sendConnectRequestWithChallenge(nodeId, token, capabilities)
                            } else if (event == "node.pair.resolved") {
                                // Pairing was approved, now reconnect with token
                                val approved = json.optJSONObject("payload")?.optBoolean("approved", false)
                                if (approved == true) {
                                    // Get the token from the response and reconnect
                                    val newToken = json.optJSONObject("payload")?.optString("token", "") ?: ""
                                    isPairingMode = false
                                    sendConnectRequestWithChallenge(nodeId, newToken, capabilities)
                                }
                            }
                        }
                        "res" -> {
                            val id = json.optString("id", "")
                            val ok = json.optBoolean("ok", false)
                            
                            if (id == pendingRequestId) {
                                if (ok) {
                                    isConnected = true
                                    isPairingMode = false
                                    _connectionStatus.trySend(ConnectionStatus(true, nodeId))
                                } else {
                                    // Check if it's a device-required or device signature error
                                    val error = json.optJSONObject("error")?.optString("message", "Connection failed")
                                    Log.d(TAG, "Connection failed with error: $error")
                                    if ((error?.contains("device-required") == true || 
                                         error?.contains("device signature") == true ||
                                         error?.contains("invalid") == true) && isPairingMode) {
                                        // Device not paired - send pairing request on THIS connection
                                        Log.d(TAG, "Device not paired, sending pairing request")
                                        sendPairingRequest(nodeId, token, capabilities)
                                        // Don't close - wait for pairing response
                                    } else {
                                        _connectionStatus.trySend(ConnectionStatus(false, nodeId))
                                    }
                                }
                                pendingRequestId = null
                            }
                        }
                        else -> {
                            // Handle commands and other messages
                            val command = parseCommand(json)
                            command?.let { _incomingCommands.trySend(it) }
                        }
                    }
                } catch (e: Exception) {
                    e.printStackTrace()
                }
            }

            override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
                Log.d(TAG, "WebSocket closing: code=$code reason=$reason")
                
                // Check if it's a device signature error - send pairing request instead of closing
                if (isPairingMode && (reason.contains("device signature") || 
                                      reason.contains("invalid") ||
                                      reason.contains("device-required"))) {
                    Log.d(TAG, "Device auth rejected, sending pairing request instead of closing")
                    // Send pairing request on this connection instead of closing
                    // Note: We need the token and capabilities, but they're stored in local vars
                    // For now, use the values from currentNodeId and a placeholder
                    sendPairingRequest(currentNodeId, currentToken, NodeCapabilities(
                        screenshot = true, touch = true, shell = true,
                        notifications = true, appControl = true, sensors = true
                    ))
                    // Don't close - wait for pairing response
                    return
                }
                
                webSocket.close(1000, null)
                // Emit connection status
                _connectionStatus.trySend(ConnectionStatus(false, nodeId))
            }

            override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
                Log.d(TAG, "WebSocket closed: code=$code reason=$reason")
                isConnected = false
                
                // Also check here in case onClosing didn't handle it
                if (isPairingMode && (reason.contains("device signature") || 
                                      reason.contains("invalid") ||
                                      reason.contains("device-required"))) {
                    Log.d(TAG, "Connection rejected due to device auth, will send pairing request on reconnect")
                    // Set flag to send pairing request after reconnect
                    pendingPairingAfterConnect = true
                    // Don't reconnect here - let the app retry naturally
                    // The next connect() call will check pendingPairingAfterConnect
                }
                
                // Always emit connection status
                _connectionStatus.trySend(ConnectionStatus(false, nodeId))
            }

            override fun onFailure(webSocket: WebSocket, t: Throwable, t2: Response?) {
                Log.e(TAG, "WebSocket failure: ${t.message}", t)
                isConnected = false
                _connectionStatus.trySend(ConnectionStatus(false, nodeId))
            }
        })
    }

    private fun sendConnectRequest(nodeId: String, token: String, capabilities: NodeCapabilities) {
        pendingRequestId = "connect-${System.currentTimeMillis()}"
        
        val connectParams = JSONObject().apply {
            put("minProtocol", 3)
            put("maxProtocol", 3)
            put("client", JSONObject().apply {
                put("id", "cli")  // Use "cli" client ID - only one that works with gateway
                put("displayName", nodeId.ifBlank { "Android Node" })
                put("version", "1.0.0")
                put("platform", "android")
                put("mode", "node")
            })
            put("role", "node")
            put("scopes", JSONArray())
            put("caps", JSONArray().apply {
                if (capabilities.screenshot) put("screenshot")
                if (capabilities.touch) put("touch")
                if (capabilities.shell) put("shell")
                if (capabilities.notifications) put("notifications")
                if (capabilities.appControl) put("appControl")
                if (capabilities.sensors) put("sensors")
            })
            put("commands", JSONArray().apply {
                put("screenshot")
                put("touch.*")
                put("shell")
                put("notifications")
                put("app.launch")
                put("device.info")
                put("apps.list")
            })
            put("permissions", JSONObject())
            put("auth", JSONObject().apply {
                put("token", token)
            })
            put("device", JSONObject().apply {
                put("id", deviceId)
                put("publicKey", devicePublicKey)
                put("si", deviceId)  // Service identity = device ID
                put("signature", signData(token))  // Sign the token
            })
            put("locale", "en-US")
            put("userAgent", "openclaw-android/1.0.0")
        }

        val connectRequest = JSONObject().apply {
            put("type", "req")
            put("id", pendingRequestId)
            put("method", "connect")
            put("params", connectParams)
        }

        sendJson(connectRequest)
    }

    private fun sendConnectRequestWithChallenge(nodeId: String, token: String, capabilities: NodeCapabilities) {
        pendingRequestId = "connect-${System.currentTimeMillis()}"
        
        // Use v3 signature format
        val scopesArray = JSONArray()
        
        // Sign the v3 payload - clientId should be deviceId
        val signature = signV3Payload(
            nonce = challengeNonce,
            clientId = deviceId,  // Use device ID for signature
            role = "node",
            scopes = scopesArray,
            token = token,
            platform = "android",
            deviceFamily = "phone"
        )
        
        val connectParams = JSONObject().apply {
            put("minProtocol", 3)
            put("maxProtocol", 3)
            put("client", JSONObject().apply {
                put("id", "cli")  // Use "cli" client ID - only one that works with gateway
                put("displayName", nodeId.ifBlank { "Android Node" })
                put("version", "1.0.0")
                put("platform", "android")
                put("mode", "node")
            })
            put("role", "node")
            put("scopes", scopesArray)
            put("caps", JSONArray().apply {
                if (capabilities.screenshot) put("screenshot")
                if (capabilities.touch) put("touch")
                if (capabilities.shell) put("shell")
                if (capabilities.notifications) put("notifications")
                if (capabilities.appControl) put("appControl")
                if (capabilities.sensors) put("sensors")
            })
            put("commands", JSONArray().apply {
                put("screenshot")
                put("touch.*")
                put("shell")
                put("notifications")
                put("app.launch")
                put("device.info")
                put("apps.list")
            })
            put("permissions", JSONObject())
            put("auth", JSONObject().apply {
                put("token", token)
            })
            put("device", JSONObject().apply {
                put("id", deviceId)
                put("publicKey", devicePublicKey)
                put("signature", signature)
                put("signedAt", challengeTimestamp)
                put("nonce", challengeNonce)
            })
            put("locale", "en-US")
            put("userAgent", "openclaw-android/1.0.0")
        }

        val connectRequest = JSONObject().apply {
            put("type", "req")
            put("id", pendingRequestId)
            put("method", "connect")
            put("params", connectParams)
        }

        Log.d(TAG, "Sending connect request with v3 signature")
        sendJson(connectRequest)
        
        // Don't send pairing request here - wait for connect response first
    }

    private fun sendPairingRequest(nodeId: String, token: String, capabilities: NodeCapabilities) {
        pendingRequestId = "pair-request-${System.currentTimeMillis()}"
        
        val pairRequest = JSONObject().apply {
            put("type", "req")
            put("id", pendingRequestId)
            put("method", "node.pair.request")
            put("params", JSONObject().apply {
                put("nodeId", nodeId)
                put("displayName", nodeId)
                put("platform", "android")
                put("version", "1.0.0")
                put("caps", JSONArray().apply {
                    if (capabilities.screenshot) put("screenshot")
                    if (capabilities.touch) put("touch")
                    if (capabilities.shell) put("shell")
                    if (capabilities.notifications) put("notifications")
                    if (capabilities.appControl) put("appControl")
                    if (capabilities.sensors) put("sensors")
                })
                put("commands", JSONArray().apply {
                    put("screenshot")
                    put("touch.*")
                    put("shell")
                    put("notifications")
                    put("app.launch")
                    put("device.info")
                    put("apps.list")
                })
            })
        }
        
        sendJson(pairRequest)
    }

    fun requestPairing(nodeId: String) {
        pendingRequestId = "pair-request-${System.currentTimeMillis()}"
        
        val pairRequest = JSONObject().apply {
            put("type", "req")
            put("id", pendingRequestId)
            put("method", "node.pair.request")
            put("params", JSONObject().apply {
                put("nodeId", nodeId)
            })
        }
        
        sendJson(pairRequest)
    }

    fun disconnect() {
        webSocket?.close(1000, "User disconnected")
        webSocket = null
        isConnected = false
    }

    fun sendResponse(response: NodeResponse) {
        val json = when (response) {
            is NodeResponse.Success -> JSONObject().apply {
                put("type", "response")
                put("id", response.id)
                put("ok", true)
                put("payload", JSONObject().apply {
                    response.data.forEach { (key, value) ->
                        put(key, value)
                    }
                })
            }
            is NodeResponse.Error -> JSONObject().apply {
                put("type", "res")
                put("id", response.id)
                put("ok", false)
                put("error", JSONObject().apply {
                    put("message", response.error)
                })
            }
            is NodeResponse.ScreenshotResponse -> JSONObject().apply {
                put("type", "res")
                put("id", response.id)
                put("ok", true)
                put("payload", JSONObject().apply {
                    put("screenshot", response.base64Image)
                })
            }
        }
        sendJson(json)
    }

    fun sendEvent(event: NodeEvent) {
        val json = when (event) {
            is NodeEvent.Notification -> JSONObject().apply {
                put("type", "event")
                put("event", "node.notification")
                put("payload", JSONObject().apply {
                    put("app", event.app)
                    put("title", event.title)
                    put("text", event.text)
                    put("timestamp", event.timestamp)
                })
            }
            is NodeEvent.Battery -> JSONObject().apply {
                put("type", "event")
                put("event", "node.battery")
                put("payload", JSONObject().apply {
                    put("level", event.level)
                    put("charging", event.charging)
                })
            }
            is NodeEvent.ForegroundAppChanged -> JSONObject().apply {
                put("type", "event")
                put("event", "node.foregroundApp")
                put("payload", JSONObject().apply {
                    put("package", event.packageName)
                    put("app_name", event.appName)
                })
            }
            is NodeEvent.ConnectionStatus -> JSONObject().apply {
                put("type", "event")
                put("event", "node.connection")
                put("payload", JSONObject().apply {
                    put("connected", event.connected)
                    put("node_id", event.nodeId)
                })
            }
        }
        sendJson(json)
    }

    private fun sendJson(json: JSONObject) {
        Log.d(TAG, "Sending: ${json.toString().take(200)}")
        webSocket?.send(json.toString())
    }

    private fun parseCommand(json: JSONObject): NodeCommand? {
        val type = json.getString("type")
        val id = json.optString("id", "")

        // Handle both "req" and direct command types
        val commandType = if (type == "req") {
            json.optString("method", "")
        } else {
            type
        }

        return when (commandType) {
            "screenshot" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                val quality = params?.optInt("quality", 80) ?: 80
                NodeCommand.Screenshot(id, quality)
            }
            "tap" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                NodeCommand.Tap(id, params?.getInt("x") ?: 0, params?.getInt("y") ?: 0)
            }
            "swipe" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                NodeCommand.Swipe(
                    id,
                    params?.getInt("x1") ?: 0, params?.getInt("y1") ?: 0,
                    params?.getInt("x2") ?: 0, params?.getInt("y2") ?: 0,
                    params?.optInt("duration", 300) ?: 300
                )
            }
            "long_press" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                NodeCommand.LongPress(
                    id,
                    params?.getInt("x") ?: 0, params?.getInt("y") ?: 0,
                    params?.optInt("duration", 500) ?: 500
                )
            }
            "launch_app" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                NodeCommand.LaunchApp(id, params?.getString("package") ?: "")
            }
            "shell" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                NodeCommand.Shell(id, params?.getString("command") ?: "")
            }
            "device_info" -> NodeCommand.DeviceInfo(id)
            "list_apps" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                val systemApps = params?.optBoolean("system_apps", false) ?: false
                NodeCommand.ListApps(id, systemApps)
            }
            "foreground_app" -> NodeCommand.GetForegroundApp(id)
            "notification_listener" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                val enabled = params?.optBoolean("enabled", true) ?: true
                NodeCommand.NotificationListener(id, enabled)
            }
            "key_event" -> {
                val params = json.optJSONObject("params") ?: json.optJSONObject("payload")
                NodeCommand.KeyEvent(id, params?.getInt("key_code") ?: 0)
            }
            "ping" -> NodeCommand.Ping
            else -> null
        }
    }
}
