shannoncoat/Sources/WindowOverlay.swift

// A small floating tag pinned to a Claude window's corner, showing which
// profile it belongs to — every profile launches the same app with the
// same icon and title, so there's otherwise no way to tell two open
// Claude windows apart at a glance. Kept in sync live via AXObserver
// notifications (move/resize/miniaturize/destroy) rather than polling —
// only a persistent process can hold a watcher like this at all, which is
// exactly the gap the old one-shot CLI script couldn't close.
import AppKit
import ApplicationServices

// The tag's position relative to its window's top-right corner, shared
// across every open tag and persisted so a drag survives a relaunch.
// Anchored from the right/top edges (not left) because that's the corner
// that stays put under the common resize gesture (dragging the bottom-right
// handle) — an offset-from-left would drift away from a user's chosen spot
// every time the window changed width.
enum WindowOverlayPosition {
    private static let rightInsetKey = "WindowOverlayRightInset"
    private static let topInsetKey = "WindowOverlayTopInset"

    // Top-right by default — top-left is where every window's close/
    // minimize/zoom controls live, which the tag must never sit over.
    static let defaultRightInset: CGFloat = 10
    static let defaultTopInset: CGFloat = 6

    static var rightInset: CGFloat {
        get { (UserDefaults.standard.object(forKey: rightInsetKey) as? Double).map { CGFloat($0) } ?? defaultRightInset }
        set { UserDefaults.standard.set(Double(newValue), forKey: rightInsetKey) }
    }

    static var topInset: CGFloat {
        get { (UserDefaults.standard.object(forKey: topInsetKey) as? Double).map { CGFloat($0) } ?? defaultTopInset }
        set { UserDefaults.standard.set(Double(newValue), forKey: topInsetKey) }
    }
}

final class WindowOverlay: NSObject, NSWindowDelegate {
    let pid: pid_t
    private let axWindow: AXUIElement
    private var observer: AXObserver?
    private let panel: NSPanel

    private static let tagSize = NSSize(width: 110, height: 20)

    init?(profileName: String, color: NSColor, pid: pid_t) {
        let axApp = AXUIElementCreateApplication(pid)
        guard let window = Self.firstWindow(of: axApp) else { return nil }
        self.pid = pid
        self.axWindow = window

        let panel = NSPanel(
            contentRect: NSRect(origin: .zero, size: Self.tagSize),
            styleMask: [.borderless, .nonactivatingPanel],
            backing: .buffered, defer: false)
        panel.isOpaque = false
        panel.backgroundColor = .clear
        panel.hasShadow = true
        panel.level = .floating
        // Draggable (so it can be moved off whatever it's obstructing),
        // but .nonactivatingPanel keeps a click/drag from stealing focus
        // away from the Claude window underneath.
        panel.isMovableByWindowBackground = true
        panel.collectionBehavior = [.stationary, .ignoresCycle]

        let background = NSView(frame: NSRect(origin: .zero, size: Self.tagSize))
        background.wantsLayer = true
        background.layer?.backgroundColor = color.cgColor
        background.layer?.cornerRadius = 5

        let label = NSTextField(labelWithString: profileName)
        label.frame = NSRect(origin: .zero, size: Self.tagSize)
        label.alignment = .center
        label.font = .systemFont(ofSize: 11, weight: .semibold)
        label.textColor = .white
        label.drawsBackground = false
        label.lineBreakMode = .byTruncatingTail
        background.addSubview(label)

        panel.contentView = background
        self.panel = panel
        super.init()
        panel.delegate = self

        guard reposition() else { return nil }
        startObserving()
    }

    // A cold Electron launch can take a second or more between the
    // process starting (which is what triggers `updateWindowOverlays`,
    // via NSWorkspace's launch notification) and its first window
    // actually existing — the same race ClaudeControl.focus's
    // waitForWindow already handles for the same reason. `init?` can't
    // retry on its own (once it returns nil, that attempt is done), so
    // this polls it every quarter second for up to ~3s before giving up.
    static func attach(
        profileName: String, color: NSColor, pid: pid_t, attemptsRemaining: Int = 12,
        completion: @escaping (WindowOverlay?) -> Void
    ) {
        if let overlay = WindowOverlay(profileName: profileName, color: color, pid: pid) {
            completion(overlay)
        } else if attemptsRemaining > 0 {
            DispatchQueue.main.asyncAfter(deadline: .now() + 0.25) {
                attach(
                    profileName: profileName, color: color, pid: pid, attemptsRemaining: attemptsRemaining - 1,
                    completion: completion)
            }
        } else {
            completion(nil)
        }
    }

    deinit {
        stopObserving()
        panel.orderOut(nil)
    }

    private static func firstWindow(of axApp: AXUIElement) -> AXUIElement? {
        var value: CFTypeRef?
        guard AXUIElementCopyAttributeValue(axApp, kAXWindowsAttribute as CFString, &value) == .success,
              let windows = value as? [AXUIElement], let first = windows.first
        else { return nil }
        return first
    }

