package com.durian.tssparty.data.remote

import android.util.Base64
import android.util.Log
import com.durian.tssparty.domain.model.Participant
import com.durian.tssparty.grpc.*
import io.grpc.ManagedChannel
import io.grpc.ManagedChannelBuilder
import io.grpc.ConnectivityState
import io.grpc.stub.StreamObserver
import kotlinx.coroutines.*
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.*
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger
import javax.inject.Inject
import javax.inject.Singleton
import kotlin.math.min
import kotlin.math.pow

/**
 * Connection state for gRPC client
 */
sealed class GrpcConnectionState {
    object Disconnected : GrpcConnectionState()
    object Connecting : GrpcConnectionState()
    object Connected : GrpcConnectionState()
    data class Reconnecting(val attempt: Int, val maxAttempts: Int) : GrpcConnectionState()
    data class Failed(val reason: String) : GrpcConnectionState()
}

/**
 * Connection events for UI notification
 */
sealed class GrpcConnectionEvent {
    data class Disconnected(val reason: String) : GrpcConnectionEvent()
    object Reconnected : GrpcConnectionEvent()
    data class ReconnectFailed(val reason: String) : GrpcConnectionEvent()
    data class PendingMessages(val count: Int) : GrpcConnectionEvent()
}

/**
 * Reconnect configuration
 */
data class ReconnectConfig(
    val maxRetries: Int = 10,
    val initialDelayMs: Long = 1000,
    val maxDelayMs: Long = 30000,
    val backoffMultiplier: Double = 2.0
)

/**
 * gRPC client for Message Router service
 *
 * Features:
 * - Automatic reconnection with exponential backoff
 * - Heartbeat mechanism with failure detection
 * - Stream auto-recovery on disconnect
 * - Connection state notifications via StateFlow
 */
@Singleton
class GrpcClient @Inject constructor() {

    companion object {
        private const val TAG = "GrpcClient"
        private const val HEARTBEAT_INTERVAL_MS = 30000L
        private const val MAX_HEARTBEAT_FAILS = 3
        private const val CONNECTION_TIMEOUT_SECONDS = 10L
        private const val REQUEST_TIMEOUT_SECONDS = 30L
        private const val MAX_REQUEST_RETRIES = 3
        private const val REQUEST_RETRY_DELAY_MS = 500L
    }

    private var channel: ManagedChannel? = null
    private var stub: MessageRouterGrpc.MessageRouterBlockingStub? = null
    private var asyncStub: MessageRouterGrpc.MessageRouterStub? = null

    // Connection state
    private val _connectionState = MutableStateFlow<GrpcConnectionState>(GrpcConnectionState.Disconnected)
    val connectionState: StateFlow<GrpcConnectionState> = _connectionState.asStateFlow()

    // Connection events (for UI notifications)
    private val _connectionEvents = MutableSharedFlow<GrpcConnectionEvent>(extraBufferCapacity = 10)
    val connectionEvents: SharedFlow<GrpcConnectionEvent> = _connectionEvents.asSharedFlow()

    // Reconnection state
    private val reconnectConfig = ReconnectConfig()
    private var currentHost: String? = null
    private var currentPort: Int? = null
    private val isReconnecting = AtomicBoolean(false)
    private val reconnectAttempts = AtomicInteger(0)
    private val shouldReconnect = AtomicBoolean(true)
    private var reconnectJob: Job? = null

    // Registration state (for re-registration after reconnect)
    private var registeredPartyId: String? = null
    private var registeredPartyRole: String? = null

    // Additional signing party registration (for imported/restored shares)
    // When signing with a restored wallet, the signing partyId differs from the device partyId
    // and must also be registered with the message-router
    private var registeredSigningPartyId: String? = null

    // Heartbeat state
    private var heartbeatJob: Job? = null
    private val heartbeatFailCount = AtomicInteger(0)

    // Stream subscriptions (for recovery after reconnect)
    private var activeMessageSubscription: MessageSubscription? = null
    private var eventStreamSubscribed = AtomicBoolean(false)
    private var eventStreamPartyId: String? = null

    // Stream version tracking (to detect stale stream events)
    private val messageStreamVersion = AtomicInteger(0)
    private val eventStreamVersion = AtomicInteger(0)

    // Reconnection callback - called when streams need to be re-established
    private var onReconnectedCallback: (() -> Unit)? = null

    // Channel state monitoring
    private var channelStateMonitorJob: Job? = null

    // Coroutine scope for background tasks
    private val scope = CoroutineScope(Dispatchers.IO + SupervisorJob())

