package com.inou.clawdnode.service

import android.util.Log
import com.google.gson.Gson
import com.google.gson.JsonObject
import com.inou.clawdnode.BuildConfig
import com.inou.clawdnode.ClawdNodeApp
import com.inou.clawdnode.protocol.Protocol
import com.inou.clawdnode.protocol.*
import com.inou.clawdnode.security.DeviceIdentity
import kotlinx.coroutines.*
import okhttp3.*
import java.util.UUID
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicInteger

/**
 * WebSocket client for Clawdbot Gateway connection.
 * 
 * Implements Gateway Protocol v3:
 * 1. Connect to /ws endpoint
 * 2. Receive connect.challenge event with nonce
 * 3. Send connect request with role, caps, commands, auth
 * 4. Handle hello-ok response
 * 
 * After handshake, handles message routing for node commands.
 */
class GatewayClient(
    private val onCommand: (NodeCommand) -> Unit,
    private val onConnectionChange: (Boolean) -> Unit,
    var onLog: ((String) -> Unit)? = null
) {
    
    private fun log(message: String) {
        Log.d(tag, message)
        onLog?.invoke(message)
    }
    
    private fun logError(message: String, t: Throwable? = null) {
        if (t != null) {
            Log.e(tag, message, t)
            onLog?.invoke("ERROR: $message - ${t.message}")
        } else {
            Log.e(tag, message)
            onLog?.invoke("ERROR: $message")
        }
    }
    private val tag = "GatewayClient"
    private val gson = Gson()
    
    private val client = OkHttpClient.Builder()
        .connectTimeout(10, TimeUnit.SECONDS)
        .readTimeout(0, TimeUnit.SECONDS) // No timeout for WebSocket
        .writeTimeout(10, TimeUnit.SECONDS)
        .pingInterval(30, TimeUnit.SECONDS)
        .build()
    
    private var webSocket: WebSocket? = null
    private var isConnected = false
    private var isHandshakeComplete = false
    private var shouldReconnect = true
    private val scope = CoroutineScope(Dispatchers.IO + SupervisorJob())
    private val requestIdCounter = AtomicInteger(0)
    
    private val auditLog get() = ClawdNodeApp.instance.auditLog
    private val tokenStore get() = ClawdNodeApp.instance.tokenStore
    
    private val deviceIdentity by lazy { DeviceIdentity(ClawdNodeApp.instance) }
    
    // Node capabilities
    private val caps = listOf("notifications", "calls", "voice")
    private val commands = listOf(
        "notification.action",
        "notification.dismiss", 
        "call.answer",
        "call.reject",
        "call.speak",
        "call.hangup"
    )
    
    fun connect() {
        val url = tokenStore.gatewayUrl ?: run {
            log("No gateway URL configured")
            return
        }
        
        if (tokenStore.gatewayToken == null) {
            log("No gateway token configured")
            return
        }
        
        shouldReconnect = true
        isHandshakeComplete = false
        
        // Build WebSocket URL - connect to /ws (not /ws/node)
        val wsUrl = buildWsUrl(url)
        log("Connecting to $wsUrl")
        
        val request = Request.Builder()
            .url(wsUrl)
            .build()
        
        webSocket = client.newWebSocket(request, object : WebSocketListener() {
            override fun onOpen(webSocket: WebSocket, response: Response) {
                log("WebSocket connected, waiting for challenge...")
                // Don't set isConnected yet - wait for handshake
            }
            
            override fun onMessage(webSocket: WebSocket, text: String) {
                log("Received: $text")
                handleMessage(text)
            }
            
            override fun onClosing(webSocket: WebSocket, code: Int, reason: String) {
                log("Connection closing: $code $reason")
                webSocket.close(1000, null)
            }
            
            override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
                log("Connection closed: $code $reason")
                handleDisconnect()
            }
            
            override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
                logError("Connection failed", t)
                auditLog.log("GATEWAY_ERROR", "Connection failed: ${t.message}")
                handleDisconnect()
            }
        })
    }
    
    fun disconnect() {
        shouldReconnect = false
        webSocket?.close(1000, "Client disconnect")
        webSocket = null
        isConnected = false
        isHandshakeComplete = false
        onConnectionChange(false)
    }
    
    fun send(event: NodeEvent) {
        val requestId = generateRequestId()
        val json = event.toProtocolFrame(requestId)
        log("Sending: $json")
        
        if (isConnected && isHandshakeComplete) {
            webSocket?.send(json)
        } else {
            log("Not connected or handshake incomplete, cannot send event")
            // TODO: Queue for retry
        }
    }
    
    /**
     * Send a response to a node.invoke request
     */
    fun sendResponse(requestId: String, success: Boolean, payload: Any? = null, error: String? = null) {
        val frame = if (success) {
            ResponseFrame(id = requestId, ok = true, payload = payload)
        } else {
            ResponseFrame(id = requestId, ok = false, error = ErrorPayload("ERROR", error ?: "Unknown error"))
        }
        
        val json = frame.toJson()
        log("Sending response: $json")
        
        if (isConnected && isHandshakeComplete) {
            webSocket?.send(json)
        }
    }
    
    private fun buildWsUrl(baseUrl: String): String {
        // Convert http(s) to ws(s)
        val wsBase = baseUrl
            .replace("http://", "ws://")
            .replace("https://", "wss://")
            .trimEnd('/')
        
        // Connect to /ws endpoint (no token in URL - sent in handshake)
        return "$wsBase/ws"
    }
    
    private fun handleMessage(json: String) {
        val frame = ProtocolFrame.fromJson(json)
        if (frame == null) {
            log("Failed to parse frame: $json")
            return
        }
        
        when (frame.type) {
            "event" -> handleEvent(frame)
            "req" -> handleRequest(frame)
            "res" -> handleResponse(frame)
            else -> log("Unknown frame type: ${frame.type}")
        }
    }
    
    private fun handleEvent(frame: ProtocolFrame) {
        when (frame.event) {
            "connect.challenge" -> handleConnectChallenge(frame)
            else -> log("Received event: ${frame.event}")
        }
    }
    
    private fun handleConnectChallenge(frame: ProtocolFrame) {
        log("Received connect.challenge, sending handshake...")
        
        val payload = frame.payload
        val nonce = payload?.get("nonce")?.asString
        
        if (nonce == null) {
            logError("No nonce in connect.challenge")
            return
        }
        
        val token = tokenStore.gatewayToken ?: run {
            logError("No token available for handshake")
            return
        }
        
        // Sign the challenge nonce and get device identity
        // Wrap in try/catch - keystore operations can fail on first run or if hardware unavailable
        val signedChallenge: DeviceIdentity.SignedChallenge?
        val deviceId: String
        val publicKey: String?
        
        try {
            deviceId = deviceIdentity.deviceId
            signedChallenge = deviceIdentity.signChallenge(nonce)
            publicKey = deviceIdentity.publicKey
            log("Device identity ready: id=${deviceId.take(8)}..., signed challenge")
        } catch (e: Exception) {
            logError("Failed to initialize device identity or sign challenge", e)
            // Cannot proceed without device identity for non-local connections
            webSocket?.close(1000, "Device identity initialization failed: ${e.message}")
            return
        }
        
        val connectRequest = ConnectRequest(
            id = generateRequestId(),
            params = ConnectParams(
                client = ClientInfo(
                    id = Protocol.CLIENT_ID,
                    version = getAppVersion(),
                    platform = Protocol.PLATFORM,
                    mode = Protocol.MODE
                ),
                role = Protocol.ROLE,
                caps = caps,
                commands = commands,
                permissions = mapOf(
                    "notifications.read" to true,
                    "notifications.action" to true,
                    "calls.answer" to true,
                    "calls.reject" to true,
                    "calls.speak" to true
                ),
                auth = AuthInfo(token = token),
                userAgent = "clawdnode-android/${getAppVersion()}",
                device = DeviceInfo(
                    id = deviceId,
                    publicKey = publicKey,
                    signature = signedChallenge.signature,
                    signedAt = signedChallenge.signedAt,
                    nonce = signedChallenge.nonce
                )
            )
        )
        
        val requestJson = connectRequest.toJson()
        log("Sending connect request: $requestJson")
        webSocket?.send(requestJson)
    }
    
    private fun handleRequest(frame: ProtocolFrame) {
        when (frame.method) {
            "node.invoke" -> handleNodeInvoke(frame)
            else -> {
                log("Unknown request method: ${frame.method}")
                // Send error response
                frame.id?.let { id ->
                    sendResponse(id, false, error = "Unknown method: ${frame.method}")
                }
            }
        }
    }
    
    private fun handleNodeInvoke(frame: ProtocolFrame) {
        val requestId = frame.id ?: run {
            logError("node.invoke missing request id")
            return
        }
        
        val params = frame.params ?: run {
            logError("node.invoke missing params")
            sendResponse(requestId, false, error = "Missing params")
            return
        }
        
        val command = params.get("command")?.asString
        val args = params.get("args")?.asJsonObject
        
        if (command == null) {
            logError("node.invoke missing command")
            sendResponse(requestId, false, error = "Missing command")
            return
        }
        
        log("Received node.invoke: $command")
        auditLog.logCommand(command, "gateway", true)
        
        val invokeParams = NodeInvokeParams(command = command, args = args)
        val nodeCommand = NodeCommand.fromInvoke(requestId, invokeParams)
        
        if (nodeCommand != null) {
            onCommand(nodeCommand)
        } else {
            log("Unknown command: $command")
            sendResponse(requestId, false, error = "Unknown command: $command")
        }
    }
    
    private fun handleResponse(frame: ProtocolFrame) {
        val ok = frame.ok ?: false
        val payload = frame.payload
        
        // Check if this is hello-ok response
        val payloadType = payload?.get("type")?.asString
        if (payloadType == "hello-ok") {
            handleHelloOk(frame)
            return
        }
        
        if (!ok) {
            val error = frame.error
            logError("Request failed: $error")
        }
    }
    
    private fun handleHelloOk(frame: ProtocolFrame) {
        log("Received hello-ok, handshake complete!")
        
        val payload = frame.payload
        val protocol = payload?.get("protocol")?.asInt ?: Protocol.VERSION
        
        // Check for device token
        val auth = payload?.get("auth")?.asJsonObject
        val deviceToken = auth?.get("deviceToken")?.asString
        if (deviceToken != null) {
            log("Received device token, storing for future connects")
            // Could store this for future use
        }
        
        isConnected = true
        isHandshakeComplete = true
        onConnectionChange(true)
        auditLog.log("GATEWAY_CONNECTED", "Protocol v$protocol handshake complete")
    }
    
    private fun handleDisconnect() {
        isConnected = false
        isHandshakeComplete = false
        onConnectionChange(false)
        
        if (shouldReconnect) {
            auditLog.log("GATEWAY_RECONNECT", "Scheduling reconnect in 5s")
            scope.launch {
                delay(5000)
                if (shouldReconnect) {
                    connect()
                }
            }
        }
    }
    
    private fun generateRequestId(): String {
        return "req_${requestIdCounter.incrementAndGet()}_${UUID.randomUUID().toString().take(8)}"
    }
    
    private fun getAppVersion(): String {
        return try {
            BuildConfig.VERSION_NAME
        } catch (e: Exception) {
            "0.1.0"
        }
    }
    
    fun isConnected() = isConnected && isHandshakeComplete
}