blackroad-watch/BlackRoadWatch/Sources/iOS/BlackRoadUDP.swift

import Foundation
import Network
import WatchConnectivity
import Combine

// MARK: - UDP Listener for M1s Firmware Broadcasts

class BlackRoadUDPManager: NSObject, ObservableObject {
    static let shared = BlackRoadUDPManager()

    @Published var isListening = false
    @Published var isConnected = false
    @Published var watchReachable = false
    @Published var sensorData: SensorData?
    @Published var aiStatus: AIStatus?
    @Published var systemHealth: SystemHealth?
    @Published var lastUpdate: Date?
    @Published var packetsReceived: UInt64 = 0
    @Published var sourceAddress: String = ""

    private var listener: NWListener?
    private var wcSession: WCSession?
    private let udpPort: UInt16 = 8420
    private let queue = DispatchQueue(label: "io.blackroad.udp", qos: .userInteractive)

    override init() {
        super.init()
        setupWatchConnectivity()
    }

    // MARK: - UDP Listener

    func startListening() {
        guard listener == nil else { return }

        do {
            let params = NWParameters.udp
            params.allowLocalEndpointReuse = true
            params.requiredInterfaceType = .wifi

            listener = try NWListener(using: params, on: NWEndpoint.Port(rawValue: udpPort)!)

            listener?.stateUpdateHandler = { [weak self] state in
                DispatchQueue.main.async {
                    switch state {
                    case .ready:
                        self?.isListening = true
                        print("[BR] UDP listener ready on port \(self?.udpPort ?? 0)")
                    case .failed(let error):
                        self?.isListening = false
                        print("[BR] UDP listener failed: \(error)")
                        DispatchQueue.main.asyncAfter(deadline: .now() + 3) {
                            self?.stopListening()
                            self?.startListening()
                        }
                    case .cancelled:
                        self?.isListening = false
                    default:
                        break
                    }
                }
            }

            listener?.newConnectionHandler = { [weak self] connection in
                self?.handleConnection(connection)
            }

            listener?.start(queue: queue)
        } catch {
            print("[BR] Failed to create UDP listener: \(error)")
        }
    }

    func stopListening() {
        listener?.cancel()
        listener = nil
        DispatchQueue.main.async {
            self.isListening = false
            self.isConnected = false
        }
    }

    private func handleConnection(_ connection: NWConnection) {
        connection.start(queue: queue)
        receiveData(on: connection)
    }

    private func receiveData(on connection: NWConnection) {
        connection.receiveMessage { [weak self] data, context, isComplete, error in
            if let data = data, !data.isEmpty {
                self?.processPacket(data, from: connection)
            }
            if error == nil {
                self?.receiveData(on: connection)
            }
        }
    }

    // MARK: - JSON Parsing

    private func processPacket(_ data: Data, from connection: NWConnection) {
        guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let type = json["type"] as? String, type == "br_watch" else {
            return
        }

        if let endpoint = connection.currentPath?.remoteEndpoint,
           case let .hostPort(host, _) = endpoint {
            let addr = "\(host)"
            DispatchQueue.main.async { self.sourceAddress = addr }
        }

        DispatchQueue.main.async { [weak self] in
            guard let self = self else { return }

            self.packetsReceived += 1
            self.isConnected = true
            self.lastUpdate = Date()

            if let sensorDict = json["sensor"] as? [String: Any] {
                self.sensorData = SensorData(
                    temperature: sensorDict["temp"] as? Double ?? 0,
                    humidity: sensorDict["hum"] as? Double ?? 0,
                    light: 0,
                    accelX: 0, accelY: 0, accelZ: 0,
                    batteryMV: sensorDict["bat"] as? Int ?? 0,
                    uptimeSec: UInt32(sensorDict["up"] as? Int ?? 0)
                )
            }

            if let aiDict = json["ai"] as? [String: Any] {
                self.aiStatus = AIStatus(
                    modelID: 1,
                    confidence: aiDict["conf"] as? Int ?? 0,
                    inferenceMS: 0,
                    totalInferences: UInt32(aiDict["infers"] as? Int ?? 0),
                    npuLoad: aiDict["load"] as? Int ?? 0,
                    npuTemp: aiDict["temp"] as? Int ?? 0,
                    classID: 0
                )
            }

            if let fleetDict = json["fleet"] as? [String: Any] {
                self.systemHealth = SystemHealth(
                    fleetOnline: fleetDict["on"] as? Int ?? 0,
                    fleetTotal: fleetDict["total"] as? Int ?? 0,
                    agentsActive: fleetDict["agents"] as? Int ?? 0,
                    trafficGreen: fleetDict["green"] as? Int ?? 0,
                    trafficYellow: 0,
                    trafficRed: 0,
                    tasksPending: 0,
                    tasksDone: fleetDict["tasks"] as? Int ?? 0,
                    memoryEntries: UInt32(fleetDict["mem"] as? Int ?? 0),
                    reposCount: fleetDict["repos"] as? Int ?? 0,
                    cfProjects: fleetDict["cf"] as? Int ?? 0,
                    cpuLoad: 0
                )
            }

            self.sendToWatch()
        }
    }

    // MARK: - Watch Connectivity

    private func setupWatchConnectivity() {
        guard WCSession.isSupported() else { return }
        wcSession = WCSession.default
        wcSession?.delegate = self
        wcSession?.activate()
    }

    private func sendToWatch() {
        guard let session = wcSession, session.activationState == .activated else { return }

        let encoder = JSONEncoder()
        var context: [String: Any] = [:]

        if let sensor = sensorData,
           let sensorJSON = try? encoder.encode(sensor),
           let sensorDict = try? JSONSerialization.jsonObject(with: sensorJSON) {
            context["sensor"] = sensorDict
        }
        if let ai = aiStatus,
           let aiJSON = try? encoder.encode(ai),
           let aiDict = try? JSONSerialization.jsonObject(with: aiJSON) {
            context["ai"] = aiDict
        }
        if let health = systemHealth,
           let healthJSON = try? encoder.encode(health),
           let healthDict = try? JSONSerialization.jsonObject(with: healthJSON) {
            context["health"] = healthDict
        }
        context["timestamp"] = Date().timeIntervalSince1970

        try? session.updateApplicationContext(context)

        if session.isReachable {
            session.sendMessage(context, replyHandler: nil)
        }

        DispatchQueue.main.async {
            self.watchReachable = session.isReachable
        }
    }
}

// MARK: - WCSessionDelegate

extension BlackRoadUDPManager: WCSessionDelegate {
    func session(_ session: WCSession, activationDidCompleteWith activationState: WCSessionActivationState, error: Error?) {
        print("[BR] Watch session activated: \(activationState.rawValue)")
        DispatchQueue.main.async { self.watchReachable = session.isReachable }
    }

    func sessionDidBecomeInactive(_ session: WCSession) {}
    func sessionDidDeactivate(_ session: WCSession) { session.activate() }

    func sessionReachabilityDidChange(_ session: WCSession) {
        DispatchQueue.main.async { self.watchReachable = session.isReachable }
    }
}