DESCracker/ios/DESCracker/Network/SlaveClient.swift

import Foundation

protocol SlaveClientDelegate: AnyObject {
    func slaveClient(_ client: SlaveClient, didChangeStatus status: String, connected: Bool)
    func slaveClient(_ client: SlaveClient, didLog line: String)
    func slaveClient(_ client: SlaveClient, didUpdateWorkers workers: [WorkerSnapshot])
    func slaveClient(_ client: SlaveClient, didHitJob jobId: String, index: String, keyHex: String)
}

final class SlaveClient: NSObject, URLSessionWebSocketDelegate {
    weak var delegate: SlaveClientDelegate?
    var cpuWorkers: Int
    var isConnected: Bool { task?.state == .running }
    var metalAvailable: Bool { metal.isAvailable }
    var metalName: String { metal.gpuName }
    var metalMemory: UInt64 { metal.memoryBytes }
    var cpuCores: Int { cpu.coreCount }

    private var devices: [ComputeDevice] = []
    private var session: URLSession!
    private var task: URLSessionWebSocketTask?
    private var pingTimer: Timer?
    private var jobs: [String: DeviceJob] = [:]
    private var completedBlocks: [String: Int] = [:]
    private var progressGates: [String: ProgressGate] = [:]
    private var rateStats: [String: RateStats] = [:]
    private let metal = MetalEngine()
    private let cpu = CpuEngine()
    private let syncQueue = DispatchQueue(label: "com.descracker.slave")
    private var connectURL = ""
    private var loggedMetalCompile = false

    override init() {
        cpuWorkers = max(1, ProcessInfo.processInfo.processorCount)
        super.init()
        let config = URLSessionConfiguration.default
        config.waitsForConnectivity = false
        session = URLSession(configuration: config, delegate: self, delegateQueue: nil)
    }

    func prepareEngines() throws {
        try cpu.prepare()
    }

    func prepareEngines(completion: @escaping (Bool, Error?) -> Void) {
        DispatchQueue.global(qos: .userInitiated).async {
            var captured: Error?
            var ok = true
            do {
                try self.prepareEngines()
            } catch {
                captured = error
                ok = false
            }
            DispatchQueue.main.async {
                completion(ok, captured)
            }
        }
    }

    func setDevices(_ devices: [ComputeDevice]) {
        syncQueue.async {
            self.devices = devices
            for device in self.devices where self.completedBlocks[device.key] == nil {
                self.completedBlocks[device.key] = 0
            }
            self.emitWorkers()
        }
    }

    func connect(to host: String, port: Int) {
        disconnect()
        let urlString = "ws://\(host):\(port)"
        connectURL = urlString
        notifyStatus("Connecting to \(urlString)…", connected: false)
        guard let url = URL(string: urlString) else {
            log("Invalid URL \(urlString)")
            notifyStatus("Disconnected", connected: false)
            return
        }
        task = session.webSocketTask(with: url)
        task?.resume()
    }

    func disconnect() {
        syncQueue.sync {
            clearPing()
            killAll()
        }
        let current = task
        task = nil
        current?.cancel(with: .normalClosure, reason: nil)
        syncQueue.async { self.emitWorkers() }
    }

