import Foundation

/// Serializes on-device Apple Foundation Model (AFM) requests through a small
/// concurrency limit, so a burst of soft-AI work can't saturate the Neural Engine and
/// stall the rest of the app. The on-device model is a single shared device resource —
/// firing several generations at once only adds contention and latency, never speed —
/// so every provider routes through the one ``shared`` queue.
///
/// Waiting requests are released in **priority** order: a generation tied to something
/// the user is watching right now (naming the chat they just started) always jumps
/// ahead of queued background work (triaging the backlog). Ties break FIFO, so same-
/// priority work still runs in arrival order. A higher priority only preempts *waiting*
/// work — it never interrupts a request already in flight, since AFM calls aren't
/// cancellable mid-generation.
public actor AFMRequestQueue {
    /// Which user experience a request serves, highest priority first. The provider maps
    /// each soft-AI feature onto one of these so the queue can order a backlog of work.
    public enum Priority: Int, Comparable, Sendable {
        /// Background backlog work the user isn't actively watching: backlog triage,
        /// to-do gisting, and the non-Bash collapsed tool-call summaries (Grep/Glob folds).
        case background = 0
        /// Wrapping up a session: classifying whether a finished turn awaits the user,
        /// and writing its glanceable summary card.
        case completion = 1
        /// Collapsed Bash tool-call summaries (the map/reduce gist of a run's commands):
        /// they fold into a line the user reads in the open transcript, so they outrank
        /// session wrap-up — but still yield to a live chat name.
        case bashSummary = 2
        /// Directly tied to a fresh user action and visible immediately: naming a chat
        /// from its first message.
        case interactive = 3

        public static func < (lhs: Priority, rhs: Priority) -> Bool {
            lhs.rawValue < rhs.rawValue
        }

        /// A short human label for the queue inspector ("Running", "Completing"…).
        public var label: String {
            switch self {
            case .background: "Background"
            case .completion: "Completion"
            case .bashSummary: "Bash summary"
            case .interactive: "Interactive"
            }
        }
    }

    /// The on-device queue every provider shares for local model work. The Neural Engine is a
    /// single device resource, so concurrency is capped at one — a burst of soft-AI work lines
    /// up here instead of contending for the chip.
    public static let shared = AFMRequestQueue(maxConcurrent: 1, isCloud: false)

    /// The Private Cloud Compute queue. PCC runs on Apple's servers, not this device, so it
    /// isn't resource-constrained the way the Neural Engine is — many requests can be in flight
    /// at once. Work still routes through a queue so the inspector can show what the cloud is
    /// generating and what it just finished; with an effectively unbounded concurrency cap it
    /// fires requests concurrently and never makes anything wait — except while frozen for
    /// inspection, where new requests line up as waiters and resume drains them all at once.
    public static let sharedCloud = AFMRequestQueue(maxConcurrent: .max, isCloud: true)

    private let maxConcurrent: Int
    /// Whether this queue's work runs in Private Cloud Compute (Apple's servers) rather than on
    /// this device's Neural Engine. Stamped into every ``Request`` so the merged inspector can
    /// flag PCC rows with a cloud glyph; it doesn't affect ordering or concurrency.
    private let isCloud: Bool
    /// Monotonic ticket, stamped on each request to break same-priority ties FIFO and to
    /// identify a request when it's released or cancelled.
    private var nextSeq: UInt64 = 0

    /// One pending or in-flight request: its priority, a human-readable description of
    /// the work (for the queue inspector), and its FIFO ticket.
    private struct Entry {
        let seq: UInt64
        let priority: Priority
        let label: String
        /// A richer, human-readable view of the actual work — typically a preview of the
        /// prompt the model is processing — surfaced when a queue row is expanded.
        let detail: String
        /// Tags a class of redundant background work where only the newest still-waiting
        /// instance is worth running: enqueuing one supersedes any earlier waiter sharing
        /// the key (re-triaging the whole backlog covers what a queued triage of an older
        /// snapshot would have). nil for one-off work that never coalesces.
        let coalesceKey: String?
        /// When the request entered the queue.
        let enqueuedAt: Date
        /// When it was handed a slot and began generating; nil while still waiting. Drives
        /// the "running for…" / "waiting…" timers in the inspector.
        var startedAt: Date?
        /// When the request finished and gave its slot back; nil while running or waiting.
        /// Set only on entries moved into ``historyEntries``.
        var finishedAt: Date?
        /// The text the request produced, captured on release for the inspector's expanded
        /// detail — what the model actually said (the Bash summary line). nil while running or
        /// waiting, and for work that returns no text (structured triage).
        var result: String?
    }
    /// Requests currently holding a slot — executing, or just handed one and about to
    /// start. A queued (still-waiting) request is NOT here; it's in `waiters`.
    private var runningEntries: [Entry] = []

    private struct Waiter {
        let entry: Entry
        /// Resumes the suspended `acquire`. `true` hands over a slot (the waiter may
        /// run); `false` means it was cancelled before getting one (no slot to release).
        let resume: @Sendable (Bool) -> Void
    }
    private var waiters: [Waiter] = []

    /// A bounded, most-recent-last log of requests that have finished, so the inspector can
    /// show what the model just worked on (the queue moves fast and rows vanish the instant
    /// they complete). Capped so a long-running app can't grow this without bound.
    private var historyEntries: [Entry] = []
    private let maxHistory = 50

    /// When `true`, the queue stops handing slots to waiting requests: nothing new starts,
    /// so the waiting list holds still while the user inspects it in the sidebar. A request
    /// already in flight isn't interrupted (AFM can't be cancelled mid-generation) — it
    /// finishes and moves to history, then its slot sits idle until the freeze lifts. New
    /// requests still enqueue while frozen; they just wait their turn like everything else.
    private var frozen = false

    /// `maxConcurrent` defaults to 1: the on-device model effectively serializes on the
    /// Neural Engine anyway, so one in-flight request at a time is the honest cap.
    public init(maxConcurrent: Int = 1, isCloud: Bool = false) {
        self.maxConcurrent = max(1, maxConcurrent)
        self.isCloud = isCloud
    }

    /// Runs `operation` once a slot is free, with waiting requests ordered by `priority`.
    /// The operation itself runs outside the queue's isolation, so a long generation
    /// never blocks the bookkeeping here — only slot accounting is serialized. `label`
    /// describes the work for the queue inspector (e.g. "Chat name", "Backlog triage").
    ///
    /// `resultText`, when given, derives the request's output from its result once it returns
    /// — for text work (a Bash summary), the model's own reply — and stores it on the request's
    /// history entry, so the inspector's expanded detail can show what the model produced. A
    /// nil/blank return leaves the entry with no recorded output.
    public func run<T: Sendable>(
        priority: Priority, label: String = "", detail: String = "",
        resultText: (@Sendable (T) -> String?)? = nil,
        _ operation: @Sendable () async -> T
    ) async -> T {
        let seq = await acquire(priority: priority, label: label, detail: detail)
        let result = await operation()
        // Only release a slot we were actually handed — a cancelled waiter never held one.
        if let seq { release(seq, result: resultText?(result)) }
        return result
    }

    /// Like ``run``, but coalesces redundant background work tagged with `coalesceKey` — a
    /// class where only the newest still-waiting instance is worth running because it
    /// supersedes the rest (re-triaging the whole current backlog covers everything an
    /// earlier, still-queued triage of an older snapshot would have). Enqueuing here cancels
    /// every *waiting* request sharing the key — each returns `nil` without running its
    /// operation — so a burst of identical background work collapses to a single pass. A
    /// request already in flight is never touched (AFM isn't cancellable mid-generation, and
    /// it's about to finish anyway). Returns the operation's result, or `nil` if this request
    /// was itself superseded by a newer one before it was handed a slot.
    public func runCoalescing<T: Sendable>(
        priority: Priority, coalesceKey: String, label: String = "", detail: String = "",
        resultText: (@Sendable (T) -> String?)? = nil,
        _ operation: @Sendable () async -> T
    ) async -> T? {
        guard let seq = await acquire(priority: priority, label: label, detail: detail,
                                      coalesceKey: coalesceKey) else { return nil }
        let result = await operation()
        release(seq, result: resultText?(result))
        return result
    }

    /// Suspends until a slot is free (immediately if one is). Returns the request's ticket
    /// when a slot was granted, `nil` if the task was cancelled while waiting (so the
    /// caller must not release). Cancellation-safe: a waiting request cancelled mid-queue
    /// is pulled out and resumed rather than left to leak its continuation.
    private func acquire(priority: Priority, label: String, detail: String,
                         coalesceKey: String? = nil) async -> UInt64? {
        // The newest keyed request wins: drop any earlier still-waiting work of the same
        // class, since this one supersedes it. (A request already holding a slot keeps it —
        // AFM can't be cancelled mid-generation.) A free slot means there are no waiters to
        // supersede, so this only ever fires on the waiting path.
        if let coalesceKey { supersedeWaiters(coalesceKey: coalesceKey) }
        let seq = nextSeq
        nextSeq += 1
        let now = Date()
        var entry = Entry(seq: seq, priority: priority, label: label, detail: detail,
                          coalesceKey: coalesceKey, enqueuedAt: now, startedAt: nil)
        // A frozen queue grants no slots, even an idle one — the request queues so the
        // user sees it line up without anything starting to run.
        if !frozen && runningEntries.count < maxConcurrent {
            entry.startedAt = now
            runningEntries.append(entry)
            return seq
        }
        return await withTaskCancellationHandler {
            await withCheckedContinuation { (continuation: CheckedContinuation<UInt64?, Never>) in
                // Already cancelled before we suspended — don't enqueue at all. (Actor
                // isolation guarantees this body runs to completion before the onCancel
                // task below can touch `waiters`, so there's no append-after-cancel gap.)
                if Task.isCancelled {
                    continuation.resume(returning: nil)
                } else {
                    waiters.append(Waiter(entry: entry,
                                          resume: { continuation.resume(returning: $0 ? seq : nil) }))
                }
            }
        } onCancel: {
            Task { await self.cancelWaiter(seq) }
        }
    }

    /// Pulls a still-waiting request out of the queue on cancellation and resumes it with
    /// "no slot". A no-op if it has already been granted a slot (it'll run and release
    /// normally; the operation can observe cancellation itself).
    private func cancelWaiter(_ seq: UInt64) {
        guard let index = waiters.firstIndex(where: { $0.entry.seq == seq }) else { return }
        let waiter = waiters.remove(at: index)
        waiter.resume(false)
    }

    /// Cancels every still-waiting request tagged with `coalesceKey`, resuming each with
    /// "no slot" so its ``runCoalescing`` returns nil and skips its operation. Called when a
    /// newer request of the same class arrives and renders the earlier waiters redundant.
    /// Only waiters are affected — a running request is mid-generation and can't be undone.
    private func supersedeWaiters(coalesceKey: String) {
        var kept: [Waiter] = []
        var superseded: [Waiter] = []
        for waiter in waiters {
            if waiter.entry.coalesceKey == coalesceKey { superseded.append(waiter) }
            else { kept.append(waiter) }
        }
        guard !superseded.isEmpty else { return }
        waiters = kept
        for waiter in superseded { waiter.resume(false) }
    }

    /// Frees the slot held by `seq`, then hands it to the highest-priority waiter (or lets
    /// the slot go idle if none wait). Removing then re-adding keeps the running count at
    /// `maxConcurrent` when a waiter is promoted, and drops it by one when none is.
    ///
    /// `result`, when non-blank, is recorded on the finished entry before it lands in history
    /// — so a text request (a Bash summary) carries the model's own output into the inspector's
    /// expanded detail. Blank/whitespace records nothing.
    private func release(_ seq: UInt64, result: String? = nil) {
        if let index = runningEntries.firstIndex(where: { $0.seq == seq }) {
            var finished = runningEntries.remove(at: index)
            finished.finishedAt = Date()
            if let result,
               !result.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
                finished.result = result
            }
            historyEntries.append(finished)
            if historyEntries.count > maxHistory {
                historyEntries.removeFirst(historyEntries.count - maxHistory)
            }
        }
        // While frozen, leave the freed slot idle so the waiting queue holds still for
        // inspection; lifting the freeze fills it again (see ``setFrozen``).
        if !frozen { promoteNextWaiter() }
    }

    /// Promotes the highest-priority waiter into a free slot, stamping the moment it starts
    /// running and resuming it. No-op when every slot is taken or no one is waiting; returns
    /// whether a waiter was promoted so ``setFrozen`` can fill all idle slots on resume.
    @discardableResult
    private func promoteNextWaiter() -> Bool {
        guard runningEntries.count < maxConcurrent, let index = nextWaiterIndex() else { return false }
        let waiter = waiters.remove(at: index)
        var promoted = waiter.entry  // promote the waiter into the freed slot…
        promoted.startedAt = Date()  // …and stamp the moment it starts running
        runningEntries.append(promoted)
        waiter.resume(true)
        return true
    }

    /// Freezes or resumes slot hand-off. While frozen, no waiting request is started, so the
    /// queue the user is inspecting stops moving. Lifting the freeze immediately fills every
    /// idle slot from the waiting queue, so processing picks up right where it left off.
    public func setFrozen(_ value: Bool) {
        guard frozen != value else { return }
        frozen = value
        if !frozen { while promoteNextWaiter() {} }
    }

    /// Whether slot hand-off is currently frozen.
    public var isFrozen: Bool { frozen }

    /// The index of the next waiter to run: highest priority, ties broken by lowest seq
    /// (earliest arrival).
    private func nextWaiterIndex() -> Int? {
        guard !waiters.isEmpty else { return nil }
        var best = waiters.startIndex
        for index in waiters.indices where index != best {
            let candidate = waiters[index].entry
            let leader = waiters[best].entry
            if candidate.priority > leader.priority
                || (candidate.priority == leader.priority && candidate.seq < leader.seq) {
                best = index
            }
        }
        return best
    }

    // MARK: - Inspection

    /// A point-in-time view of the queue for the inspector UI: what's running now and
    /// what's waiting (already ordered by release order — highest priority first, ties
    /// FIFO). One queue is on-device, one is Private Cloud Compute; the inspector merges
    /// both and marks each request by the backend it ran on (see ``Request/isCloud``).
    public struct Snapshot: Sendable {
        public var running: [Request]
        public var waiting: [Request]
        /// Recently finished requests, most-recent first. Surfaced as a grayed-out history
        /// so the user can see what the model just did; doesn't count toward ``isEmpty``,
        /// which tracks only live (running or waiting) work.
        public var history: [Request]
        /// Whether slot hand-off is frozen right now (see ``setFrozen``), so the inspector
        /// can reflect the paused state and offer to resume.
        public var frozen: Bool
        public var isEmpty: Bool { running.isEmpty && waiting.isEmpty }

        public init(running: [Request], waiting: [Request], history: [Request] = [],
                    frozen: Bool = false) {
            self.running = running
            self.waiting = waiting
            self.history = history
            self.frozen = frozen
        }
    }

    /// One request in a ``Snapshot``: its priority, the work's short label, a richer
    /// detail string (the prompt preview, for the expanded inspector), and the timestamps
    /// that drive its live "running for…" / "waiting…" timers.
    public struct Request: Sendable, Identifiable {
        public let id: UInt64
        public let priority: Priority
        public let label: String
        public let detail: String
        public let enqueuedAt: Date
        public let startedAt: Date?
        /// When the request finished; non-nil only for history rows.
        public let finishedAt: Date?
        /// The text the request produced — the model's reply — shown in the expanded detail
        /// when a finished row is clicked. nil while running/waiting and for non-text work.
        public let result: String?
        /// Whether this request ran in Private Cloud Compute rather than on-device — drives
        /// the cloud glyph the merged inspector puts on PCC rows.
        public let isCloud: Bool

        public init(id: UInt64, priority: Priority, label: String, detail: String,
                    enqueuedAt: Date, startedAt: Date?, finishedAt: Date? = nil,
                    result: String? = nil, isCloud: Bool = false) {
            self.id = id
            self.priority = priority
            self.label = label
            self.detail = detail
            self.enqueuedAt = enqueuedAt
            self.startedAt = startedAt
            self.finishedAt = finishedAt
            self.result = result
            self.isCloud = isCloud
        }

        /// A key unique across both queues. The per-queue `id` alone collides between the
        /// on-device and cloud queues (they share no counter), so the merged inspector keys
        /// row identity and expansion on this instead.
        public var traceID: String { "\(isCloud ? "pcc" : "afm"):\(id)" }

        /// Seconds since the request entered the queue.
        public func queuedFor(now: Date) -> TimeInterval { max(0, now.timeIntervalSince(enqueuedAt)) }

        /// Seconds since it started generating, or nil if it hasn't started yet.
        public func runningFor(now: Date) -> TimeInterval? {
            startedAt.map { max(0, now.timeIntervalSince($0)) }
        }

        /// How long the request ran end to end, or nil if it never started or hasn't
        /// finished. Drives the "ran for…" line on history rows.
        public var ranFor: TimeInterval? {
            guard let startedAt, let finishedAt else { return nil }
            return max(0, finishedAt.timeIntervalSince(startedAt))
        }
    }

    /// A snapshot of the queue right now — running requests in start order, waiting ones
    /// in the order they'll be released, and recently finished ones most-recent first.
    public func snapshot() -> Snapshot {
        func request(_ e: Entry) -> Request {
            Request(id: e.seq, priority: e.priority, label: e.label, detail: e.detail,
                    enqueuedAt: e.enqueuedAt, startedAt: e.startedAt, finishedAt: e.finishedAt,
                    result: e.result, isCloud: isCloud)
        }
        let running = runningEntries.map(request)
        let ordered = waiters.map(\.entry).sorted {
            $0.priority != $1.priority ? $0.priority > $1.priority : $0.seq < $1.seq
        }
        let waiting = ordered.map(request)
        let history = historyEntries.reversed().map(request)  // most-recent first
        return Snapshot(running: running, waiting: waiting, history: history, frozen: frozen)
    }
}
