adds lights UI

XIONControlPanel/Servers/HomeAssistantServer.swift

@@ -6,14 +6,13 @@
 //  Copyright © 2025 XION. All rights reserved.
 //
 
+import Combine
 import Foundation
 
 class HomeAssistantServer: Server
 {
     var connectionStatus: ConnectionStatus = .disconnected
     
-    private typealias Reply = (Data) -> Void
-    
     // Configure:
     private let authToken: String = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiI2ZTAwODI1N2Y4N2Q0OWIxYWRkN2ExNTBhOWRmNGZiZCIsImlhdCI6MTc1MDMwMDg5NiwiZXhwIjoyMDY1NjYwODk2fQ.KRf9GpRZdpW9Z_vkbL3sl74rgKS7eAyc8kO0a5jLTgg"
     
@@ -21,59 +20,60 @@ class HomeAssistantServer: Server
     private static let filterLabel = "control_panel"
     
     private let url: URL
-    private var websocketTask: URLSessionWebSocketTask? = nil
-    private var pendingMessages: [MessageID: Reply] = [:]
+    private let liaison: Liaison
     private var connectionContinuations: [CheckedContinuation<Bool, Never>] = []
+    private var eventTask: Task<(), any Error>? = nil
+    
+    typealias NoResult = [String: String]
     
     required init(_ url: URL) {
         self.url = url
+        
+        let websocketURL = url.appending(component: "websocket").websocket()
+        self.liaison = Liaison(websocketURL: websocketURL)
     }
     
     func connect(_ completion: @escaping ((any Error)?) -> Void) {
-        guard websocketTask == nil else { return }
-        
-        let url = url.appending(component: "websocket").websocket()
-        let websocketTask = URLSession.shared.webSocketTask(with: url)
-        websocketTask.resume()
-        
-        self.websocketTask = websocketTask
-        Task {
-            do {
-                while websocketTask.state == .running {
-                    switch try await websocketTask.receive() {
-                    case .string(let string):
-                        await decodeEvent(string.data(using: .utf8)!)
-                    case .data(let data):
-                        await decodeEvent(data)
-                    default:
-                        break
-                    }
-                }
-                
-                print("Websocket not running: \(String(describing: websocketTask.error))")
-            } catch {
-                encounteredWebsocketError(error)
-            }
+        guard self.eventTask == nil else {
+            completion(nil)
+            return
         }
         
+        self.eventTask = Task {
+            for await event in try await liaison.partialEvents() {
+                do {
+                    try await handleLiaisonEvent(event)
+                } catch {
+                    encounteredWebsocketError(error)
+                }
+            }
+            
+            print("Event loop exited")
+            connectionStatus = .error
+        }
         completion(nil)
     }
     
     func disconnect(_ completion: ((any Error)?) -> Void) {
-        websocketTask?.cancel()
-        websocketTask = nil
+        Task { await liaison.disconnect() }
         completion(nil)
     }
     
     func fetchDevices(_ completion: @escaping (Result<[AnyDevice], any Error>) -> Void) {
+        guard connectionStatus != .error else {
+            // Need to report this state to the UI.
+            completion(.failure(ServerError.disconnected))
+            return
+        }
+
         Task {
             await waitConnected()
             
             do {
-                let registryEntries: [RegistryEntry] = try await sendMessage(RegistryEntriesMessage())
+                let registryEntries: [RegistryEntry] = try await liaison.sendMessage(RegistryEntriesMessage())
                 let filteredEntries = registryEntries.filter { $0.labels.contains(Self.filterLabel) }
                 
-                let states: [EntityState] = try await sendMessage(GetStatesMessage())
+                let states: [EntityState] = try await liaison.sendMessage(GetStatesMessage())
                 let stateMap = states.reduce(into: [String: DeviceState]()) { partialResult, state in
                     partialResult[state.entityId] = switch state.state {
                     case .on:  .on
@@ -95,6 +95,11 @@ class HomeAssistantServer: Server
     }
     
     func toggleDevice(_ device: AnyDevice, state: DeviceState, completion: @escaping ((any Error)?) -> Void) {
+        guard connectionStatus != .error else {
+            completion(ServerError.disconnected)
+            return
+        }
+        
         Task {
             await waitConnected()
             
@@ -105,7 +110,19 @@ class HomeAssistantServer: Server
                     case .on:  .turnOn
                 }
                 
-                let _: Bool = try await sendMessage(CallServiceMessage(entityID: entityID, service: service))
+                var domain: CallServiceMessage.ServiceDomain = .switch
+                if entityID.starts(with: "light.") {
+                    // device.type is a UI concept, not related to what HA knows about it.
+                    domain = .light
+                }
+                
+                let message = CallServiceMessage(
+                    entityID: entityID,
+                    service: service,
+                    domain: domain
+                )
+                
+                let _: NoResult = try await liaison.sendMessage(message)
                 Task { @MainActor in completion(nil) }
             } catch {
                 Task { @MainActor in completion(error) }
@@ -125,77 +142,40 @@ class HomeAssistantServer: Server
             return await withCheckedContinuation { continuation in
                 self.connectionContinuations.append(continuation)
             }
+        } else if connectionStatus == .error {
+            return false
         }
         
         return true
     }
     
-    @discardableResult
-    private func sendMessage<R: Decodable>(_ message: Message) async throws -> R {
-        return try await withCheckedThrowingContinuation { continuation in
-            Task { @MainActor in
-                self.pendingMessages[message.id] = { response in
-                    do {
-                        let decodedResponse = try JSONDecoder().decode(Event<R>.self, from: response)
-                        if let result = decodedResponse.result {
-                            continuation.resume(returning: result)
-                        } else {
-                            let error = decodedResponse.error ?? ErrorResult(code: "unknown", message: "unknown")
-                            continuation.resume(throwing: error)
-                        }
-                    } catch {
-                        print("Response decoding error: \(error)")
-                        continuation.resume(throwing: error)
-                    }
-                    
-                    self.pendingMessages.removeValue(forKey: message.id)
-                }
-            }
-            
-            Task {
-                do {
-                    let encodedMessage = try JSONEncoder().encode(message)
-                    let jsonString = String(data: encodedMessage, encoding: .utf8)!
-                    try await websocketTask?.send(.string(jsonString))
-                } catch {
-                    print("Error sending message: \(error)")
-                    encounteredWebsocketError(error)
-                }
-            }
-        }
-    }
-    
     private func encounteredWebsocketError(_ error: any Error) {
         print("Websocket Error: \(error)")
         self.connectionStatus = .error
-        self.websocketTask?.cancel()
-        self.websocketTask = nil
-        // reconnects the next time we're refreshed. 
+        self.eventTask?.cancel()
+        self.eventTask = nil
     }
     
-    private func decodeEvent(_ data: Data) async {
-        do {
-            let decoder = JSONDecoder()
-            let partialEvent = try decoder.decode(PartialEvent.self, from: data)
-            try await handlePartialEvent(partialEvent, data: data)
-        } catch {
-            print("HA Decoding error: \(error)")
-        }
-    }
-    
-    private func handlePartialEvent(_ partialEvent: PartialEvent, data: Data) async throws {
-        Task {
+    private func handleLiaisonEvent(_ event: Liaison.Event) async throws {
+        switch event {
+        case .partialEvent(let partialEvent, _):
             switch partialEvent.type {
             case .authRequired:
-                await authenticate()
+                Task { await authenticate() }
             case .authOK:
-                didConnectToWebsocket()
-            case .result:
-                guard let id = partialEvent.id else { return }
-                Task { @MainActor in
-                    self.pendingMessages[id]?(data)
-                }
+                Task { didConnectToWebsocket() }
+            default:
+                break
             }
+            
+        case .pong:
+            didConnectToWebsocket()
+
+        case .disconnected(let error):
+            encounteredWebsocketError(error)
+            
+        case .error(let error):
+            encounteredWebsocketError(error)
         }
     }
     
@@ -210,12 +190,188 @@ class HomeAssistantServer: Server
     
     private func authenticate() async {
         do {
-            let _: Bool = try await sendMessage(AuthMessage(token: authToken))
+            let _: NoResult = try await liaison.sendMessage(AuthMessage(token: authToken))
         } catch {
             print("Error authenticating: \(error)")
         }
     }
     
+    // MARK: - Liaison
+    
+    actor Liaison
+    {
+        let websocketURL: URL
+        
+        private typealias Continuation = AsyncStream<Event>.Continuation
+        private typealias Reply = (Data) -> Void
+        
+        private var websocketTask: URLSessionWebSocketTask? = nil
+        private var eventTask: Task<(), Never>? = nil
+        private var pingTask: Task<(), Never>? = nil
+        
+        private var pendingMessages: [MessageID: Reply] = [:]
+        private var continuations: [UUID: Continuation] = [:]
+        
+        fileprivate init(websocketURL: URL) {
+            self.websocketURL = websocketURL
+        }
+        
+        fileprivate func connect() {
+            spawnWebsocketTaskIfNecessary()
+        }
+        
+        fileprivate func disconnect() {
+            self.websocketTask?.cancel()
+            self.websocketTask = nil
+        }
+        
+        @discardableResult
+        fileprivate func sendMessage<R: Decodable>(_ message: Message) async throws -> R {
+            return try await withCheckedThrowingContinuation { continuation in
+                Task {
+                    var outgoingMessage = message
+                    outgoingMessage.id = outgoingMessage.nextMessageID() // important that we update this here.
+                    
+                    self.pendingMessages[outgoingMessage.id] = { response in
+                        do {
+                            let decodedResponse = try JSONDecoder().decode(HomeAssistantServer.Event<R>.self, from: response)
+                            if let result = decodedResponse.result {
+                                continuation.resume(returning: result)
+                            } else if R.self != NoResult.self {
+                                let error = decodedResponse.error ?? ErrorResult(code: "unknown", message: "unknown")
+                                continuation.resume(throwing: error)
+                            }
+                        } catch {
+                            print("Response decoding error: \(error)")
+                            continuation.resume(throwing: error)
+                        }
+                        
+                        self.pendingMessages.removeValue(forKey: outgoingMessage.id)
+                    }
+
+                    let encodedMessage = try JSONEncoder().encode(outgoingMessage)
+                    let jsonString = String(data: encodedMessage, encoding: .utf8)!
+                    try await websocketTask?.send(.string(jsonString))
+                }
+            }
+        }
+        
+        fileprivate func partialEvents() throws -> AsyncStream<Liaison.Event> {
+            return AsyncStream { continuation in
+                let uuid = UUID()
+                self.continuations[uuid] = continuation
+                continuation.onTermination = { termination in
+                    Task { await self.removeContinuation(uuid: uuid) }
+                }
+                
+                spawnWebsocketTaskIfNecessary()
+            }
+        }
+        
+        private func removeContinuation(uuid: UUID) {
+            print("removing continuation: \(uuid)")
+            self.continuations.removeValue(forKey: uuid)
+        }
+        
+        private func restartWebsocket() {
+            disconnect()
+            spawnWebsocketTaskIfNecessary()
+        }
+
+        private func spawnWebsocketTaskIfNecessary() {
+            guard self.websocketTask == nil else { return }
+            
+            let websocketTask = URLSession.shared.webSocketTask(with: websocketURL)
+            self.websocketTask = websocketTask
+            websocketTask.resume()
+            
+            // Event Task
+            let decoder = JSONDecoder()
+            
+            eventTask?.cancel()
+            eventTask = Task.detached { [weak self] in
+                guard let self else { return }
+                do {
+                    while websocketTask.state == .running {
+                        switch try await websocketTask.receive() {
+                        case .string(let string):
+                            let data = string.data(using: .utf8)!
+                            let partialEvent = try decoder.decode(PartialEvent.self, from: data)
+                            await handlePartialEvent(partialEvent, data: data)
+                        case .data(let data):
+                            let partialEvent = try decoder.decode(PartialEvent.self, from: data)
+                            await handlePartialEvent(partialEvent, data: data)
+                        default:
+                            break
+                        }
+                    }
+                    
+                    print("Websocket not running: \(String(describing: websocketTask.error))")
+                } catch {
+                    await yield(.error(error))
+                }
+            }
+            
+            // Ping Task
+            pingTask?.cancel()
+            pingTask = Task.detached { [weak self] in
+                guard let self else { return }
+                guard let websocket = await self.websocketTask else { return }
+                
+                var pinging = true
+                while pinging {
+                    do {
+                        try await Task.sleep(for: .seconds(5))
+                    } catch { print("Task cancelled.") }
+                    
+                    websocket.sendPing { error in
+                        if let error {
+                            pinging = false
+                            print("Ping error: \(error). Disconnecting.")
+                            Task {
+                                await self.disconnect()
+                                await self.yield(.disconnected(error))
+                            }
+                        } else {
+                            print("pong!")
+                            Task {
+                                await self.yield(.pong)
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        
+        private func handlePartialEvent(_ event: PartialEvent, data: Data) {
+            if event.type == .result {
+                guard let id = event.id else { return }
+                if let resultHandler = self.pendingMessages[id] {
+                    resultHandler(data)
+                } else {
+                    print("No result handler for id: \(id)")
+                }
+            } else {
+                yield(.partialEvent(event, data))
+            }
+        }
+        
+        private func yield(_ event: Event) {
+            for (_, continuation) in continuations {
+                continuation.yield(event)
+            }
+        }
+        
+        enum Event
+        {
+            case partialEvent(PartialEvent, Data)
+            case pong
+            
+            case disconnected(any Error)
+            case error(any Error)
+        }
+    }
+    
     // MARK: - Types
     
     struct PartialEvent: Codable
@@ -239,12 +395,17 @@ class HomeAssistantServer: Server
         case result             = "result"
     }
         
-    struct ErrorResult: Codable, Error
+    struct ErrorResult: Codable, Swift.Error
     {
         let code: String
         let message: String
     }
     
+    enum ServerError: Swift.Error
+    {
+        case disconnected
+    }
+    
     // MARK: - Result Types
     
     struct RegistryEntry: Codable
@@ -305,10 +466,15 @@ class HomeAssistantServer: Server
 extension HomeAssistantDevice
 {
     convenience init(entity: HomeAssistantServer.RegistryEntry, state: DeviceState) {
+        var type: DeviceType = .switch
+        if entity.labels.contains("control_panel_light") {
+            type = .light
+        }
+        
         self.init(
             name: entity.name?.replacingOccurrences(of: " Switch", with: "") ?? entity.entityId,
             serial: entity.entityId,
-            type: .switch,
+            type: type,
             state: state
         )
     }
@@ -329,28 +495,38 @@ extension MessageID
 
 fileprivate protocol Message: Encodable
 {
-    var id: MessageID { get }
+    var id: MessageID { get set }
     var type: String { get }
+    
+    func nextMessageID() -> MessageID
+}
+
+extension Message
+{
+    func nextMessageID() -> MessageID { .next() }
 }
 
 fileprivate struct RegistryEntriesMessage: Message
 {
-    let id: MessageID = .next()
+    var id: MessageID = nil
     let type: String = "config/entity_registry/list"
 }
 
 fileprivate struct GetStatesMessage: Message
 {
-    let id: MessageID = .next()
+    var id: MessageID = nil
     let type: String = "get_states"
 }
 
 fileprivate struct AuthMessage: Message
 {
-    let id: MessageID = nil
+    var id: MessageID = nil
     let type: String = "auth"
     let token: String
     
+    // No ID for auth messages
+    func nextMessageID() -> MessageID { nil }
+    
     enum CodingKeys: String, CodingKey {
         case id
         case type
@@ -360,7 +536,7 @@ fileprivate struct AuthMessage: Message
 
 fileprivate struct CallServiceMessage: Message
 {
-    let id: MessageID = .next()
+    var id: MessageID = nil
     let type: String = "call_service"
     
     let service: Service
@@ -376,6 +552,7 @@ fileprivate struct CallServiceMessage: Message
     enum ServiceDomain: String, Codable
     {
         case `switch`
+        case light
     }
     
     enum Service: String, Codable