    /**
     * Connect to the Message Router server
     */
    fun connect(host: String, port: Int) {
        Log.d(TAG, "=== connect() called ===")
        Log.d(TAG, "  host: $host, port: $port")
        Log.d(TAG, "  isReconnecting before reset: ${isReconnecting.get()}")

        // Save connection params for reconnection
        currentHost = host
        currentPort = port
        shouldReconnect.set(true)
        reconnectAttempts.set(0)

        // 重要：初次连接时确保 isReconnecting 为 false
        // 这样 waitForConnection() 能正确区分初次连接和重连
        isReconnecting.set(false)
        Log.d(TAG, "  isReconnecting after reset: ${isReconnecting.get()} (should be false for first connect)")

        doConnect(host, port)
    }

    private fun doConnect(host: String, port: Int) {
        Log.d(TAG, "doConnect: $host:$port, isReconnecting=${isReconnecting.get()}")
        _connectionState.value = GrpcConnectionState.Connecting

        try {
            // Disconnect existing connection
            cleanupConnection()

            val builder = ManagedChannelBuilder
                .forAddress(host, port)
                .keepAliveTime(30, TimeUnit.SECONDS)
                .keepAliveTimeout(10, TimeUnit.SECONDS)
                .keepAliveWithoutCalls(true)
                .idleTimeout(5, TimeUnit.MINUTES)

            // Use TLS for port 443, plaintext for other ports (like local development)
            if (port == 443) {
                builder.useTransportSecurity()
            } else {
                builder.usePlaintext()
            }

            channel = builder.build()

            // Create stubs (deadline set per-call via getStubWithDeadline())
            stub = MessageRouterGrpc.newBlockingStub(channel)
            asyncStub = MessageRouterGrpc.newStub(channel)

            // Wait for connection to be ready
            scope.launch {
                waitForConnection()
            }

        } catch (e: Exception) {
            Log.e(TAG, "Connection failed: ${e.message}")
            _connectionState.value = GrpcConnectionState.Failed(e.message ?: "Unknown error")
            triggerReconnect("Connection failed: ${e.message}")
        }
    }

    private suspend fun waitForConnection() {
        val ch = channel ?: return

        try {
            withTimeout(CONNECTION_TIMEOUT_SECONDS * 1000) {
                while (true) {
                    val state = ch.getState(true)
                    Log.d(TAG, "Channel state: $state")

                    when (state) {
                        ConnectivityState.READY -> {
                            // 关键修复：先读取 isReconnecting 再重置，用于区分初次连接和重连
                            // - 初次连接：isReconnecting = false（由 connect() 触发）
                            // - 重连：isReconnecting = true（由 triggerReconnect() 触发，包括后台唤醒）
                            val wasReconnecting = isReconnecting.getAndSet(false)

                            Log.d(TAG, "=== Channel READY ===")
                            Log.d(TAG, "  wasReconnecting: $wasReconnecting")
                            Log.d(TAG, "  registeredPartyId: $registeredPartyId")
                            Log.d(TAG, "  eventStreamSubscribed: ${eventStreamSubscribed.get()}")
                            Log.d(TAG, "  eventStreamPartyId: $eventStreamPartyId")

                            _connectionState.value = GrpcConnectionState.Connected
                            reconnectAttempts.set(0)
                            heartbeatFailCount.set(0)

                            // Start channel state monitoring
                            startChannelStateMonitor()

                            // 只有重连时才需要恢复注册和订阅
                            // 初次连接时，registerParty() 和 subscribeSessionEvents() 会在外部显式调用
                            if (wasReconnecting) {
                                Log.d(TAG, ">>> RECONNECT: Restoring registration and streams")
                                // Re-register if we were registered before
                                reRegisterIfNeeded()
                                // Re-subscribe to streams
                                reSubscribeStreams()
                            } else {
                                Log.d(TAG, ">>> FIRST CONNECT: Skipping restore (will be done by caller)")
                            }

                            // Restart heartbeat (both first connect and reconnect need this)
                            startHeartbeat()

                            return@withTimeout
                        }
                        ConnectivityState.TRANSIENT_FAILURE, ConnectivityState.SHUTDOWN -> {
                            throw Exception("Connection failed with state: $state")
                        }
                        else -> {
                            delay(100)
                        }
                    }
                }
            }
        } catch (e: Exception) {
            Log.e(TAG, "Wait for connection failed: ${e.message}")
            _connectionState.value = GrpcConnectionState.Failed(e.message ?: "Connection timeout")
            triggerReconnect("Connection timeout")
        }
    }

    /**
     * Disconnect from the server (won't auto-reconnect)
     */
    fun disconnect() {
        Log.d(TAG, "Disconnecting")
        shouldReconnect.set(false)
        cleanupConnection()
        _connectionState.value = GrpcConnectionState.Disconnected
    }

