import Foundation
import os

/// Lock-lifecycle tracing. Every grant/release/queue is logged so a stuck lock can be diagnosed
/// from `log stream --predicate 'subsystem == "com.nucleic"'` (or Console.app, category "locks").
/// Messages are prefixed `LOCKTRACE` for easy grep; session ids and paths are `.public` (a local
/// project's own data, not secrets). Cheap — these events are infrequent.
let lockLog = Logger(subsystem: "com.nucleic", category: "locks")

/// The top-level project branch a session's lineage flows to (its `rootRef`) — the key for a
/// **lock domain** (LOCKING §2). Every session sharing a domain contends on files; sessions in
/// different domains never do.
public struct LockDomain: Hashable, Sendable {
    public let ref: String
    public init(_ ref: String) { self.ref = ref }
}

/// What a blocked `acquire` resumes with (LOCKING §4.2).
public enum AcquireOutcome: Sendable, Equatable {
    /// The session holds every requested file — the agent may edit.
    case granted
    /// Granted, but update-on-grant pulled a parent change into a file this edit targets, so the
    /// agent's in-context view is stale: the host blocks this one edit and feeds the agent `diff`
    /// (scoped to `files`) so it re-grounds and re-issues. The lock is held, so the retry proceeds
    /// (LOCKING §4.5). `diff` may be empty when the change was too large to inline — re-read instead.
    case grantedNeedsReground(files: [String], diff: String)
    /// The asking session went away / was reaped while queued — its call unblocks without a lock.
    case cancelled
    /// Granted, but updating the worktree to the parent (update-on-grant) hit a conflict; the
    /// session keeps its locks and is paused for resolution (LOCKING §4.5).
    case updateConflict([String])
}

/// Result of update-on-grant — bring the granted session's worktree up to `<parentRef>`
/// (LOCKING §4.5). Supplied by the host via an injected collaborator.
public enum LockUpdateResult: Sendable, Equatable {
    /// Up to date, or nothing relevant moved — the agent may proceed unchanged.
    case clean
    /// The parent moved a file this edit targets: `files` (and the scoped `diff`) tell the agent
    /// what changed underneath it so it re-grounds before editing.
    case changed(files: [String], diff: String)
    case conflicted([String])
}

/// A lock-lifecycle moment the host can surface in a session's chat log (the same trio the
/// `LOCKTRACE` os_log records, but in-chat). Paths are repo-relative and sorted.
public enum LockEvent: Sendable, Equatable {
    /// Queued — blocked waiting for these paths, which another session holds.
    case waiting([String])
    /// Newly granted these paths — the session now holds them and may edit.
    case acquired([String])
    /// Released these paths.
    case released([String])
}

/// A point-in-time view of one domain's holders and waiters, for the UI (LOCKING §8). Raw
/// (session ids + paths); the host enriches it with titles into a `LockQueueSnapshot`.
public struct DomainLockState: Sendable, Equatable {
    public struct Holder: Sendable, Equatable {
        public let session: SessionID
        public let paths: [String]
    }
    public struct Waiter: Sendable, Equatable {
        public let session: SessionID
        public let delta: [String]
    }
    public let holders: [Holder]
    public let waiters: [Waiter]
}

/// A persisted lock record — domain + path + holder — used to rebuild held locks at launch
/// (LOCKING §6). Queue order is *not* persisted: it's transient (a restart's agents are gone and
/// re-acquire on their next edit), and held ownership is reconstructable from git anyway.
public struct PersistedLock: Sendable, Equatable {
    public let domain: String
    public let path: String
    public let holder: SessionID
    public init(domain: String, path: String, holder: SessionID) {
        self.domain = domain
        self.path = path
        self.holder = holder
    }
}