    private static func frame(of window: AXUIElement) -> CGRect? {
        var posValue: CFTypeRef?
        var sizeValue: CFTypeRef?
        guard AXUIElementCopyAttributeValue(window, kAXPositionAttribute as CFString, &posValue) == .success,
              AXUIElementCopyAttributeValue(window, kAXSizeAttribute as CFString, &sizeValue) == .success,
              CFGetTypeID(posValue) == AXValueGetTypeID(), CFGetTypeID(sizeValue) == AXValueGetTypeID()
        else { return nil }
        var point = CGPoint.zero
        var size = CGSize.zero
        guard AXValueGetValue(posValue as! AXValue, .cgPoint, &point),
              AXValueGetValue(sizeValue as! AXValue, .cgSize, &size)
        else { return nil }
        return CGRect(origin: point, size: size)
    }

    private static func screen(containing axFrame: CGRect) -> NSScreen? {
        NSScreen.screens.first { $0.frame.minX <= axFrame.midX && axFrame.midX <= $0.frame.maxX } ?? NSScreen.main
    }

    // AX coordinates are top-left-origin (screen top = y 0); AppKit screen
    // coordinates are bottom-left-origin. Returns the window's frame
    // translated into AppKit's space.
    private static func appKitFrame(of axFrame: CGRect, on screen: NSScreen) -> NSRect {
        NSRect(x: axFrame.origin.x, y: screen.frame.maxY - axFrame.origin.y - axFrame.height,
               width: axFrame.width, height: axFrame.height)
    }

    // Positions the tag using the persisted (or default) offset from the
    // window's top-right corner. Returns false (and hides the tag) if the
    // window has no readable frame right now — minimized, or Accessibility
    // not granted.
    @discardableResult
    private func reposition() -> Bool {
        guard let axFrame = Self.frame(of: axWindow), let screen = Self.screen(containing: axFrame) else {
            panel.orderOut(nil)
            return false
        }
        let windowFrame = Self.appKitFrame(of: axFrame, on: screen)
        let origin = NSPoint(
            x: windowFrame.maxX - WindowOverlayPosition.rightInset - Self.tagSize.width,
            y: windowFrame.maxY - WindowOverlayPosition.topInset - Self.tagSize.height)
        panel.setFrameOrigin(origin)
        if !panel.isVisible { panel.orderFrontRegardless() }
        return true
    }

    // Fires for every panel move — ours (from `reposition`, when the
    // tracked window itself moves/resizes) and the user's (dragging the
    // tag). Either way, re-deriving the offset from the panel's actual
    // current position and persisting it is correct: our own moves just
    // recompute the same offset they were placed with, a no-op; a user
    // drag captures the new spot they chose so it survives the next
    // reposition and the next launch. Clamped to stay within the window
    // so a sloppy drag can't wander off it entirely.
    func windowDidMove(_ notification: Notification) {
        guard let axFrame = Self.frame(of: axWindow), let screen = Self.screen(containing: axFrame) else { return }
        let windowFrame = Self.appKitFrame(of: axFrame, on: screen)
        let panelOrigin = panel.frame.origin
        let rightInset = (windowFrame.maxX - (panelOrigin.x + Self.tagSize.width))
            .clamped(to: 0...(max(0, windowFrame.width - Self.tagSize.width)))
        let topInset = (windowFrame.maxY - (panelOrigin.y + Self.tagSize.height))
            .clamped(to: 0...(max(0, windowFrame.height - Self.tagSize.height)))
        WindowOverlayPosition.rightInset = rightInset
        WindowOverlayPosition.topInset = topInset
    }

    // MARK: - Live tracking

    private func startObserving() {
        var newObserver: AXObserver?
        guard AXObserverCreate(pid, WindowOverlay.axCallback, &newObserver) == .success, let newObserver
        else { return }
        observer = newObserver

        let refcon = Unmanaged.passUnretained(self).toOpaque()
        for name in [
            kAXMovedNotification, kAXResizedNotification, kAXUIElementDestroyedNotification,
            kAXWindowMiniaturizedNotification, kAXWindowDeminiaturizedNotification,
        ] {
            AXObserverAddNotification(newObserver, axWindow, name as CFString, refcon)
        }
        CFRunLoopAddSource(CFRunLoopGetCurrent(), AXObserverGetRunLoopSource(newObserver), .defaultMode)
    }

    private func stopObserving() {
        guard let observer else { return }
        CFRunLoopRemoveSource(CFRunLoopGetCurrent(), AXObserverGetRunLoopSource(observer), .defaultMode)
        self.observer = nil
    }

    private static let axCallback: AXObserverCallback = { _, _, notification, refcon in
        guard let refcon else { return }
        let overlay = Unmanaged<WindowOverlay>.fromOpaque(refcon).takeUnretainedValue()
        switch notification as String {
        case kAXUIElementDestroyedNotification, kAXWindowMiniaturizedNotification:
            overlay.panel.orderOut(nil)
        default:
            overlay.reposition()
        }
    }
}

extension Comparable {
    func clamped(to range: ClosedRange<Self>) -> Self {
        min(max(self, range.lowerBound), range.upperBound)
    }
}