    func urlSession(
        _ session: URLSession,
        webSocketTask: URLSessionWebSocketTask,
        didOpenWithProtocol protocol: String?
    ) {
        guard webSocketTask === task else { return }
        syncQueue.async {
            let deviceJSON = self.devices.map { $0.jsonObject() }
            let hostname = ProcessInfo.processInfo.hostName
            self.send([
                "type": "hello",
                "hostname": hostname,
                "platform": "ios arm64",
                "devices": deviceJSON,
            ])
            let exposed = self.devices.isEmpty
                ? "no workers"
                : self.devices.map(\.name).joined(separator: ", ")
            self.notifyStatus("Connected to \(self.connectURL)", connected: true)
            self.log("Hello sent (protocol \(Wire.protocolVersion)) — exposing \(exposed)")
            self.emitWorkers()
            DispatchQueue.main.async {
                self.pingTimer?.invalidate()
                self.pingTimer = Timer.scheduledTimer(withTimeInterval: Wire.pingInterval, repeats: true) { [weak self] _ in
                    self?.send(["type": "ping"])
                }
            }
            self.listen()
        }
    }

    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        if task !== self.task && self.task != nil { return }
        syncQueue.async {
            self.clearPing()
            self.killAll()
            if self.task === task {
                self.task = nil
            }
            if let error, (error as NSError).code != NSURLErrorCancelled {
                self.log("WebSocket error: \(error.localizedDescription)")
            }
            self.notifyStatus("Disconnected", connected: false)
            self.emitWorkers()
        }
    }

    private func listen() {
        guard let task else { return }
        task.receive { [weak self] result in
            guard let self, task === self.task else { return }
            switch result {
            case .failure:
                return
            case .success(let message):
                if case .string(let text) = message, let msg = Wire.decodeMessage(text) {
                    self.onServer(msg)
                }
                self.listen()
            }
        }
    }

    private func onServer(_ msg: [String: Any]) {
        let type = msg["type"] as? String
        switch type {
        case "hello_ack":
            log("Master ack v\(msg["protocolVersion"] ?? "?")")
        case "assign":
            guard let block = AssignBlock.fromJSON(msg) else {
                log("Ignored malformed assign")
                return
            }
            runAssign(block)
        case "cancel":
            syncQueue.async {
                self.killAll()
                self.log("Cancel received")
                self.notifyStatus("Idle (cancelled)", connected: true)
                self.emitWorkers()
            }
        case "shutdown":
            DispatchQueue.main.async { self.disconnect() }
        default:
            break
        }
    }

    private func engine(for device: ComputeDevice) -> BruteEngine {
        device.type == "metal" ? metal : cpu
    }

    private func device(forKey key: String) -> ComputeDevice? {
        if let match = devices.first(where: { $0.key == key }) { return match }
        return devices.first(where: { $0.type == "metal" }) ?? devices.first
    }

    private func runAssign(_ block: AssignBlock) {
        syncQueue.async { [weak self] in
            guard let self else { return }
            let matched = self.devices.first(where: { $0.key == block.deviceKey })
            guard let device = matched ?? self.device(forKey: block.deviceKey) else {
                self.send(["type": "error", "jobId": block.jobId, "message": "no compute device selected"])
                return
            }
            if matched == nil {
                self.log("No device \(block.deviceKey); using \(device.name)")
            }
            self.jobs[device.key]?.kill?()
            self.clearProgressGate(device.key)
            self.log(String(
                format: "Assigned %@  len=%u pad=0x%02x start=%llu count=%llu",
                device.name,
                block.keyLen,
                block.padByte,
                block.start,
                block.count
            ))
            if device.type == "metal" && !self.loggedMetalCompile {
                self.loggedMetalCompile = true
                self.log("Compiling Metal kernel for this GPU (first job may pause)…")
            }
            if device.type == "cpu" && block.cpuWorkers == nil {
                block.cpuWorkers = self.cpuWorkers
            }
            let started = Date()
            let hitCount = Locked(0)
            let engine = self.engine(for: device)
            let deviceKey = device.key
            let jobId = block.jobId
            let job = DeviceJob(jobId: jobId, count: block.count)
            job.kill = { engine.cancel() }
            self.jobs[deviceKey] = job
            self.beginRate(deviceKey)
            self.updateStatus()
            self.emitWorkers()

            engine.runAssign(
                block,
                onProgress: { done, count, rate in
                    self.syncQueue.async {
                        if let live = self.jobs[deviceKey], live.jobId == jobId {
                            live.done = done
                            live.count = count
                            live.rate = rate
                        }
                        self.sendProgress(jobId: jobId, deviceKey: deviceKey, done: done, count: count, rate: rate)
                        self.updateStatus()
                        self.emitWorkers()
                    }
                },
                onHit: { index, key, plain in
                    self.syncQueue.async {
                        hitCount.value += 1
                        let indexS = Wire.formatDec64(index)
                        let keyHex = Wire.formatHex64(key)
                        let plainHex = Wire.formatHex64(plain)
                        self.send([
                            "type": "hit",
                            "jobId": jobId,
                            "index": indexS,
                            "keyHex": keyHex,
                            "plainHex": plainHex,
                        ])
                        DispatchQueue.main.async {
                            self.delegate?.slaveClient(self, didHitJob: jobId, index: indexS, keyHex: keyHex)
                        }
                    }
                },
                completion: { _, elapsed, cancelled, error in
                    self.syncQueue.async {
                        let still = self.jobs[deviceKey]
                        let same = still?.jobId == jobId
                        if cancelled {
                            if same, let still {
                                self.endRate(deviceKey, keys: still.done)
                                self.jobs.removeValue(forKey: deviceKey)
                                self.clearProgressGate(deviceKey)
                            }
                            self.updateStatus()
                            self.emitWorkers()
                            return
                        }
                        if let error {
                            if same, let still {
                                self.endRate(deviceKey, keys: still.done)
                                self.jobs.removeValue(forKey: deviceKey)
                                self.clearProgressGate(deviceKey)
                            }
                            self.send([
                                "type": "error",
                                "jobId": jobId,
                                "message": error.localizedDescription,
                            ])
                            self.log(error.localizedDescription)
                            self.updateStatus()
                            self.emitWorkers()
                            return
                        }
                        if same {
                            self.jobs.removeValue(forKey: deviceKey)
                            self.endRate(deviceKey, keys: block.count)
                            self.completedBlocks[deviceKey] = (self.completedBlocks[deviceKey] ?? 0) + 1
                            self.clearProgressGate(deviceKey)
                            let elapsedOut = elapsed > 0 ? elapsed : Date().timeIntervalSince(started)
                            self.send([
                                "type": "block_complete",
                                "jobId": jobId,
                                "deviceKey": deviceKey,
                                "hits": hitCount.value,
                                "elapsed": elapsedOut,
                            ])
                        }
                        self.updateStatus()
                        self.emitWorkers()
                    }
                }
            )
        }
    }

    private func send(_ msg: [String: Any]) {
        guard let json = Wire.encodeMessage(msg),
              let task, task.state == .running
        else { return }
        task.send(.string(json)) { [weak self] error in
            if let error {
                self?.log("Send failed: \(error.localizedDescription)")
            }
        }
    }

    private func rate(for deviceKey: String) -> RateStats {
        if let stats = rateStats[deviceKey] { return stats }
        let stats = RateStats()
        rateStats[deviceKey] = stats
        return stats
    }

    private func beginRate(_ deviceKey: String) {
        let stats = rate(for: deviceKey)
        if stats.busyStarted == nil {
            stats.busyStarted = Date()
        }
    }

    private func endRate(_ deviceKey: String, keys: UInt64) {
        let stats = rate(for: deviceKey)
        if let started = stats.busyStarted {
            stats.activeMs += Date().timeIntervalSince(started) * 1000
            stats.busyStarted = nil
        }
        if keys > 0 {
            stats.keysDone += keys
        }
    }

    private func killAll() {
        for (deviceKey, job) in jobs {
            endRate(deviceKey, keys: job.done)
            job.kill?()
            clearProgressGate(deviceKey)
        }
        jobs.removeAll()
    }

    private func clearPing() {
        DispatchQueue.main.async {
            self.pingTimer?.invalidate()
            self.pingTimer = nil
        }
    }

    private func progressGate(_ deviceKey: String) -> ProgressGate {
        if let gate = progressGates[deviceKey] { return gate }
        let gate = ProgressGate()
        progressGates[deviceKey] = gate
        return gate
    }

    private func sendProgress(jobId: String, deviceKey: String, done: UInt64, count: UInt64, rate: Double) {
        let gate = progressGate(deviceKey)
        gate.payload = [
            "type": "progress",
            "jobId": jobId,
            "deviceKey": deviceKey,
            "done": Wire.formatDec64(done),
            "count": Wire.formatDec64(count),
            "rate": rate,
        ]
        let wait = Wire.progressInterval - (Date().timeIntervalSince1970 - gate.lastSent)
        if wait <= 0 {
            flushProgress(deviceKey)
            return
        }
        if !gate.timerPending {
            gate.timerPending = true
            let key = deviceKey
            syncQueue.asyncAfter(deadline: .now() + wait) {
                self.progressGates[key]?.timerPending = false
                self.flushProgress(key)
            }
        }
    }

    private func flushProgress(_ deviceKey: String) {
        guard let gate = progressGates[deviceKey], let payload = gate.payload else { return }
        let live = jobs[deviceKey]
        if live == nil || live?.jobId != payload["jobId"] as? String {
            gate.payload = nil
            return
        }
        gate.payload = nil
        gate.lastSent = Date().timeIntervalSince1970
        send(payload)
    }

    private func clearProgressGate(_ deviceKey: String) {
        guard let gate = progressGates[deviceKey] else { return }
        gate.timerPending = false
        gate.payload = nil
    }

    private func emitWorkers() {
        let rows: [WorkerSnapshot] = devices.map { device in
            let live = jobs[device.key]
            let done = live?.done ?? 0
            let count = live?.count ?? 0
            let pct = live != nil && count > 0 ? Double(done) * 100.0 / Double(count) : 0
            let row = WorkerSnapshot()
            row.deviceKey = device.key
            row.deviceName = device.name
            row.deviceType = device.type
            row.current = live == nil ? "idle" : "\(Wire.formatCount(done)) / \(Wire.formatCount(count))"
            row.pct = pct
            row.completedBlocks = completedBlocks[device.key] ?? 0
            row.rate = live?.rate ?? 0
            row.busy = live != nil
            return row
        }
        DispatchQueue.main.async {
            self.delegate?.slaveClient(self, didUpdateWorkers: rows)
        }
    }

    private func updateStatus() {
        let busy = Array(jobs.values)
        let connected = task != nil
        if busy.isEmpty {
            notifyStatus(connected ? "Idle" : "Disconnected", connected: connected)
            return
        }
        if busy.count == 1, let job = busy.first {
            let shortId = job.jobId.count > 8 ? String(job.jobId.prefix(8)) : job.jobId
            notifyStatus("Working  \(shortId)  \(job.done)/\(job.count)", connected: true)
            return
        }
        notifyStatus("\(busy.count) workers assigned", connected: true)
    }

    private func notifyStatus(_ status: String, connected: Bool) {
        DispatchQueue.main.async {
            self.delegate?.slaveClient(self, didChangeStatus: status, connected: connected)
        }
    }

    private func log(_ line: String) {
        DispatchQueue.main.async {
            self.delegate?.slaveClient(self, didLog: line)
        }
    }
}

private final class DeviceJob {
    var jobId: String
    var done: UInt64 = 0
    var count: UInt64
    var rate = 0.0
    var kill: (() -> Void)?

    init(jobId: String, count: UInt64) {
        self.jobId = jobId
        self.count = count
    }
}

private final class RateStats {
    var keysDone: UInt64 = 0
    var activeMs: TimeInterval = 0
    var busyStarted: Date?
}

private final class ProgressGate {
    var lastSent: TimeInterval = 0
    var timerPending = false
    var payload: [String: Any]?
}