/// Authoritative, in-memory file-lock system (LOCKING §4). Replaces the emergent
/// `unmergedFiles`-derived "lock": locks are explicit records keyed by `(domain, path)`, scoped
/// to a session's root tree (`rootRef`), so every session in one tree contends on a file and at
/// most one holds it at a time. Acquisition is **all-or-nothing** with a fair, demoting waiter
/// queue. Release is driven by the host (a mediated merge, lifecycle, the manual button) and —
/// once wired (Phase C) — a git-landed poll. An `actor`, so every mutation is serialized: no two
/// `acquire`s can both observe a file as free and both take it.
public actor LockManager {
    // MARK: Injected collaborators (optional; tests/host supply fakes/real ones)

    /// Whether a queued session is gone/terminal and should be reaped (resumed `.cancelled`) so
    /// its `acquire` call doesn't hang forever. Unset ⇒ never reap (the host drives removal).
    private var reapCheck: (@Sendable (SessionID) async -> Bool)?
    /// Update-on-grant collaborator (LOCKING §4.5): bring a granted session's worktree up to its
    /// parent. Passed the session and the files it is acquiring (so the merge result can be scoped
    /// to them). Unset ⇒ grants resume immediately (pure mechanics, for unit tests).
    private var updateOnGrant: (@Sendable (SessionID, [String]) async -> LockUpdateResult)?
    /// Host hook to record lock-lifecycle events in a session's chat log. Unset ⇒ no logging
    /// (the default for unit tests). Fired fire-and-forget so it never blocks the lock table.
    private var onLockEvent: (@Sendable (SessionID, LockEvent) async -> Void)?

    public init() {}

    public func setReapCheck(_ f: @escaping @Sendable (SessionID) async -> Bool) { reapCheck = f }
    public func setUpdateOnGrant(_ f: @escaping @Sendable (SessionID, [String]) async -> LockUpdateResult) { updateOnGrant = f }
    public func setLockEventHandler(_ f: @escaping @Sendable (SessionID, LockEvent) async -> Void) { onLockEvent = f }

    /// Surface a lock-lifecycle event to the host without suspending the actor: spawned as an
    /// actor-isolated task so it runs after the current mutation, never interleaving an `await`
    /// into a critical section (the table stays consistent and a release can't be missed).
    private func emitLockEvent(_ session: SessionID, _ event: LockEvent) {
        guard let onLockEvent else { return }
        Task { await onLockEvent(session, event) }
    }

    // MARK: State (LOCKING §4.1)

    struct LockKey: Hashable { let domain: LockDomain; let path: String }

    private struct Waiter {
        let session: SessionID
        let domain: LockDomain
        /// The full new set it needs, granted atomically (all-or-nothing). While queued the
        /// session holds only what it *already* held — never a new file (LOCKING §4.2).
        let delta: Set<String>
        let continuation: CheckedContinuation<AcquireOutcome, Never>
    }

    /// `(domain, path) → holder`. The authoritative lock table.
    private var holderOf: [LockKey: SessionID] = [:]
    /// Reverse index for O(1) release.
    private var heldBySession: [SessionID: Set<LockKey>] = [:]
    /// FIFO waiter queue per domain, maintained with one-place demotion (LOCKING §4.3).
    private var queues: [LockDomain: [Waiter]] = [:]
    /// Producer side of the activity stream the host's release-reconcile loop parks on. Yields
    /// on every grant/enqueue so the loop wakes the instant there's something to watch, and stays
    /// parked (no polling) while nothing is held (LOCKING §4.4). nil until the host subscribes.
    private var reconcileSignal: AsyncStream<Void>.Continuation?

    // MARK: Acquisition (LOCKING §4.2)

    /// Acquire `files` for `session` in `domain` (= its `rootRef`). Returns `.granted` once the
    /// session holds them all, or `.cancelled` if it was reaped while queued.
    ///
    /// All-or-nothing: a session with no locks yet takes its whole requested set or nothing — it
    /// never holds a subset while blocked. A session that already holds locks keeps them and
    /// queues for the full new set, so it never holds a *new* file while blocked on another
    /// ("hold edited, queue for delta").
    public func acquire(_ session: SessionID, domain: LockDomain, files: [String]) async -> AcquireOutcome {
        let paths = Self.normalize(files)
        let mine = Set((heldBySession[session] ?? []).map(\.path))
        let needed = paths.subtracting(mine)
        if needed.isEmpty {
            lockLog.info("LOCKTRACE acquire noop session=\(session.rawValue, privacy: .public) already-holds \(paths.sorted().joined(separator: ","), privacy: .public)")
            return .granted          // already holds everything requested
        }
        if isFree(needed, in: domain, for: session) {
            grant(needed, to: session, in: domain)
            // Grant is recorded (holderOf) synchronously BEFORE this await, so a concurrent acquire
            // for an overlapping file sees it held and queues — the merge can't cause a double-grant.
            // Immediate grants run update-on-grant too (LOCKING §4.5): the parent may have advanced
            // even with no waiting, so a freshly-granted session is never left on a stale parent.
            let outcome = await completeGrant(session, files: needed)
            lockLog.info("LOCKTRACE acquire granted session=\(session.rawValue, privacy: .public) domain=\(domain.ref, privacy: .public) files=\(needed.sorted().joined(separator: ","), privacy: .public)")
            emitLockEvent(session, .acquired(needed.sorted()))
            return outcome
        }
        lockLog.info("LOCKTRACE acquire queued session=\(session.rawValue, privacy: .public) domain=\(domain.ref, privacy: .public) delta=\(needed.sorted().joined(separator: ","), privacy: .public)")
        return await withCheckedContinuation { cont in
            queues[domain, default: []].append(
                Waiter(session: session, domain: domain, delta: needed, continuation: cont))
            signalActivity()  // a waiter now exists — keep the reconcile loop alive to grant it
            emitLockEvent(session, .waiting(needed.sorted()))  // logged after enqueue, never before
        }
    }

    /// Whether `acquire(session, domain, files)` would block right now (a needed file is held by
    /// another). Lets the host set its "waiting" indicator before the suspending `acquire`.
    public func wouldBlock(_ session: SessionID, domain: LockDomain, files: [String]) -> Bool {
        let needed = Self.normalize(files).subtracting(Set((heldBySession[session] ?? []).map(\.path)))
        return !needed.isEmpty && !isFree(needed, in: domain, for: session)
    }

    // MARK: Release (LOCKING §4.4)

    /// Release specific paths held by `session`, then re-poll so waiters advance.
    public func release(_ session: SessionID, paths: [String]) async {
        let norm = Set(Self.normalize(paths))
        guard let keys = heldBySession[session] else { return }
        let dropped = keys.filter { norm.contains($0.path) }
        guard !dropped.isEmpty else { return }
        var remaining = keys
        for k in dropped { holderOf[k] = nil; remaining.remove(k) }
        heldBySession[session] = remaining.isEmpty ? nil : remaining
        lockLog.notice("LOCKTRACE release session=\(session.rawValue, privacy: .public) paths=\(dropped.map(\.path).sorted().joined(separator: ","), privacy: .public) remaining=\(remaining.count)")
        emitLockEvent(session, .released(dropped.map(\.path).sorted()))
        for d in Set(dropped.map(\.domain)) { await poll(d) }
    }

    /// Release everything `session` holds and drop it from any queue (resuming `.cancelled` if it
    /// was waiting). The mediated-merge / lifecycle / manual-button path.
    public func releaseAll(_ session: SessionID) async {
        var domains = Set<LockDomain>()
        if let keys = heldBySession[session] {
            lockLog.notice("LOCKTRACE releaseAll session=\(session.rawValue, privacy: .public) paths=\(keys.map(\.path).sorted().joined(separator: ","), privacy: .public)")
            for k in keys { holderOf[k] = nil; domains.insert(k.domain) }
            heldBySession[session] = nil
            emitLockEvent(session, .released(keys.map(\.path).sorted()))
        } else {
            lockLog.notice("LOCKTRACE releaseAll session=\(session.rawValue, privacy: .public) held-nothing")
        }
        for (d, q) in queues {
            guard q.contains(where: { $0.session == session }) else { continue }
            var survivors: [Waiter] = []
            for w in q where w.session != session { survivors.append(w) }
            for w in q where w.session == session { w.continuation.resume(returning: .cancelled) }
            queues[d] = survivors.isEmpty ? nil : survivors
            domains.insert(d)
        }
        for d in domains { await poll(d) }
    }

    /// Seed held locks at launch (LOCKING §6) from the host's reconstruction (each live session's
    /// not-yet-landed files in its domain). Ownership only; grants each path to the first claimant
    /// and skips a path already held, so a stale double-claim can't overwrite.
    public func restore(_ locks: [PersistedLock]) {
        for l in locks {
            let key = LockKey(domain: LockDomain(l.domain), path: l.path)
            if holderOf[key] == nil { grant([l.path], to: l.holder, in: LockDomain(l.domain)) }
        }
    }

    // MARK: Queue maintenance (LOCKING §4.3)

    /// One scheduling pass over a domain: reap dead waiters, grant every waiter whose full delta
    /// is free (readiness — not position — gates a grant, so a ready waiter is never starved
    /// behind a blocked one), reserving granted paths within the pass so two waiters can't take
    /// the same file. A still-blocked head is demoted one place. Runs on every release and may be
    /// called on a timer backstop.
    public func poll(_ domain: LockDomain) async {
        await reapDeadWaiters(in: domain)
        let granted = grantReadyWaiters(in: domain)
        for w in granted {
            let outcome = await completeGrant(w.session, files: w.delta)
            // It now holds its delta — `.granted`, `.grantedNeedsReground`, or kept-on-update-
            // conflict (LOCKING §4.5).
            emitLockEvent(w.session, .acquired(w.delta.sorted()))
            w.continuation.resume(returning: outcome)
        }
    }

    /// Run update-on-grant for a just-granted session (LOCKING §4.5): bring its worktree up to the
    /// parent so it sees the latest. Clean → `.granted`. Parent moved a file this edit targets →
    /// `.grantedNeedsReground` (lock kept; the host re-grounds the agent). Conflict → keep the lock
    /// and tell the caller (`.updateConflict`). No hook set → `.granted` (the pure-mechanics default,
    /// so the LockManager unit-tests exercise the queue without any git). `files` is the acquire
    /// delta — the files being granted — so the merge result can be scoped to them.
    private func completeGrant(_ session: SessionID, files: Set<String>) async -> AcquireOutcome {
        guard let updateOnGrant else { return .granted }
        switch await updateOnGrant(session, files.sorted()) {
        case .clean: return .granted
        case .changed(let f, let d): return .grantedNeedsReground(files: f, diff: d)
        case .conflicted(let paths): return .updateConflict(paths)
        }
    }

    private func reapDeadWaiters(in domain: LockDomain) async {
        guard let reapCheck, let q = queues[domain] else { return }
        var survivors: [Waiter] = []
        for w in q {
            if await reapCheck(w.session) {
                w.continuation.resume(returning: .cancelled)
            } else {
                survivors.append(w)
            }
        }
        queues[domain] = survivors.isEmpty ? nil : survivors
    }

    /// Synchronously decide and record grants for `domain`, returning the granted waiters to
    /// resume. (Resuming is done by the caller, after the table mutation, so the queue isn't
    /// mutated mid-iteration.)
    private func grantReadyWaiters(in domain: LockDomain) -> [Waiter] {
        guard let q = queues[domain], !q.isEmpty else { return [] }
        var reserved = Set<String>()
        var granted: [Waiter] = []
        var remaining: [Waiter] = []
        for w in q {
            if isFree(w.delta, in: domain, for: w.session, alsoBlockedBy: reserved) {
                grant(w.delta, to: w.session, in: domain)
                reserved.formUnion(w.delta)
                granted.append(w)
            } else {
                remaining.append(w)
            }
        }
        // "Demote by a place": a still-blocked next-up yields so the following waiter gets a turn.
        if remaining.count > 1, let head = remaining.first,
           !isFree(head.delta, in: domain, for: head.session) {
            remaining.removeFirst()
            remaining.insert(head, at: 1)
        }
        queues[domain] = remaining.isEmpty ? nil : remaining
        return granted
    }

    // MARK: Introspection (LOCKING §8)

    /// Raw holders + waiters for `domain`, for the host to enrich into a `LockQueueSnapshot`.
    public func state(in domain: LockDomain) -> DomainLockState {
        var byHolder: [SessionID: [String]] = [:]
        for (key, holder) in holderOf where key.domain == domain {
            byHolder[holder, default: []].append(key.path)
        }
        let holders = byHolder
            .map { DomainLockState.Holder(session: $0.key, paths: $0.value.sorted()) }
            .sorted { $0.session.rawValue < $1.session.rawValue }
        let waiters = (queues[domain] ?? [])
            .map { DomainLockState.Waiter(session: $0.session, delta: $0.delta.sorted()) }
        return DomainLockState(holders: holders, waiters: waiters)
    }

    /// Whether `session` is currently queued anywhere (drives the "waiting" UI indicator).
    public func isWaiting(_ session: SessionID) -> Bool {
        queues.values.contains { q in q.contains { $0.session == session } }
    }

    /// Repo-relative paths `session` currently holds (any domain), sorted.
    public func heldPaths(_ session: SessionID) -> [String] {
        (heldBySession[session] ?? []).map(\.path).sorted()
    }

    public struct RawHold: Sendable, Equatable {
        public let session: SessionID
        public let domain: LockDomain
        public let path: String
    }
    public struct RawWait: Sendable, Equatable {
        public let session: SessionID
        public let domain: LockDomain
        public let delta: [String]
    }
    /// All holds and waits across every domain, for the host's lock-queue viewer (LOCKING §8).
    public func rawState() -> (holds: [RawHold], waits: [RawWait]) {
        let holds = holderOf.map { RawHold(session: $0.value, domain: $0.key.domain, path: $0.key.path) }
        let waits = queues.flatMap { domain, q in
            q.map { RawWait(session: $0.session, domain: domain, delta: $0.delta.sorted()) }
        }
        return (holds, waits)
    }
    /// Every session currently holding at least one lock (for the host's release-reconcile poll).
    public func sessionsHoldingLocks() -> [SessionID] { Array(heldBySession.keys) }

    // MARK: Reconcile-loop liveness (LOCKING §4.4)

    /// True when nothing is held and nobody waits — the host's release-reconcile loop parks on
    /// this. Authoritative (reads the live table, never a stale snapshot), so a lock granted
    /// concurrently with the loop's idle check is always seen: the loop can never stop while a
    /// lock is held, which is what stranded locks (a held lock with no loop watching it for its
    /// landing) under the previous snapshot-based stop.
    public func isIdle() -> Bool { holderOf.isEmpty && queues.isEmpty }

    /// The stream the host's reconcile loop parks on while idle. Yields once per grant/enqueue.
    /// `.bufferingNewest(1)` coalesces bursts into a single pending wake-up *and* guarantees a
    /// signal raised between the loop's `isIdle()` check and its park is still delivered — closing
    /// the lost-wake-up window. Single-consumer (the loop); re-subscribing replaces the producer.
    public func reconcileSignals() -> AsyncStream<Void> {
        AsyncStream(bufferingPolicy: .bufferingNewest(1)) { continuation in
            reconcileSignal = continuation
        }
    }

    /// Wake the host's reconcile loop. A no-op before the host subscribes (the loop's drain-on-
    /// start instead covers any lock already held when it first runs).
    private func signalActivity() { reconcileSignal?.yield(()) }

    // MARK: Deadlock detection (LOCKING §8)

    /// Sessions caught in a wait-for cycle — a genuine deadlock that the queue can never clear on
    /// its own. Because a queued session keeps the locks it already holds while waiting for its
    /// delta ("hold edited, queue for delta", LOCKING §4.2), two sessions that each hold a file the
    /// other is queued for will wait on each other forever: A holds X and wants Y, B holds Y and
    /// wants X. Neither releases until granted, neither is granted until the other releases. Only an
    /// outside release (the manual button, a merge, a reap) breaks it — so the UI flags it.
    ///
    /// Builds the wait-for graph (waiter → every *other* session whose held lock overlaps the
    /// waiter's delta, same domain) and returns every session lying on a directed cycle. A waiter
    /// holding no locks can point into a cycle but never be on one, so it's correctly excluded.
    public func deadlockedSessions() -> Set<SessionID> {
        var edges: [SessionID: Set<SessionID>] = [:]
        for (domain, q) in queues {
            for w in q {
                for (key, holder) in holderOf
                where key.domain == domain && holder != w.session
                    && w.delta.contains(where: { ConflictDetector.pathsOverlap(key.path, $0) }) {
                    edges[w.session, default: []].insert(holder)
                }
            }
        }
        return Self.sessionsOnCycle(edges)
    }

    /// Every node that lies on at least one directed cycle of `edges`. A node is on a cycle iff it
    /// can reach itself by following edges. Graphs here are tiny (a handful of contending
    /// sessions), so a per-node reachability walk is plenty.
    private static func sessionsOnCycle(_ edges: [SessionID: Set<SessionID>]) -> Set<SessionID> {
        var onCycle = Set<SessionID>()
        for start in edges.keys {
            var stack = Array(edges[start] ?? [])
            var visited = Set<SessionID>()
            while let node = stack.popLast() {
                if node == start { onCycle.insert(start); break }
                guard visited.insert(node).inserted else { continue }
                stack.append(contentsOf: edges[node] ?? [])
            }
        }
        return onCycle
    }

    // MARK: Helpers

    /// True when every path in `paths` is free in `domain` for `session`: no *other* session
    /// holds an overlapping lock, and nothing reserved earlier in this pass overlaps. Overlap is
    /// directory-aware (`ConflictDetector.pathsOverlap`), so `src/` blocks `src/x.swift`.
    private func isFree(
        _ paths: Set<String>, in domain: LockDomain, for session: SessionID,
        alsoBlockedBy reserved: Set<String> = []
    ) -> Bool {
        for p in paths {
            for (key, holder) in holderOf where key.domain == domain && holder != session {
                if ConflictDetector.pathsOverlap(key.path, p) { return false }
            }
            for r in reserved where ConflictDetector.pathsOverlap(r, p) { return false }
        }
        return true
    }

    private func grant(_ paths: Set<String>, to session: SessionID, in domain: LockDomain) {
        for p in paths {
            let key = LockKey(domain: domain, path: p)
            holderOf[key] = session
            heldBySession[session, default: []].insert(key)
        }
        signalActivity()  // a lock is now held — wake the reconcile loop to watch for its landing
    }

    private static func normalize(_ files: [String]) -> Set<String> {
        Set(files.map(ConflictDetector.normalize).filter { !$0.isEmpty })
    }
}