    /**
     * Cleanup all connection resources
     */
    private fun cleanupConnection() {
        // Cancel reconnect job
        reconnectJob?.cancel()
        reconnectJob = null

        // Stop channel state monitor
        channelStateMonitorJob?.cancel()
        channelStateMonitorJob = null

        // Stop heartbeat
        stopHeartbeat()

        // Increment stream versions to invalidate old stream callbacks
        messageStreamVersion.incrementAndGet()
        eventStreamVersion.incrementAndGet()

        // Shutdown channel
        channel?.let { ch ->
            try {
                ch.shutdown()
                val terminated = ch.awaitTermination(2, TimeUnit.SECONDS)
                if (!terminated) {
                    ch.shutdownNow()
                }
            } catch (e: Exception) {
                Log.e(TAG, "Error shutting down channel: ${e.message}")
            }
            Unit
        }
        channel = null
        stub = null
        asyncStub = null
    }

    /**
     * Start monitoring channel state for disconnections
     */
    private fun startChannelStateMonitor() {
        channelStateMonitorJob?.cancel()
        channelStateMonitorJob = scope.launch {
            val ch = channel ?: return@launch
            var lastState = ConnectivityState.READY

            while (isActive && _connectionState.value == GrpcConnectionState.Connected) {
                delay(5000) // Check every 5 seconds

                try {
                    val currentState = ch.getState(false)
                    if (currentState != lastState) {
                        Log.d(TAG, "Channel state changed: $lastState -> $currentState")
                        lastState = currentState

                        when (currentState) {
                            ConnectivityState.TRANSIENT_FAILURE, ConnectivityState.SHUTDOWN -> {
                                Log.w(TAG, "Channel entered failure state: $currentState")
                                triggerReconnect("Channel state: $currentState")
                            }
                            ConnectivityState.IDLE -> {
                                // Try to reconnect if idle
                                Log.d(TAG, "Channel idle, requesting connection")
                                ch.getState(true) // Request connection
                            }
                            else -> { /* READY, CONNECTING - OK */ }
                        }
                    }
                } catch (e: Exception) {
                    Log.e(TAG, "Error checking channel state: ${e.message}")
                }
            }
        }
    }

    /**
     * Trigger reconnection with exponential backoff
     */
    private fun triggerReconnect(reason: String) {
        Log.d(TAG, "[IDLE_CRASH_DEBUG] triggerReconnect called: $reason")
        Log.d(TAG, "[IDLE_CRASH_DEBUG] shouldReconnect=${shouldReconnect.get()}, isReconnecting=${isReconnecting.get()}")

        if (!shouldReconnect.get() || isReconnecting.getAndSet(true)) {
            Log.d(TAG, "[IDLE_CRASH_DEBUG] triggerReconnect skipped (already reconnecting or disabled)")
            return
        }

        val host = currentHost
        val port = currentPort
        if (host == null || port == null) {
            Log.d(TAG, "[IDLE_CRASH_DEBUG] triggerReconnect skipped (no host/port)")
            isReconnecting.set(false)
            return
        }

        Log.d(TAG, "[IDLE_CRASH_DEBUG] Triggering reconnect to $host:$port")

        // Emit disconnected event
        _connectionEvents.tryEmit(GrpcConnectionEvent.Disconnected(reason))

        reconnectJob?.cancel()
        reconnectJob = scope.launch {
            val attempt = reconnectAttempts.incrementAndGet()

            if (attempt > reconnectConfig.maxRetries) {
                Log.e(TAG, "Max reconnect attempts reached")
                _connectionState.value = GrpcConnectionState.Failed("Max retries exceeded")
                _connectionEvents.tryEmit(GrpcConnectionEvent.ReconnectFailed("Max retries exceeded"))
                isReconnecting.set(false)
                return@launch
            }

            _connectionState.value = GrpcConnectionState.Reconnecting(attempt, reconnectConfig.maxRetries)

            // Calculate delay with exponential backoff
            val delay = min(
                (reconnectConfig.initialDelayMs * reconnectConfig.backoffMultiplier.pow(attempt - 1.0)).toLong(),
                reconnectConfig.maxDelayMs
            )

            Log.d(TAG, "Reconnecting in ${delay}ms (attempt $attempt/${reconnectConfig.maxRetries})")
            delay(delay)

            // 注意：不要在这里重置 isReconnecting！
            // isReconnecting 会在 waitForConnection() 的 READY 分支中被重置
            // 这样 waitForConnection() 才能知道这是重连而非初次连接
            Log.d(TAG, ">>> Starting reconnect, isReconnecting=$isReconnecting (should be true)")
            doConnect(host, port)
        }
    }

