import Foundation
import NucleicProtocol

/// The mesh-wide AFM job queue (docs/COVALENCE_RUNNER.md §5, item 6): the host-side scheduler
/// that fans delegated intelligence work out across every connected mesh member that advertised
/// `WireClientCapabilities.canProvideIntelligence`. Jobs run **in parallel across devices** (one
/// in flight per worker — each device's own `AFMRequestQueue` serializes on its Neural Engine
/// anyway, so queuing more there only hides work from better placement here).
///
/// Placement rules:
/// - **Class restriction.** A job's `IntelligencePriority` gates which devices may run it: the
///   two highest tiers (`interactive`, `bashSummary`) feed things the user is actively reading
///   and need desktop-class tok/s, so they only ever run on laptop/desktop workers;
///   `completion`/`background` may run on iPhones/iPads too.
/// - **Power bias.** Pending jobs are dispatched highest-priority-first, and each picks the
///   most powerful *idle* eligible worker by the Apple Silicon hierarchy
///   (`IntelligenceWorkerProfile.powerRank`: M-Ultra > M-Max > M-Pro > M > A-series) — so the
///   hottest work always lands on the strongest chip currently free.
///
/// A job whose deadline lapses (or that exhausts its re-dispatch attempts after a worker
/// disconnect) resolves nil, and the submitting provider falls back to heuristics — a slow or
/// vanishing device can delay soft-AI polish, never block it.
public actor MeshIntelligenceQueue {
    /// Pushes one request at a registered worker's connection (a `ConnectionHandler` seal+send).
    public typealias SendRequest = @Sendable (WireIntelligenceRequest) async -> Void

    private struct Worker {
        let deviceID: String
        /// Which connection registered this worker — a reconnect replaces the registration, and
        /// the *old* connection's teardown must not unregister the new one (the same dedup
        /// hazard `SyncHost.handlersByDevice` guards).
        let connectionID: String
        let profile: IntelligenceWorkerProfile
        let send: SendRequest
        /// Request ids currently dispatched to this device.
        var inFlight: Set<String> = []
    }

    private struct Job {
        let request: WireIntelligenceRequest
        /// FIFO tie-breaker within a priority tier.
        let seq: UInt64
        /// Dispatch attempts so far — a worker disconnect re-queues the job until this hits
        /// ``MeshIntelligenceQueue/maxAttempts``.
        var attempts: Int = 0
        let continuation: CheckedContinuation<WireIntelligenceResult?, Never>
    }

    /// How many jobs one worker holds at a time. One: the device-side AFM queue is serial, so
    /// a second queued job would just sit there while a stronger worker goes idle here.
    private let perWorkerCap = 1
    /// How many times a job is dispatched before giving up (first send + one re-dispatch after
    /// a worker disconnect).
    private let maxAttempts = 2
    /// Deadline when the request carries none.
    private let defaultDeadline: Double = 60

    private var workers: [String: Worker] = [:]  // by deviceID
    private var pending: [Job] = []
    private var inFlight: [String: (job: Job, deviceID: String)] = [:]  // by request id
    private var deadlines: [String: Task<Void, Never>] = [:]
    private var nextSeq: UInt64 = 0

    public init() {}

    // MARK: - Worker registry (driven by SyncHost registration)

    /// A device that advertised `canProvideIntelligence` connected (or reconnected — the newest
    /// registration wins). Kicks dispatch, so work queued while no eligible device was around
    /// starts the moment one appears.
    public func registerWorker(
        deviceID: String, connectionID: String,
        profile: IntelligenceWorkerProfile, send: @escaping SendRequest
    ) {
        // A reconnect replaces the slot; anything the *old* connection was running can't answer
        // on the new socket, so put it back in line.
        if let old = workers[deviceID] {
            workers[deviceID] = nil
            for id in old.inFlight { requeue(id) }
        }
        workers[deviceID] = Worker(
            deviceID: deviceID, connectionID: connectionID, profile: profile, send: send)
        dispatch()
    }

    /// The registering connection closed. Guarded by `connectionID` so a superseded (deduped)
    /// connection's teardown can't unregister the replacement that's already live.
    public func unregisterWorker(deviceID: String, connectionID: String) {
        guard let worker = workers[deviceID], worker.connectionID == connectionID else { return }
        workers[deviceID] = nil
        for id in worker.inFlight { requeue(id) }
        dispatch()
    }

    // MARK: - Submit / results

    /// Run one delegated job on the mesh: queues it, dispatches per the placement rules, and
    /// resolves with the executor's answer — or nil on deadline, disconnect exhaustion, an
    /// executor error, or when no device eligible for its priority is connected at all (fail
    /// fast, so the caller's heuristic fallback renders immediately instead of waiting out a
    /// deadline nothing could meet).
    public func submit(_ request: WireIntelligenceRequest) async -> WireIntelligenceResult? {
        guard hasEligibleWorker(for: request.priority) else { return nil }
        let result = await withCheckedContinuation {
            (continuation: CheckedContinuation<WireIntelligenceResult?, Never>) in
            let job = Job(request: request, seq: nextSeq, continuation: continuation)
            nextSeq += 1
            pending.append(job)
            armDeadline(for: request)
            dispatch()
        }
        // An executor that couldn't serve the kind (model off, unknown kind) answers with
        // `error` set — normalize to nil so every caller reads one "no result" shape.
        if let result, result.error != nil { return nil }
        return result
    }

    /// The answer a worker sent back (`ClientMsg.intelligenceResult` → the bridge). Correlated
    /// by request id; a stale id (deadline already fired, job re-queued elsewhere) or an answer
    /// from a device the job wasn't dispatched to is dropped.
    public func receiveResult(_ result: WireIntelligenceResult, from deviceID: String) {
        guard let entry = inFlight[result.id], entry.deviceID == deviceID else { return }
        inFlight[result.id] = nil
        deadlines[result.id]?.cancel()
        deadlines[result.id] = nil
        workers[deviceID]?.inFlight.remove(result.id)
        entry.job.continuation.resume(returning: result)
        dispatch()
    }

    /// Whether any connected worker may run jobs of `priority` (ignoring current load).
    public func hasEligibleWorker(for priority: IntelligencePriority?) -> Bool {
        workers.values.contains { eligible($0, for: priority) }
    }

    /// Connected worker count, for logs/inspection.
    public var workerCount: Int { workers.count }

    // MARK: - Scheduling

    private func eligible(_ worker: Worker, for priority: IntelligencePriority?) -> Bool {
        (priority ?? .background).allowsMobileExecution
            || worker.profile.deviceClass == .desktop
    }

    /// Hand every pending job it can place to a worker: jobs highest-priority-first (FIFO
    /// within a tier), each onto the most powerful idle eligible device — which is exactly
    /// what biases the hottest tiers toward the strongest chips.
    private func dispatch() {
        guard !pending.isEmpty else { return }
        pending.sort {
            let l = ($0.request.priority ?? .background).rank
            let r = ($1.request.priority ?? .background).rank
            return l != r ? l > r : $0.seq < $1.seq
        }
        var index = 0
        while index < pending.count {
            let job = pending[index]
            guard let choice = bestIdleWorker(for: job.request.priority) else {
                index += 1  // no slot for this tier right now; a lower tier may still fit a phone
                continue
            }
            pending.remove(at: index)
            var dispatched = job
            dispatched.attempts += 1
            inFlight[job.request.id] = (dispatched, choice.deviceID)
            workers[choice.deviceID]?.inFlight.insert(job.request.id)
            let send = choice.send
            let request = job.request
            Task { await send(request) }
        }
    }

    private func bestIdleWorker(for priority: IntelligencePriority?) -> Worker? {
        workers.values
            .filter { eligible($0, for: priority) && $0.inFlight.count < perWorkerCap }
            .max {
                $0.profile.powerRank != $1.profile.powerRank
                    ? $0.profile.powerRank < $1.profile.powerRank
                    : $0.deviceID > $1.deviceID  // deterministic tie-break
            }
    }

    /// Put a job whose worker vanished back in line — or give up (resolve nil) once its
    /// dispatch attempts are spent, so a flapping device can't hold a summary hostage.
    private func requeue(_ requestID: String) {
        guard let entry = inFlight.removeValue(forKey: requestID) else { return }
        if entry.job.attempts >= maxAttempts {
            deadlines[requestID]?.cancel()
            deadlines[requestID] = nil
            entry.job.continuation.resume(returning: nil)
        } else {
            pending.append(entry.job)  // the deadline keeps running — it bounds total wait
        }
    }

    private func armDeadline(for request: WireIntelligenceRequest) {
        let seconds = request.deadlineSeconds ?? defaultDeadline
        deadlines[request.id] = Task { [weak self] in
            try? await Task.sleep(for: .seconds(seconds))
            guard !Task.isCancelled else { return }
            await self?.expire(request.id)
        }
    }

    /// The deadline lapsed: resolve nil wherever the job currently is. A result that arrives
    /// later is dropped by ``receiveResult``'s correlation guard; the worker's slot is freed so
    /// the straggler doesn't wedge future placement.
    private func expire(_ requestID: String) {
        deadlines[requestID] = nil
        if let index = pending.firstIndex(where: { $0.request.id == requestID }) {
            let job = pending.remove(at: index)
            job.continuation.resume(returning: nil)
            return
        }
        if let entry = inFlight.removeValue(forKey: requestID) {
            workers[entry.deviceID]?.inFlight.remove(requestID)
            entry.job.continuation.resume(returning: nil)
            dispatch()
        }
    }
}
