Xaibatsu-Control-Panel/XIONControlPanel/Servers/HomeAssistantServer.swift

//
//  HomeAssistantServer.swift
//  XIONControlPanel
//
//  Created by James Magahern on 6/18/25.
//  Copyright © 2025 XION. All rights reserved.
//

import Combine
import Foundation

class HomeAssistantServer: Server
{
    var connectionStatus: ConnectionStatus = .disconnected
    
    // Configure:
    private let authToken: String = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiI2ZTAwODI1N2Y4N2Q0OWIxYWRkN2ExNTBhOWRmNGZiZCIsImlhdCI6MTc1MDMwMDg5NiwiZXhwIjoyMDY1NjYwODk2fQ.KRf9GpRZdpW9Z_vkbL3sl74rgKS7eAyc8kO0a5jLTgg"
    
    // Configure:
    private static let filterLabel = "control_panel"
    
    private let url: URL
    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 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) {
        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 liaison.sendMessage(RegistryEntriesMessage())
                let filteredEntries = registryEntries.filter { $0.labels.contains(Self.filterLabel) }
                
                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
                    case .off: .off
                    case .other(_): .off
                    }
                }
                
                let devices = filteredEntries.map { entry in
                    HomeAssistantDevice(entity: entry, state: stateMap[entry.entityId] ?? .on)
                        .eraseToAnyDevice()
                }
                
                Task { @MainActor in completion(.success(devices)) }
            } catch {
                Task { @MainActor in completion(.failure(error)) }
            }
        }
    }
    
    func toggleDevice(_ device: AnyDevice, state: DeviceState, completion: @escaping ((any Error)?) -> Void) {
        guard connectionStatus != .error else {
            completion(ServerError.disconnected)
            return
        }
        
        Task {
            await waitConnected()
            
            do {
                let entityID = device.serial
                let service: CallServiceMessage.Service = switch state {
                    case .off: .turnOff
                    case .on:  .turnOn
                }
                
                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) }
            }
        }
    }
    
    func responsibleForDevice(_ device: AnyDevice) -> Bool {
        return true
    }
    
    // MARK: -
    
    @discardableResult
    private func waitConnected() async -> Bool {
        if connectionStatus == .connecting || connectionStatus == .disconnected {
            return await withCheckedContinuation { continuation in
                self.connectionContinuations.append(continuation)
            }
        } else if connectionStatus == .error {
            return false
        }
        
        return true
    }
    
    private func encounteredWebsocketError(_ error: any Error) {
        print("Websocket Error: \(error)")
        self.connectionStatus = .error
        self.eventTask?.cancel()
        self.eventTask = nil
    }
    
    private func handleLiaisonEvent(_ event: Liaison.Event) async throws {
        switch event {
        case .partialEvent(let partialEvent, _):
            switch partialEvent.type {
            case .authRequired:
                Task { await authenticate() }
            case .authOK:
                Task { didConnectToWebsocket() }
            default:
                break
            }
            
        case .pong:
            didConnectToWebsocket()

        case .disconnected(let error):
            encounteredWebsocketError(error)
            
        case .error(let error):
            encounteredWebsocketError(error)
        }
    }
    
    private func didConnectToWebsocket() {
        self.connectionStatus = .connected
        
        connectionContinuations.forEach { continuation in
            continuation.resume(returning: true)
        }
        connectionContinuations.removeAll()
    }
    
    private func authenticate() async {
        do {
            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
    {
        let id: Int64?
        let type: EventType
    }
    
    struct Event<Result: Decodable>: Decodable
    {
        let id: Int?
        let type: EventType
        let result: Result?
        let error: ErrorResult?
    }
    
    enum EventType: String, Codable
    {
        case authRequired       = "auth_required"
        case authOK             = "auth_ok"
        case result             = "result"
    }
        
    struct ErrorResult: Codable, Swift.Error
    {
        let code: String
        let message: String
    }
    
    enum ServerError: Swift.Error
    {
        case disconnected
    }
    
    // MARK: - Result Types
    
    struct RegistryEntry: Codable
    {
        let id: String
        let entityId: String
        let name: String?
        let labels: [String]
        
        enum CodingKeys: String, CodingKey {
            case id
            case name
            case labels
            case entityId = "entity_id"
        }
    }
    
    struct EntityState: Decodable
    {
        let entityId: String
        let state: State
        
        enum State: Decodable {
            case on
            case off
            case other(String)
            
            init(from decoder: Decoder) throws {
                let container = try decoder.singleValueContainer()
                let stringValue = try container.decode(String.self)
                
                switch stringValue.lowercased() {
                case "on":
                    self = .on
                case "off":
                    self = .off
                default:
                    self = .other(stringValue)
                }
            }
        }
        
        struct Attributes: Decodable {
            let friendlyName: String?
            
            enum CodingKeys: String, CodingKey {
                case friendlyName = "friendly_name"
            }
        }
        
        enum CodingKeys: String, CodingKey {
            case entityId = "entity_id"
            case state
        }
    }
}

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: type,
            state: state
        )
    }
}

typealias MessageID = Int64?

extension MessageID
{
    private static var counter: Int64 = 1
    static func next() -> Self {
        OSAtomicIncrement64(&Self.counter)
        return Self.counter
    }
}

// MARK: - Message Types

fileprivate protocol Message: Encodable
{
    var id: MessageID { get set }
    var type: String { get }
    
    func nextMessageID() -> MessageID
}

extension Message
{
    func nextMessageID() -> MessageID { .next() }
}

fileprivate struct RegistryEntriesMessage: Message
{
    var id: MessageID = nil
    let type: String = "config/entity_registry/list"
}

fileprivate struct GetStatesMessage: Message
{
    var id: MessageID = nil
    let type: String = "get_states"
}

fileprivate struct AuthMessage: Message
{
    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
        case token = "access_token"
    }
}

fileprivate struct CallServiceMessage: Message
{
    var id: MessageID = nil
    let type: String = "call_service"
    
    let service: Service
    let domain: ServiceDomain
    let target: Target
    
    init(entityID: String, service: Service, domain: ServiceDomain = .switch) {
        self.service = service
        self.domain = domain
        self.target = Target(entityID: entityID)
    }
    
    enum ServiceDomain: String, Codable
    {
        case `switch`
        case light
    }
    
    enum Service: String, Codable
    {
        case toggle
        case turnOn   = "turn_on"
        case turnOff  = "turn_off"
    }
    
    struct Target: Codable
    {
        let entityID: String
        
        enum CodingKeys: String, CodingKey {
            case entityID = "entity_id"
        }
    }
}