saltylab-ios/SulTee/SulTee/WebSocketClient.swift

import Foundation
import UIKit

/// Thin WebSocket wrapper around URLSessionWebSocketTask.
/// Reconnects automatically on disconnect.
final class WebSocketClient: NSObject, ObservableObject {

    enum ConnectionState {
        case disconnected, connecting, connected
    }

    @Published var state: ConnectionState = .disconnected

    private let url: URL
    private var session: URLSession!
    private var task: URLSessionWebSocketTask?
    private var shouldRun = false

    init(url: URL) {
        self.url = url
        super.init()
        self.session = URLSession(configuration: .default, delegate: self, delegateQueue: nil)
    }

    func connect() {
        shouldRun = true
        guard state == .disconnected else { return }
        openConnection()
    }

    func disconnect() {
        shouldRun = false
        task?.cancel(with: .normalClosure, reason: nil)
        task = nil
        DispatchQueue.main.async { self.state = .disconnected }
    }

    func send(_ message: [String: Any]) {
        guard state == .connected,
              let data = try? JSONSerialization.data(withJSONObject: message),
              let json = String(data: data, encoding: .utf8) else { return }

        task?.send(.string(json)) { error in
            if let error {
                print("[WebSocket] send error: \(error)")
            }
        }
    }

    // MARK: - Private

    private func openConnection() {
        DispatchQueue.main.async { self.state = .connecting }
        task = session.webSocketTask(with: url)
        task?.resume()
        scheduleReceive()
    }

    private func scheduleReceive() {
        task?.receive { [weak self] result in
            guard let self else { return }
            switch result {
            case .success(let message):
                self.handle(message)
                self.scheduleReceive()
            case .failure(let error):
                print("[WebSocket] receive error: \(error)")
                self.reconnectIfNeeded()
            }
        }
    }

    private func handle(_ message: URLSessionWebSocketTask.Message) {
        switch message {
        case .string(let text):
            guard let data = text.data(using: .utf8),
                  let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
                  let type = json["type"] as? String else { return }

            if type == "haptic" {
                DispatchQueue.main.async {
                    UIImpactFeedbackGenerator(style: .medium).impactOccurred()
                }
            }
        default:
            break
        }
    }

    private func reconnectIfNeeded() {
        DispatchQueue.main.async { self.state = .disconnected }
        guard shouldRun else { return }
        DispatchQueue.global().asyncAfter(deadline: .now() + 2) { [weak self] in
            self?.openConnection()
        }
    }
}

// MARK: - URLSessionWebSocketDelegate

extension WebSocketClient: URLSessionWebSocketDelegate {
    func urlSession(_ session: URLSession,
                    webSocketTask: URLSessionWebSocketTask,
                    didOpenWithProtocol protocol: String?) {
        DispatchQueue.main.async { self.state = .connected }
    }

    func urlSession(_ session: URLSession,
                    webSocketTask: URLSessionWebSocketTask,
                    didCloseWith closeCode: URLSessionWebSocketTask.CloseCode,
                    reason: Data?) {
        reconnectIfNeeded()
    }
}