soniox-mobile-to-mac/mac/src/RelayClient.swift

import Foundation
import Cocoa

public final class RelayClient: NSObject, URLSessionWebSocketDelegate {
    public static let shared = RelayClient()
    
    // Connects through SSH local port forward 18999 -> Linux Server 8999
    private let primaryUrl = URL(string: "ws://127.0.0.1:18999/ws/mac")!
    
    private var webSocketTask: URLSessionWebSocketTask?
    private var urlSession: URLSession!
    
    private var isRunning = false
    public private(set) var isConnected = false
    private var reconnectTimer: Timer?
    private var pingTimer: Timer?
    private var monitorTimerSource: DispatchSourceTimer?
    private let monitorQueue = DispatchQueue(label: "com.soniox.macsync.monitor", qos: .userInteractive)
    
    public var onRemoteUpdateReceived: ((String, Int?, Bool) -> Void)?
    
    private var lastInjectedText: String = ""
    private var lastInjectedTime: Double = 0
    
    override private init() {
        super.init()
        let config = URLSessionConfiguration.default
        config.waitsForConnectivity = true
        config.timeoutIntervalForRequest = 30
        config.timeoutIntervalForResource = 300
        self.urlSession = URLSession(configuration: config, delegate: self, delegateQueue: .main)
    }
    
    public func start() {
        guard !isRunning else { return }
        isRunning = true
        print("RelayClient: Starting real-time bi-directional sync engine...")
        connect()
        startInputMonitoring()
    }
    
    public func stop() {
        isRunning = false
        reconnectTimer?.invalidate()
        reconnectTimer = nil
        pingTimer?.invalidate()
        pingTimer = nil
        monitorTimerSource?.cancel()
        monitorTimerSource = nil
        webSocketTask?.cancel(with: .goingAway, reason: nil)
        webSocketTask = nil
        isConnected = false
    }
    
    private func connect() {
        guard isRunning else { return }
        webSocketTask?.cancel()
        
        var request = URLRequest(url: primaryUrl)
        request.timeoutInterval = 6
        request.setValue("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) SonioxVoice/1.0", forHTTPHeaderField: "User-Agent")
        
        webSocketTask = urlSession.webSocketTask(with: request)
        webSocketTask?.resume()
        
        listenForMessages()
        startPingTimer()
    }
    
    private func listenForMessages() {
        webSocketTask?.receive { [weak self] result in
            guard let self = self, self.isRunning else { return }
            
            switch result {
            case .success(let message):
                self.isConnected = true
                switch message {
                case .string(let text):
                    self.handleMessageString(text)
                case .data(let data):
                    if let text = String(data: data, encoding: .utf8) {
                        self.handleMessageString(text)
                    }
                @unknown default:
                    break
                }
                self.listenForMessages()
                
            case .failure(let error):
                print("RelayClient: Connection error: \(error.localizedDescription)")
                self.isConnected = false
                self.scheduleReconnect()
            }
        }
    }
    
    private func handleMessageString(_ text: String) {
        guard let data = text.data(using: .utf8),
              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            return
        }
        
        let action = json["action"] as? String ?? json["type"] as? String ?? ""
        let source = json["source"] as? String ?? ""
        
        // Ignore echo messages originated by Mac itself
        if source == "mac" { return }
        
        let updateText = json["text"] as? String ?? json["insert_text"] as? String ?? ""
        let cursor = json["cursor"] as? Int ?? json["cursor_pos"] as? Int
        
        if action == "sync_state" || action == "phone_input_edit" || action == "update_input" || action == "force_replace" || action == "insert_speech" || action == "paste" {
            let now = Date().timeIntervalSince1970
            
            // Drop rapid duplicate transmissions within 150ms (unless force_replace)
            if updateText == lastInjectedText && (now - lastInjectedTime) < 0.15 && action != "force_replace" {
                return
            }
            lastInjectedText = updateText
            lastInjectedTime = now
            
            let isFullReplace = (action != "insert_speech" && action != "speech_insert")
            print("RelayClient: ⚡ Applying phone edit: '\(updateText.prefix(30))...' (replace: \(isFullReplace), len: \(updateText.count))")
            DispatchQueue.main.async {
                self.onRemoteUpdateReceived?(updateText, cursor, isFullReplace)
            }
        }
    }
    
    private func startInputMonitoring() {
        monitorTimerSource?.cancel()
        
        let timer = DispatchSource.makeTimerSource(queue: monitorQueue)
        // Poll every 50ms on dedicated user-interactive queue
        timer.schedule(deadline: .now(), repeating: .milliseconds(50))
        timer.setEventHandler { [weak self] in
            guard let self = self, self.isRunning, self.isConnected else { return }
            if let state = FocusedInputSync.shared.inspectCurrentState() {
                self.sendStateToRelay(state)
            }
        }
        timer.resume()
        self.monitorTimerSource = timer
    }
    
    private func sendStateToRelay(_ state: MacInputState) {
        guard let task = webSocketTask else { return }
        
        let payload: [String: Any] = [
            "type": "sync_state",
            "source": "mac",
            "app": state.app,
            "text": state.text,
            "cursor": state.cursor,
            "selection": state.selection,
            "revision": state.revision,
            "timestamp": state.timestamp
        ]
        
        if let data = try? JSONSerialization.data(withJSONObject: payload),
           let jsonStr = String(data: data, encoding: .utf8) {
            task.send(.string(jsonStr)) { error in
                if let error = error {
                    print("Error sending sync_state:", error)
                }
            }
        }
    }
    
    private func startPingTimer() {
        pingTimer?.invalidate()
        let timer = Timer(timeInterval: 10.0, repeats: true) { [weak self] _ in
            guard let self = self, self.isRunning else { return }
            self.webSocketTask?.send(.string("{\"type\":\"ping\"}")) { _ in }
        }
        RunLoop.main.add(timer, forMode: .common)
        self.pingTimer = timer
    }
    
    private func scheduleReconnect() {
        guard isRunning else { return }
        pingTimer?.invalidate()
        pingTimer = nil
        
        if reconnectTimer == nil || !(reconnectTimer?.isValid ?? false) {
            let timer = Timer(timeInterval: 2.0, repeats: false) { [weak self] _ in
                self?.reconnectTimer = nil
                self?.connect()
            }
            RunLoop.main.add(timer, forMode: .common)
            self.reconnectTimer = timer
        }
    }
    
    public func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask, didOpenWithProtocol protocol: String?) {
        print("RelayClient: 🟢 Persistent WebSocket Connected to Server Gateway!")
        self.isConnected = true
    }
    
    public func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask, didCloseWith closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?) {
        print("RelayClient: 🔴 WebSocket Closed (code: \(closeCode.rawValue))")
        self.isConnected = false
        self.scheduleReconnect()
    }
}