soniox-mobile-to-mac/android/app/src/main/java/com/soniox/remotemic/StreamDictationClient.kt

package com.soniox.remotemic

import android.annotation.SuppressLint
import android.media.AudioFormat
import android.media.AudioRecord
import android.media.MediaRecorder
import android.os.Handler
import android.os.Looper
import android.util.Log
import okhttp3.*
import okio.ByteString.Companion.toByteString
import org.json.JSONObject
import java.util.UUID
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean
import kotlin.math.sqrt

data class SyncStatePayload(
    val source: String,
    val app: String,
    val text: String,
    val cursor: Int,
    val selection: Int,
    val revision: Long,
    val timestamp: Double
)

class StreamDictationClient(
    private val host: String,
    private val onConnectionStateChanged: (Boolean) -> Unit,
    private val onSyncStateReceived: (SyncStatePayload) -> Unit,
    private val onPartialSpeechText: (String) -> Unit,
    private val onAudioLevel: (Float) -> Unit,
    private val onSpeechCompleted: (String) -> Unit,
    private val onError: (String) -> Unit
) {
    private val tag = "StreamDictationClient"
    private val sampleRate = 16000
    private val channelConfig = AudioFormat.CHANNEL_IN_MONO
    private val audioFormat = AudioFormat.ENCODING_PCM_16BIT

    private var audioRecord: AudioRecord? = null
    private var recordingThread: Thread? = null
    private val isRecording = AtomicBoolean(false)
    private val mainHandler = Handler(Looper.getMainLooper())

    private val okHttpClient = OkHttpClient.Builder()
        .connectTimeout(12, TimeUnit.SECONDS)
        .readTimeout(0, TimeUnit.MILLISECONDS) // Persistent WebSocket
        .writeTimeout(8, TimeUnit.SECONDS)
        .pingInterval(8, TimeUnit.SECONDS)
        .retryOnConnectionFailure(true)
        .build()

    private var activeWebSocket: WebSocket? = null
    private val isConnected = AtomicBoolean(false)
    private val isConnecting = AtomicBoolean(false)
    private var currentSessionId: String = ""
    private var isSessionActive = AtomicBoolean(false)
    private var retryAttempt = 0

    init {
        connectWebSocket()
    }

    @Synchronized
    fun connectWebSocket() {
        if (isConnected.get() && activeWebSocket != null) return
        if (isConnecting.getAndSet(true)) return

        val cleanHost = host.removePrefix("http://").removePrefix("https://").removePrefix("ws://").removePrefix("wss://")
        val wsUrl = "ws://$cleanHost/ws/stream"
        AppLogger.log(tag, "اتصال به سوکت همگام‌سازی بلادرنگ: $wsUrl")

        val req = Request.Builder()
            .url(wsUrl)
            .header("User-Agent", "SonioxAndroidRemote/5.2")
            .build()

        activeWebSocket = okHttpClient.newWebSocket(req, object : WebSocketListener() {
            override fun onOpen(webSocket: WebSocket, response: Response) {
                AppLogger.log(tag, "🟢 سوکت همگام‌سازی بلادرنگ متصل شد")
                isConnected.set(true)
                isConnecting.set(false)
                retryAttempt = 0
                mainHandler.post { onConnectionStateChanged(true) }
            }

            override fun onMessage(webSocket: WebSocket, text: String) {
                try {
                    val json = JSONObject(text)
                    val type = json.optString("type")
                    val source = json.optString("source", "")

                    if (type == "sync_state" || type == "mac_input_state") {
                        val app = json.optString("app", "Mac")
                        val txt = json.optString("text", "")
                        val cursor = json.optInt("cursor", txt.length)
                        val sel = json.optInt("selection", 0)
                        val rev = json.optLong("revision", 0L)
                        val ts = json.optDouble("timestamp", System.currentTimeMillis() / 1000.0)
                        
                        val payload = SyncStatePayload(source, app, txt, cursor, sel, rev, ts)
                        mainHandler.post { onSyncStateReceived(payload) }
                        return
                    }

                    val sid = json.optString("session_id")
                    when (type) {
                        "live", "partial" -> {
                            if (sid.isEmpty() || sid == currentSessionId || !isSessionActive.get()) {
                                val liveText = json.optString("text")
                                mainHandler.post { onPartialSpeechText(liveText) }
                            }
                        }
                        "final" -> {
                            val finalText = json.optString("text")
                            isSessionActive.set(false)
                            AppLogger.log(tag, "⚡ صوت پردازش شد: '$finalText'")
                            mainHandler.post { onSpeechCompleted(finalText) }
                        }
                        "error" -> {
                            val msg = json.optString("message", "خطای سرور")
                            AppLogger.log(tag, "❌ خطای سرور: $msg")
                            mainHandler.post { onError(msg) }
                        }
                    }
                } catch (e: Exception) {
                    Log.e(tag, "Message parse error", e)
                }
            }

            override fun onFailure(webSocket: WebSocket, t: Throwable, response: Response?) {
                AppLogger.log(tag, "🔴 قطع سوکت: ${t.message}. تلاش مجدد...")
                handleDisconnect()
            }

            override fun onClosed(webSocket: WebSocket, code: Int, reason: String) {
                AppLogger.log(tag, "سوکت بسته شد ($code). تلاش مجدد...")
                handleDisconnect()
            }

            private fun handleDisconnect() {
                isConnected.set(false)
                isConnecting.set(false)
                activeWebSocket = null
                mainHandler.post { onConnectionStateChanged(false) }
                
                retryAttempt++
                val delayMs = minOf(800L * (1L shl minOf(retryAttempt, 3)), 5000L)
                mainHandler.postDelayed({ connectWebSocket() }, delayMs)
            }
        })
    }

    /**
     * Sends speech chunk directly for cursor insertion on Mac (Cmd+V)
     */
    fun sendSpeechInsert(speechText: String, cursor: Int) {
        if (!isConnected.get() || activeWebSocket == null) {
            connectWebSocket()
        }
        val payload = JSONObject().apply {
            put("type", "insert_speech")
            put("action", "insert_speech")
            put("source", "android")
            put("text", speechText)
            put("cursor", cursor)
            put("timestamp", System.currentTimeMillis() / 1000.0)
        }.toString()
        activeWebSocket?.send(payload)
    }

    /**
     * Broadcasts phone's full updated text & cursor to Mac for replace (Cmd+A + Cmd+V)
     */
    fun sendForceReplace(text: String, cursor: Int, selection: Int = 0) {
        if (!isConnected.get() || activeWebSocket == null) {
            connectWebSocket()
        }
        val payload = JSONObject().apply {
            put("type", "update_input")
            put("action", "update_input")
            put("source", "android")
            put("text", text)
            put("cursor", cursor)
            put("selection", selection)
            put("timestamp", System.currentTimeMillis() / 1000.0)
        }.toString()
        activeWebSocket?.send(payload)
    }

    /**
     * Broadcasts phone's updated text & cursor to Mac
     */
    fun sendLocalSyncState(text: String, cursor: Int, selection: Int = 0) {
        if (!isConnected.get() || activeWebSocket == null) {
            connectWebSocket()
        }
        val payload = JSONObject().apply {
            put("type", "sync_state")
            put("source", "android")
            put("text", text)
            put("cursor", cursor)
            put("selection", selection)
            put("timestamp", System.currentTimeMillis() / 1000.0)
        }.toString()
        activeWebSocket?.send(payload)
    }

    @SuppressLint("MissingPermission")
    fun startRecording(cursorPos: Int): String {
        if (isRecording.get()) return currentSessionId

        currentSessionId = "sess_${System.currentTimeMillis()}_${UUID.randomUUID().toString().take(6)}"
        isRecording.set(true)
        isSessionActive.set(true)

        if (!isConnected.get() || activeWebSocket == null) {
            connectWebSocket()
        }

        val startFrame = JSONObject().apply {
            put("type", "start")
            put("session_id", currentSessionId)
            put("cursor_pos", cursorPos)
        }.toString()
        activeWebSocket?.send(startFrame)
        AppLogger.log(tag, "🎙️ شروع استریم گفتار در موقعیت نشانگر $cursorPos...")

        val minBufferSize = AudioRecord.getMinBufferSize(sampleRate, channelConfig, audioFormat)
        val bufferSize = maxOf(minBufferSize, 2048)

        try {
            audioRecord = AudioRecord(
                MediaRecorder.AudioSource.MIC,
                sampleRate,
                channelConfig,
                audioFormat,
                bufferSize
            )

            if (audioRecord?.state != AudioRecord.STATE_INITIALIZED) {
                AppLogger.log(tag, "❌ سخت‌افزار میکروفون راه‌اندازی نشد")
                onError("خطا در راه‌اندازی سخت‌افزار میکروفون")
                isRecording.set(false)
                isSessionActive.set(false)
                return currentSessionId
            }

            audioRecord?.startRecording()

            recordingThread = Thread {
                val chunk = ByteArray(2048)
                val shortBuffer = ShortArray(1024)

                while (isRecording.get()) {
                    val bytesRead = audioRecord?.read(chunk, 0, chunk.size) ?: -1
                    if (bytesRead > 0) {
                        val slice = if (bytesRead == chunk.size) chunk.clone() else chunk.copyOf(bytesRead)
                        activeWebSocket?.send(slice.toByteString())

                        var sum = 0.0
                        val samplesCount = bytesRead / 2
                        for (i in 0 until samplesCount) {
                            val sample = (slice[i * 2].toInt() and 0xFF) or (slice[i * 2 + 1].toInt() shl 8)
                            shortBuffer[i] = sample.toShort()
                            sum += (shortBuffer[i] * shortBuffer[i]).toDouble()
                        }
                        val rms = sqrt(sum / samplesCount) / 32768.0
                        val level = minOf(maxOf(rms * 4.5, 0.0), 1.0).toFloat()
                        mainHandler.post { onAudioLevel(level) }
                    }
                }
            }.apply {
                priority = Thread.MAX_PRIORITY
                start()
            }

        } catch (e: Exception) {
            AppLogger.log(tag, "خطای ضبط: ${e.message}")
            onError("خطای میکروفون: ${e.localizedMessage}")
            isRecording.set(false)
            isSessionActive.set(false)
        }

        return currentSessionId
    }

    fun stopRecording(cursorPos: Int) {
        if (!isRecording.get()) return
        isRecording.set(false)

        try {
            audioRecord?.stop()
            audioRecord?.release()
            audioRecord = null
            recordingThread?.join(150)
            recordingThread = null
        } catch (e: Exception) {
            Log.e(tag, "Error releasing audio hardware", e)
        }

        val stopFrame = JSONObject().apply {
            put("type", "stop")
            put("session_id", currentSessionId)
            put("cursor_pos", cursorPos)
        }.toString()
        activeWebSocket?.send(stopFrame)
        AppLogger.log(tag, "⏹️ پایان ضبط. دریافت متن نهایی...")
    }

    fun release() {
        isRecording.set(false)
        try {
            audioRecord?.release()
            activeWebSocket?.close(1000, "Client Shutdown")
        } catch (e: Exception) {}
    }
}