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 Foundation

class HomeAssistantServer: Server
{
    var connectionStatus: ConnectionStatus = .disconnected
    
    private typealias Reply = (Data) -> Void
    
    // Configure:
    private let authToken: String = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiI2ZTAwODI1N2Y4N2Q0OWIxYWRkN2ExNTBhOWRmNGZiZCIsImlhdCI6MTc1MDMwMDg5NiwiZXhwIjoyMDY1NjYwODk2fQ.KRf9GpRZdpW9Z_vkbL3sl74rgKS7eAyc8kO0a5jLTgg"
    
    // Configure:
    private static let filterLabel = "control_panel"
    
    private let url: URL
    private var websocketTask: URLSessionWebSocketTask? = nil
    private var pendingMessages: [MessageID: Reply] = [:]
    private var connectionContinuations: [CheckedContinuation<Bool, Never>] = []
    
    required init(_ url: URL) {
        self.url = url
    }
    
    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)
            }
        }
        
        completion(nil)
    }
    
    func disconnect(_ completion: ((any Error)?) -> Void) {
        websocketTask?.cancel()
        websocketTask = nil
        completion(nil)
    }
    
    func fetchDevices(_ completion: @escaping (Result<[AnyDevice], any Error>) -> Void) {
        Task {
            await waitConnected()
            
            do {
                let registryEntries: [RegistryEntry] = try await sendMessage(RegistryEntriesMessage())
                let filteredEntries = registryEntries.filter { $0.labels.contains(Self.filterLabel) }
                
                let states: [EntityState] = try await 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) {
        Task {
            await waitConnected()
            
            do {
                let entityID = device.serial
                let service: CallServiceMessage.Service = switch state {
                    case .off: .turnOff
                    case .on:  .turnOn
                }
                
                let _: Bool = try await sendMessage(CallServiceMessage(entityID: entityID, service: service))
                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)
            }
        }
        
        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. 
    }
    
    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 {
            switch partialEvent.type {
            case .authRequired:
                await authenticate()
            case .authOK:
                didConnectToWebsocket()
            case .result:
                guard let id = partialEvent.id else { return }
                Task { @MainActor in
                    self.pendingMessages[id]?(data)
                }
            }
        }
    }
    
    private func didConnectToWebsocket() {
        self.connectionStatus = .connected
        
        connectionContinuations.forEach { continuation in
            continuation.resume(returning: true)
        }
        connectionContinuations.removeAll()
    }
    
    private func authenticate() async {
        do {
            let _: Bool = try await sendMessage(AuthMessage(token: authToken))
        } catch {
            print("Error authenticating: \(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, Error
    {
        let code: String
        let message: String
    }
    
    // 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) {
        self.init(
            name: entity.name?.replacingOccurrences(of: " Switch", with: "") ?? entity.entityId,
            serial: entity.entityId,
            type: .switch,
            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 }
    var type: String { get }
}

fileprivate struct RegistryEntriesMessage: Message
{
    let id: MessageID = .next()
    let type: String = "config/entity_registry/list"
}

fileprivate struct GetStatesMessage: Message
{
    let id: MessageID = .next()
    let type: String = "get_states"
}

fileprivate struct AuthMessage: Message
{
    let id: MessageID = nil
    let type: String = "auth"
    let token: String
    
    enum CodingKeys: String, CodingKey {
        case id
        case type
        case token = "access_token"
    }
}

fileprivate struct CallServiceMessage: Message
{
    let id: MessageID = .next()
    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`
    }
    
    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"
        }
    }
}