    /**
     * Start heartbeat mechanism
     */
    private fun startHeartbeat() {
        stopHeartbeat()

        val partyId = registeredPartyId ?: return

        heartbeatJob = scope.launch {
            while (isActive && _connectionState.value == GrpcConnectionState.Connected) {
                delay(HEARTBEAT_INTERVAL_MS)

                try {
                    val request = HeartbeatRequest.newBuilder()
                        .setPartyId(partyId)
                        .setTimestamp(System.currentTimeMillis())
                        .build()

                    val response = withContext(Dispatchers.IO) {
                        getStubWithDeadline()?.heartbeat(request)
                    }

                    if (response != null) {
                        heartbeatFailCount.set(0)
                        val pending = response.pendingMessages
                        Log.d(TAG, "Heartbeat OK, pending messages: $pending")

                        // Notify if there are pending messages (may have missed events)
                        if (pending > 0) {
                            Log.w(TAG, "Has $pending pending messages - may have missed events")
                            _connectionEvents.tryEmit(GrpcConnectionEvent.PendingMessages(pending))
                        }
                    } else {
                        handleHeartbeatFailure("No response")
                    }
                } catch (e: Exception) {
                    handleHeartbeatFailure(e.message ?: "Unknown error")
                }
            }
        }
    }

    private fun handleHeartbeatFailure(reason: String) {
        val fails = heartbeatFailCount.incrementAndGet()
        Log.w(TAG, "[IDLE_CRASH_DEBUG] Heartbeat failed ($fails/$MAX_HEARTBEAT_FAILS): $reason")
        Log.w(TAG, "[IDLE_CRASH_DEBUG] Connection state: ${_connectionState.value}")
        Log.w(TAG, "[IDLE_CRASH_DEBUG] Channel state: ${channel?.getState(false)}")

        if (fails >= MAX_HEARTBEAT_FAILS) {
            Log.e(TAG, "[IDLE_CRASH_DEBUG] Too many heartbeat failures, triggering reconnect")
            triggerReconnect("Heartbeat failed")
        }
    }

    private fun stopHeartbeat() {
        Log.d(TAG, "[IDLE_CRASH_DEBUG] stopHeartbeat called")
        heartbeatJob?.cancel()
        heartbeatJob = null
        heartbeatFailCount.set(0)
    }

    /**
     * Re-register party after reconnection
     */
    private suspend fun reRegisterIfNeeded() {
        val partyId = registeredPartyId
        val role = registeredPartyRole

        if (partyId != null && role != null) {
            Log.d(TAG, "Re-registering party: $partyId")
            try {
                registerPartyInternal(partyId, role, "1.0.0")
            } catch (e: Exception) {
                Log.e(TAG, "Re-registration failed: ${e.message}")
            }
        }

        // Also re-register the signing partyId if active (for imported/restored shares)
        val signingId = registeredSigningPartyId
        if (signingId != null && signingId != partyId) {
            Log.d(TAG, "Re-registering signing party: $signingId")
            try {
                registerPartyInternal(signingId, "temporary", "1.0.0")
            } catch (e: Exception) {
                Log.e(TAG, "Re-registration of signing party failed: ${e.message}")
            }
        }
    }

    /**
     * Wait for channel to be ready (with timeout)
     */
    private suspend fun waitForChannelReady(timeoutMs: Long = 2000): Boolean {
        val ch = channel ?: return false
        val startTime = System.currentTimeMillis()

        while (System.currentTimeMillis() - startTime < timeoutMs) {
            if (ch.getState(false) == ConnectivityState.READY) {
                return true
            }
            delay(50)
        }
        return false
    }

