jarvis/Jervis/Networking/WebSocketManager.swift

// WebSocketManager.swift
// Jarvis — WebSocket connection, event parsing, reconnect with backoff

import Foundation
import Combine

enum ConnectionState {
    case disconnected             
    case connecting
    case connected
    case reconnecting(attempt: Int)
}

@MainActor
final class WebSocketManager: ObservableObject {
    static let shared = WebSocketManager()

    @Published var connectionState: ConnectionState = .disconnected

    let events = PassthroughSubject<WSEvent, Never>()

    private var task: URLSessionWebSocketTask?
    private var pingTimer: Timer?
    private var reconnectTask: Task<Void, Never>?
    private var host: String?

    private let maxBackoff: TimeInterval = 60
    private var attempt = 0

    private init() {}

    func connect(host: String) {
        self.host = host
        self.attempt = 0
        openConnection()
    }

    func disconnect() {
        reconnectTask?.cancel()
        pingTimer?.invalidate()
        task?.cancel(with: .goingAway, reason: nil)
        task = nil
        connectionState = .disconnected
    }

    // MARK: - Private

    private func openConnection() {
        guard let host, let url = URL(string: "ws://\(host)/ws") else { return }
        connectionState = attempt == 0 ? .connecting : .reconnecting(attempt: attempt)

        let session = URLSession(configuration: .default)
        task = session.webSocketTask(with: url)
        task?.resume()

        connectionState = .connected
        attempt = 0
        startPing()
        listen()
    }

    private func listen() {
        // Capture task locally so the Sendable receive closure doesn't reference
        // the @MainActor-isolated property directly.
        guard let t = task else { return }
        t.receive { [weak self] result in
            switch result {
            case .failure:
                Task { @MainActor [weak self] in self?.handleDisconnect() }
            case .success(let message):
                Task { @MainActor [weak self] in
                    guard let self else { return }
                    self.handle(message: message)
                    self.listen()
                }
            }
        }
    }

    private func handle(message: URLSessionWebSocketTask.Message) {
        guard case .string(let text) = message,
              let data = text.data(using: .utf8) else { return }

        let decoder = JSONDecoder()
        decoder.dateDecodingStrategy = .iso8601
        decoder.keyDecodingStrategy = .convertFromSnakeCase

        guard let event = try? decoder.decode(WSEvent.self, from: data) else { return }

        if event.type == .ping {
            sendPong()
            return
        }

        events.send(event)
    }

    private func sendPong() {
        task?.send(.string(#"{"type":"pong"}"#)) { _ in }
    }

    private func startPing() {
        pingTimer?.invalidate()
        pingTimer = Timer.scheduledTimer(withTimeInterval: 30, repeats: true) { [weak self] _ in
            // Hop to MainActor before touching the isolated `task` property.
            Task { @MainActor [weak self] in
                guard let self else { return }
                self.task?.sendPing { [weak self] error in
                    if error != nil {
                        Task { @MainActor [weak self] in self?.handleDisconnect() }
                    }
                }
            }
        }
    }

    private func handleDisconnect() {
        pingTimer?.invalidate()
        task = nil
        connectionState = .disconnected
        scheduleReconnect()
    }

    private func scheduleReconnect() {
        attempt += 1
        let delay = min(pow(2.0, Double(attempt - 1)), maxBackoff)
        connectionState = .reconnecting(attempt: attempt)

        reconnectTask = Task { [weak self] in
            try? await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000))
            guard !Task.isCancelled else { return }
            await MainActor.run { self?.openConnection() }
        }
    }
}