    /**
     * Re-subscribe to streams after reconnection
     * Notifies the repository layer to re-establish message/event subscriptions
     */
    private fun reSubscribeStreams() {
        Log.d(TAG, "[IDLE_CRASH_DEBUG] reSubscribeStreams called")
        val needsResubscribe = eventStreamSubscribed.get() || activeMessageSubscription != null

        if (needsResubscribe) {
            Log.d(TAG, "[IDLE_CRASH_DEBUG] Triggering stream re-subscription callback")
            Log.d(TAG, "[IDLE_CRASH_DEBUG]   - Event stream: ${eventStreamSubscribed.get()}, partyId: $eventStreamPartyId")
            Log.d(TAG, "[IDLE_CRASH_DEBUG]   - Message stream: ${activeMessageSubscription?.sessionId}")

            // Notify repository to re-establish streams
            scope.launch {
                Log.d(TAG, "[IDLE_CRASH_DEBUG] Waiting for channel to be ready...")
                // Wait for channel to be fully ready instead of fixed delay
                if (waitForChannelReady()) {
                    Log.d(TAG, "[IDLE_CRASH_DEBUG] Channel ready, invoking reconnect callback")
                    try {
                        onReconnectedCallback?.invoke()
                        Log.d(TAG, "[IDLE_CRASH_DEBUG] Reconnect callback completed")
                        // Emit reconnected event
                        _connectionEvents.tryEmit(GrpcConnectionEvent.Reconnected)
                    } catch (e: Exception) {
                        Log.e(TAG, "[IDLE_CRASH_DEBUG] Reconnect callback failed: ${e.message}")
                        Log.e(TAG, "[IDLE_CRASH_DEBUG] Stack trace: ${e.stackTraceToString()}")
                        // Don't let callback failure affect the connection state
                    }
                } else {
                    Log.w(TAG, "Channel not ready for stream re-subscription, skipping")
                }
            }
        } else {
            // No streams to restore, still emit reconnected event
            _connectionEvents.tryEmit(GrpcConnectionEvent.Reconnected)
        }
    }

    /**
     * Set callback for reconnection events
     * Called when connection is restored and streams need to be re-established
     */
    fun setOnReconnectedCallback(callback: () -> Unit) {
        onReconnectedCallback = callback
    }

    /**
     * Get active message subscription info (for recovery)
     */
    fun getActiveMessageSubscription(): Pair<String, String>? {
        return activeMessageSubscription?.let { Pair(it.sessionId, it.partyId) }
    }

    /**
     * Check if event stream was subscribed (for recovery)
     */
    fun wasEventStreamSubscribed(): Boolean = eventStreamSubscribed.get()

    /**
     * Get event stream party ID (for recovery)
     */
    fun getEventStreamPartyId(): String? = eventStreamPartyId

    /**
     * Check if connected
     */
    fun isConnected(): Boolean = _connectionState.value == GrpcConnectionState.Connected

    /**
     * Get stub with per-call deadline (fixes timeout issue)
     */
    private fun getStubWithDeadline(): MessageRouterGrpc.MessageRouterBlockingStub? {
        return stub?.withDeadlineAfter(REQUEST_TIMEOUT_SECONDS, TimeUnit.SECONDS)
    }

    /**
     * Execute an RPC call with retry logic for transient failures
     */
    private suspend fun <T> withRetry(
        operation: String,
        maxRetries: Int = MAX_REQUEST_RETRIES,
        block: suspend () -> T
    ): T {
        var lastException: Exception? = null

        repeat(maxRetries) { attempt ->
            try {
                return block()
            } catch (e: Exception) {
                lastException = e
                val isRetryable = isRetryableError(e)

                if (attempt < maxRetries - 1 && isRetryable) {
                    val delay = REQUEST_RETRY_DELAY_MS * (attempt + 1)
                    Log.w(TAG, "$operation failed (attempt ${attempt + 1}/$maxRetries), retrying in ${delay}ms: ${e.message}")
                    delay(delay)
                } else {
                    Log.e(TAG, "$operation failed after ${attempt + 1} attempts: ${e.message}")
                    throw e
                }
            }
        }

        throw lastException ?: Exception("$operation failed after $maxRetries attempts")
    }

    /**
     * Check if an error is retryable (transient network issues)
     */
    private fun isRetryableError(e: Exception): Boolean {
        val message = e.message.orEmpty().lowercase()
        return message.contains("unavailable") ||
                message.contains("deadline exceeded") ||
                message.contains("connection") ||
                message.contains("timeout") ||
                message.contains("reset") ||
                message.contains("broken pipe")
    }

    /**
     * Register party with the router
     */
    suspend fun registerParty(
        partyId: String,
        partyRole: String = "temporary",
        version: String = "1.0.0"
    ): Result<Boolean> = withContext(Dispatchers.IO) {
        // 必须等待 channel READY 后才能注册
        Log.d(TAG, "registerParty: Waiting for channel READY...")
        val isReady = waitForChannelReady(CONNECTION_TIMEOUT_SECONDS * 1000)
        if (!isReady) {
            Log.e(TAG, "registerParty: Channel not ready after timeout")
            return@withContext Result.failure(Exception("Channel not ready"))
        }
        Log.d(TAG, "registerParty: Channel is READY, proceeding with registration")

        // Save for re-registration
        registeredPartyId = partyId
        registeredPartyRole = partyRole

        registerPartyInternal(partyId, partyRole, version)
    }

    private suspend fun registerPartyInternal(
        partyId: String,
        partyRole: String,
        version: String
    ): Result<Boolean> = withContext(Dispatchers.IO) {
        try {
            withRetry("RegisterParty") {
                val request = RegisterPartyRequest.newBuilder()
                    .setPartyId(partyId)
                    .setPartyRole(partyRole)
                    .setVersion(version)
                    .build()

                val response = getStubWithDeadline()?.registerParty(request)
                if (response?.success == true) {
                    Log.d(TAG, "Party registered: $partyId")
                    startHeartbeat()
                    true
                } else {
                    throw Exception("Registration failed")
                }
            }
            Result.success(true)
        } catch (e: Exception) {
            Log.e(TAG, "Register party failed: ${e.message}")
            Result.failure(e)
        }
    }

    /**
     * Register an additional signing partyId with the message-router.
     * Used when signing with imported/restored shares where the signing partyId
     * differs from the device's own partyId. Does not overwrite the device registration.
     */
    suspend fun registerSigningParty(signingPartyId: String): Result<Boolean> = withContext(Dispatchers.IO) {
        if (signingPartyId == registeredPartyId) {
            // Same as device partyId, already registered
            return@withContext Result.success(true)
        }
        Log.d(TAG, "Registering signing partyId: $signingPartyId (device partyId: $registeredPartyId)")
        registeredSigningPartyId = signingPartyId
        registerPartyInternal(signingPartyId, "temporary", "1.0.0")
    }

    /**
     * Clear the additional signing party registration.
     * Called when signing completes or fails.
     */
    fun clearSigningPartyRegistration() {
        registeredSigningPartyId = null
    }

    /**
     * Join a session
     */
    suspend fun joinSession(
        sessionId: String,
        partyId: String,
        joinToken: String
    ): Result<JoinSessionData> = withContext(Dispatchers.IO) {
        try {
            val data = withRetry("JoinSession") {
                // Match Electron behavior: don't send device_info
                val request = JoinSessionRequest.newBuilder()
                    .setSessionId(sessionId)
                    .setPartyId(partyId)
                    .setJoinToken(joinToken)
                    .build()

                val response = getStubWithDeadline()?.joinSession(request)
                if (response?.success == true) {
                    val sessionInfo = response.sessionInfo
                    val participants = response.otherPartiesList.map { party ->
                        Participant(party.partyId, party.partyIndex)
                    }

                    JoinSessionData(
                        sessionId = sessionInfo.sessionId,
                        sessionType = sessionInfo.sessionType,
                        thresholdN = sessionInfo.thresholdN,
                        thresholdT = sessionInfo.thresholdT,
                        partyIndex = response.partyIndex,
                        participants = participants,
                        messageHash = if (sessionInfo.messageHash.isEmpty) null
                        else Base64.encodeToString(sessionInfo.messageHash.toByteArray(), Base64.NO_WRAP),
                        sessionStatus = if (sessionInfo.status.isNullOrEmpty()) null else sessionInfo.status
                    )
                } else {
                    throw Exception("Failed to join session")
                }
            }
            Result.success(data)
        } catch (e: Exception) {
            Log.e(TAG, "Join session failed: ${e.message}")
            Result.failure(e)
        }
    }

    /**
     * Mark party as ready
     */
    suspend fun markPartyReady(
        sessionId: String,
        partyId: String
    ): Result<Boolean> = withContext(Dispatchers.IO) {
        try {
            val request = MarkPartyReadyRequest.newBuilder()
                .setSessionId(sessionId)
                .setPartyId(partyId)
                .build()

            val response = getStubWithDeadline()?.markPartyReady(request)
            Result.success(response?.allReady ?: false)
        } catch (e: Exception) {
            Log.e(TAG, "Mark party ready failed: ${e.message}")
            Result.failure(e)
        }
    }

    /**
     * Route a message to other parties
     */
    suspend fun routeMessage(
        sessionId: String,
        fromParty: String,
        toParties: List<String>,
        roundNumber: Int,
        messageType: String,
        payload: ByteArray
    ): Result<String> = withContext(Dispatchers.IO) {
        try {
            val messageId = withRetry("RouteMessage") {
                val request = RouteMessageRequest.newBuilder()
                    .setSessionId(sessionId)
                    .setFromParty(fromParty)
                    .addAllToParties(toParties)
                    .setRoundNumber(roundNumber)
                    .setMessageType(messageType)
                    .setPayload(com.google.protobuf.ByteString.copyFrom(payload))
                    .build()

                val response = getStubWithDeadline()?.routeMessage(request)
                if (response?.success == true) {
                    response.messageId
                } else {
                    throw Exception("Failed to route message")
                }
            }
            Result.success(messageId)
        } catch (e: Exception) {
            Log.e(TAG, "Route message failed: ${e.message}")
            Result.failure(e)
        }
    }

    /**
     * Subscribe to messages for a party (with auto-recovery)
     */
    fun subscribeMessages(sessionId: String, partyId: String): Flow<IncomingMessage> = callbackFlow {
        // Save subscription for recovery
        activeMessageSubscription = MessageSubscription(sessionId, partyId)

        // Capture current stream version to detect stale callbacks
        val streamVersion = messageStreamVersion.incrementAndGet()

        val request = SubscribeMessagesRequest.newBuilder()
            .setSessionId(sessionId)
            .setPartyId(partyId)
            .build()

        val observer = object : StreamObserver<MPCMessage> {
            override fun onNext(message: MPCMessage) {
                // Ignore events from stale streams
                if (messageStreamVersion.get() != streamVersion) {
                    Log.d(TAG, "Ignoring message from stale stream (v$streamVersion, current v${messageStreamVersion.get()})")
                    return
                }

                val incoming = IncomingMessage(
                    messageId = message.messageId,
                    fromParty = message.fromParty,
                    isBroadcast = message.isBroadcast,
                    roundNumber = message.roundNumber,
                    payload = Base64.encodeToString(message.payload.toByteArray(), Base64.NO_WRAP)
                )
                trySend(incoming)
            }

            override fun onError(t: Throwable) {
                Log.e(TAG, "Message stream error: ${t.message}")

                // Ignore events from stale streams - close without exception to avoid crash
                if (messageStreamVersion.get() != streamVersion) {
                    Log.d(TAG, "Ignoring error from stale message stream")
                    close()
                    return
                }

                // Don't trigger reconnect for CANCELLED or channel shutdown errors
                val errorMessage = t.message.orEmpty()
                if (!errorMessage.contains("CANCELLED") && !errorMessage.contains("shutdownNow")) {
                    triggerReconnect("Message stream error: ${t.message}")
                }
                close(t)
            }

            override fun onCompleted() {
                Log.d(TAG, "Message stream completed")

                // Ignore events from stale streams
                if (messageStreamVersion.get() != streamVersion) {
                    Log.d(TAG, "Ignoring completion from stale message stream")
                    close()
                    return
                }

                // Stream ended unexpectedly - trigger reconnect if we should still be subscribed
                if (activeMessageSubscription != null && shouldReconnect.get()) {
                    Log.w(TAG, "Message stream ended unexpectedly, triggering reconnect")
                    triggerReconnect("Message stream ended")
                }
                close()
            }
        }

        asyncStub?.subscribeMessages(request, observer)

        awaitClose {
            activeMessageSubscription = null
        }
    }

    /**
     * Subscribe to session events (with auto-recovery)
     */
    fun subscribeSessionEvents(partyId: String): Flow<SessionEventData> = callbackFlow {
        // Save subscription for recovery
        eventStreamSubscribed.set(true)
        eventStreamPartyId = partyId

        // Capture current stream version to detect stale callbacks
        val streamVersion = eventStreamVersion.incrementAndGet()

        val request = SubscribeSessionEventsRequest.newBuilder()
            .setPartyId(partyId)
            .build()

        val observer = object : StreamObserver<SessionEvent> {
            override fun onNext(event: SessionEvent) {
                // Ignore events from stale streams
                if (eventStreamVersion.get() != streamVersion) {
                    Log.d(TAG, "Ignoring event from stale stream (v$streamVersion, current v${eventStreamVersion.get()})")
                    return
                }

                val eventData = SessionEventData(
                    eventId = event.eventId,
                    eventType = event.eventType,
                    sessionId = event.sessionId,
                    thresholdN = event.thresholdN,
                    thresholdT = event.thresholdT,
                    selectedParties = event.selectedPartiesList,
                    joinTokens = event.joinTokensMap,
                    messageHash = if (event.messageHash.isEmpty) null
                    else Base64.encodeToString(event.messageHash.toByteArray(), Base64.NO_WRAP)
                )
                trySend(eventData)
            }

            override fun onError(t: Throwable) {
                Log.e(TAG, "Session event stream error: ${t.message}")

                // Ignore events from stale streams - close without exception to avoid crash
                if (eventStreamVersion.get() != streamVersion) {
                    Log.d(TAG, "Ignoring error from stale event stream")
                    close()
                    return
                }

                // Don't trigger reconnect for CANCELLED or channel shutdown errors
                val errorMessage = t.message.orEmpty()
                if (!errorMessage.contains("CANCELLED") && !errorMessage.contains("shutdownNow")) {
                    triggerReconnect("Event stream error: ${t.message}")
                }
                close(t)
            }

            override fun onCompleted() {
                Log.d(TAG, "Session event stream completed")

                // Ignore events from stale streams
                if (eventStreamVersion.get() != streamVersion) {
                    Log.d(TAG, "Ignoring completion from stale event stream")
                    close()
                    return
                }

                // Stream ended unexpectedly - trigger reconnect if we should still be subscribed
                if (eventStreamSubscribed.get() && shouldReconnect.get()) {
                    Log.w(TAG, "Event stream ended unexpectedly, triggering reconnect")
                    triggerReconnect("Event stream ended")
                }
                close()
            }
        }

        val currentAsyncStub = asyncStub
        if (currentAsyncStub == null) {
            Log.e(TAG, "subscribeSessionEvents: asyncStub is null! Cannot subscribe.")
            close(Exception("gRPC not connected - asyncStub is null"))
            return@callbackFlow
        }

        Log.d(TAG, "subscribeSessionEvents: Starting subscription for partyId=$partyId")
        currentAsyncStub.subscribeSessionEvents(request, observer)
        Log.d(TAG, "subscribeSessionEvents: Subscription request sent")

        awaitClose {
            Log.d(TAG, "subscribeSessionEvents: Flow closed for partyId=$partyId")
            eventStreamSubscribed.set(false)
            eventStreamPartyId = null
        }
    }

    /**
     * Unsubscribe from session events
     */
    fun unsubscribeSessionEvents() {
        eventStreamSubscribed.set(false)
        eventStreamPartyId = null
    }

    /**
     * Unsubscribe from messages
     */
    fun unsubscribeMessages() {
        activeMessageSubscription = null
    }

    /**
     * Report completion
     */
    suspend fun reportCompletion(
        sessionId: String,
        partyId: String,
        publicKey: ByteArray? = null,
        signature: ByteArray? = null
    ): Result<Boolean> = withContext(Dispatchers.IO) {
        try {
            val builder = ReportCompletionRequest.newBuilder()
                .setSessionId(sessionId)
                .setPartyId(partyId)

            publicKey?.let {
                builder.setPublicKey(com.google.protobuf.ByteString.copyFrom(it))
            }
            signature?.let {
                builder.setSignature(com.google.protobuf.ByteString.copyFrom(it))
            }

            val response = getStubWithDeadline()?.reportCompletion(builder.build())
            Result.success(response?.allCompleted ?: false)
        } catch (e: Exception) {
            Log.e(TAG, "Report completion failed: ${e.message}")
            Result.failure(e)
        }
    }

    /**
     * Send heartbeat (manual call for testing)
     */
    suspend fun heartbeat(partyId: String): Result<Int> = withContext(Dispatchers.IO) {
        try {
            val request = HeartbeatRequest.newBuilder()
                .setPartyId(partyId)
                .setTimestamp(System.currentTimeMillis())
                .build()

            val response = getStubWithDeadline()?.heartbeat(request)
            Result.success(response?.pendingMessages ?: 0)
        } catch (e: Exception) {
            Log.e(TAG, "Heartbeat failed: ${e.message}")
            Result.failure(e)
        }
    }

    /**
     * Force reconnect (for manual recovery)
     */
    fun forceReconnect() {
        if (currentHost != null && currentPort != null) {
            reconnectAttempts.set(0)
            triggerReconnect("Manual reconnect requested")
        }
    }

    /**
     * Get current connection parameters
     */
    fun getConnectionInfo(): Pair<String, Int>? {
        val host = currentHost
        val port = currentPort
        return if (host != null && port != null) Pair(host, port) else null
    }
}

/**
 * Message subscription info
 */
private data class MessageSubscription(
    val sessionId: String,
    val partyId: String
)

/**
 * Data class for join session response
 */
data class JoinSessionData(
    val sessionId: String,
    val sessionType: String,
    val thresholdN: Int,
    val thresholdT: Int,
    val partyIndex: Int,
    val participants: List<Participant>,
    val messageHash: String?,
    val sessionStatus: String? = null
)

/**
 * Data class for incoming MPC message
 */
data class IncomingMessage(
    val messageId: String,
    val fromParty: String,
    val isBroadcast: Boolean,
    val roundNumber: Int,
    val payload: String  // base64 encoded
)

/**
 * Data class for session event
 */
data class SessionEventData(
    val eventId: String,
    val eventType: String,
    val sessionId: String,
    val thresholdN: Int,
    val thresholdT: Int,
    val selectedParties: List<String>,
    val joinTokens: Map<String, String>,
    val messageHash: String?
)
