byre/mac/Sources/Byre/Composer.swift

import AppKit
import ByreAPI
import OpenAPIRuntime
import SwiftUI
import UniformTypeIdentifiers

let maxImageBytes = 5 << 20

let models = [
    "claude-opus-5-5",
    "claude-sonnet-5-5",
    "claude-haiku-5-5",
    "claude-fable-5-1",
    "gpt-6.1-sol",
    "gpt-6-astra",
    "gpt-6-sol",
    "gpt-6-luna",
    "gpt-5.6-sol",
    "gpt-5.6-terra",
    "gpt-5.6-luna",
    "gemini-4-argon-high",
    "gemini-4-argon-medium",
    "gemini-4-argon-low",
    "gemini-3.8-flash-high",
    "gemini-3.8-flash-medium",
    "gemini-3.8-flash-low",
]

struct PendingImage: Identifiable {
    let id = UUID()
    let data: Data
    let image: NSImage
}

func imageType(_ data: Data) -> String? {
    let b = [UInt8](data.prefix(12))
    if b.starts(with: [0x89, 0x50, 0x4E, 0x47]) {
        return "png"
    }
    if b.starts(with: [0xFF, 0xD8, 0xFF]) {
        return "jpeg"
    }
    if b.starts(with: Array("GIF8".utf8)) {
        return "gif"
    }
    if b.count == 12, b.starts(with: Array("RIFF".utf8)), Array(b[8..<12]) == Array("WEBP".utf8) {
        return "webp"
    }
    return nil
}

func pngFromTIFF(_ data: Data) -> Data? {
    NSBitmapImageRep(data: data)?.representation(using: .png, properties: [:])
}

func uploadImages(_ api: API, _ imgs: [PendingImage]) async throws -> [String]? {
    if imgs.isEmpty {
        return nil
    }
    var ids: [String] = []
    for img in imgs {
        let body = HTTPBody(img.data)
        let typed: Operations.UploadBlob.Input.Body = switch imageType(img.data) {
        case "png": .png(body)
        case "jpeg": .jpeg(body)
        case "gif": .imageGif(body)
        default: .imageWebp(body)
        }
        ids.append(try await api.client.uploadBlob(body: typed).created.body.json.id)
    }
    return ids
}

@MainActor
var modelChoices: [String: String] = [:]

func vendorOf(_ model: String) -> String {
    if model.hasPrefix("gpt-") {
        return "codex"
    }
    if model.hasPrefix("gemini-") {
        return "antigravity"
    }
    return "claude"
}

@MainActor
@Observable
final class ImageDrafts {
    private var byKey: [String: [PendingImage]] = [:]

    func images(_ key: String) -> [PendingImage] {
        byKey[key] ?? []
    }

    func remove(_ key: String, _ id: UUID) {
        let rest = images(key).filter { $0.id != id }
        byKey[key] = rest.isEmpty ? nil : rest
    }

    func take(_ key: String) -> [PendingImage] {
        let imgs = images(key)
        byKey[key] = nil
        return imgs
    }

    func restore(_ key: String, _ imgs: [PendingImage]) {
        if !imgs.isEmpty {
            byKey[key] = imgs + images(key)
        }
    }

    @discardableResult
    func add(_ key: String, _ data: Data, name: String?, _ overlays: Overlays) -> Bool {
        let file = imageType(data) != nil ? data : pngFromTIFF(data)
        guard let file, let image = NSImage(data: file) else {
            overlays.showAlert("\(name ?? "That file") isn't a PNG, JPEG, GIF or WebP image.")
            return false
        }
        if file.count > maxImageBytes {
            overlays.showAlert("\(name ?? "That image") is over \(maxImageBytes >> 20) MiB.")
            return false
        }
        byKey[key] = images(key) + [PendingImage(data: file, image: image)]
        return true
    }

    func add(_ key: String, pasteboard pb: NSPasteboard, _ overlays: Overlays) -> Bool {
        if let urls = pb.readObjects(forClasses: [NSURL.self], options: [.urlReadingFileURLsOnly: true]) as? [URL], !urls.isEmpty {
            let images = urls.filter { (try? $0.resourceValues(forKeys: [.contentTypeKey]).contentType?.conforms(to: .image)) == true }
            for url in images {
                if let data = try? Data(contentsOf: url) {
                    add(key, data, name: url.lastPathComponent, overlays)
                }
            }
            return !images.isEmpty
        }
        for t in [NSPasteboard.PasteboardType.png, NSPasteboard.PasteboardType(UTType.jpeg.identifier), NSPasteboard.PasteboardType(UTType.gif.identifier), NSPasteboard.PasteboardType(UTType.webP.identifier), .tiff] {
            if let data = pb.data(forType: t) {
                add(key, data, name: nil, overlays)
                return true
            }
        }
        return false
    }

    func add(_ key: String, providers: [NSItemProvider], _ overlays: Overlays) {
        for p in providers {
            if p.hasItemConformingToTypeIdentifier(UTType.fileURL.identifier) {
                _ = p.loadDataRepresentation(for: .fileURL) { data, _ in
                    Task { @MainActor in
                        guard let data, let url = URL(dataRepresentation: data, relativeTo: nil), let file = try? Data(contentsOf: url) else {
                            return
                        }
                        self.add(key, file, name: url.lastPathComponent, overlays)
                    }
                }
                continue
            }
            guard let t = [UTType.png, .jpeg, .gif, .webP, .tiff].first(where: { p.hasItemConformingToTypeIdentifier($0.identifier) }) else {
                continue
            }
            _ = p.loadDataRepresentation(for: t) { data, _ in
                Task { @MainActor in
                    if let data {
                        self.add(key, data, name: nil, overlays)
                    }
                }
            }
        }
    }
}

struct ComposerTarget {
    let paste: @MainActor (NSPasteboard) -> Bool
}

struct ComposerFocusKey: FocusedValueKey {
    typealias Value = ComposerTarget
}

extension FocusedValues {
    var composer: ComposerTarget? {
        get {
            self[ComposerFocusKey.self]
        }
        set {
            self[ComposerFocusKey.self] = newValue
        }
    }
}

struct PasteboardCommands: Commands {
    @FocusedValue(\.composer) private var composer

    var body: some Commands {
        CommandGroup(replacing: .pasteboard) {
            Button("Cut") {
                NSApp.sendAction(#selector(NSText.cut(_:)), to: nil, from: nil)
            }
            .keyboardShortcut("x")
            Button("Copy") {
                NSApp.sendAction(#selector(NSText.copy(_:)), to: nil, from: nil)
            }
            .keyboardShortcut("c")
            Button("Paste") {
                if let composer, composer.paste(.general) {
                    return
                }
                NSApp.sendAction(#selector(NSText.paste(_:)), to: nil, from: nil)
            }
            .keyboardShortcut("v")
            Button("Select All") {
                NSApp.sendAction(#selector(NSText.selectAll(_:)), to: nil, from: nil)
            }
            .keyboardShortcut("a")
        }
    }
}

struct Thumbnails: View {
    let key: String
    @Environment(Store.self) private var store

    var body: some View {
        let drafts = store.images
        let imgs = drafts.images(key)
        if !imgs.isEmpty {
            FlowLayout(spacing: 8) {
                ForEach(imgs) { img in
                    Image(nsImage: img.image)
                        .resizable()
                        .scaledToFit()
                        .frame(maxWidth: 120, maxHeight: 90)
                        .clipShape(RoundedRectangle(cornerRadius: 6))
                        .overlay(RoundedRectangle(cornerRadius: 6).stroke(Palette.border))
                        .overlay(alignment: .topTrailing) {
                            Button("×") {
                                drafts.remove(key, img.id)
                            }
                            .buttonStyle(ThumbRemoveStyle())
                            .help("Remove image")
                            .offset(x: 6, y: -6)
                        }
                }
            }
            .padding(.top, 6)
        }
    }
}

struct ThumbRemoveStyle: ButtonStyle {
    func makeBody(configuration: ButtonStyleConfiguration) -> some View {
        ThumbRemove(configuration: configuration)
    }
}

private struct ThumbRemove: View {
    let configuration: ButtonStyleConfiguration
    @ViewState private var hover = false

    var body: some View {
        configuration.label
            .font(Fonts.font(13))
            .foregroundStyle(Palette.text)
            .padding(.horizontal, 6)
            .frame(height: 18)
            .padding(1)
            .background(Palette.raised, in: RoundedRectangle(cornerRadius: 9))
            .overlay(RoundedRectangle(cornerRadius: 9).strokeBorder(hover ? Palette.muted : Palette.border))
            .contentShape(RoundedRectangle(cornerRadius: 9))
            .onHover { hover = $0 }
    }
}

struct ComposerView: View {
    let chat: ChatModel
    let key: String
    @Environment(Store.self) private var store
    @Environment(Overlays.self) private var overlays
    @Environment(\.paneAccent) private var accent
    @Environment(\.inRightPane) private var inRightPane
    @ViewState private var text = ""
    @ViewState private var pick: String?
    @ViewState private var targeted = false
    @FocusState private var focused: Bool

    private var draftKey: String {
        "message:\(key)"
    }

    private struct ModelState {
        let live: AgentSession?
        let current: String
        let options: [String]
        let modelKey: String
        let selected: String
    }

    private func modelState() -> ModelState {
        let live = chat.liveSession
        let current = live.map { $0.model.isEmpty ? chat.builder.context[$0.id]?.model ?? "" : $0.model } ?? ""
        let offered = live.map { s in models.filter { vendorOf($0) == s.vendor.rawValue } } ?? models
        let options = !current.isEmpty && !offered.contains(current) ? [current] + offered : offered
        let modelKey = "message:\(live?.id ?? key)"
        let fresh = live == nil ? store.defaultModel : options[0]
        return ModelState(live: live, current: current, options: options, modelKey: modelKey, selected: pick ?? modelChoices[modelKey] ?? (current.isEmpty ? fresh : current))
    }

    var body: some View {
        let m = modelState()
        let live = m.live
        let sessionState = live?.state ?? .idle
        let options = m.options
        let modelKey = m.modelKey
        let selected = m.selected
        let lost = store.drafts?.lostImageCount(draftKey) ?? 0

        VStack(alignment: .leading, spacing: 8) {
            if lost > 0 {
                HStack(spacing: 8) {
                    Text("\(lost == 1 ? "1 image" : "\(lost) images") from before the reload need re-attaching.")
                        .font(Fonts.font(12))
                    Spacer()
                    Button("×") {
                        store.drafts?.forgetLostImages(draftKey)
                    }
                    .buttonStyle(WebButtonStyle())
                    .help("Dismiss")
                }
                .padding(.horizontal, 10)
                .padding(.vertical, 6)
                .background(Palette.warnTint, in: RoundedRectangle(cornerRadius: 6))
                .overlay(RoundedRectangle(cornerRadius: 6).stroke(Palette.border))
            }
            Thumbnails(key: draftKey)
            editor
            HStack(spacing: 8) {
                if !chat.team {
                    Picker("", selection: Binding(get: { selected }, set: { v in
                        pick = v
                        modelChoices[modelKey] = v
                    })) {
                        ForEach(options, id: \.self) { m in
                            Text(m).tag(m)
                        }
                    }
                    .labelsHidden()
                    .pickerStyle(.menu)
                    .fixedSize()
                    .help("Model for the next turn")
                }
                Spacer()
                Button("Interrupt") {
                    if let id = live?.id {
                        chat.interrupt(id, overlays)
                    }
                }
                .buttonStyle(WebButtonStyle(danger: true))
                .disabled(sessionState != .running)
                Button("Send") {
                    submit(modelState())
                }
                .buttonStyle(WebButtonStyle())
            }
        }
        .padding(.horizontal, 16)
        .padding(.vertical, 12)
        .background(Palette.surface)
        .overlay(alignment: .top) {
            Rectangle().fill(Palette.border).frame(height: 1)
        }
        .onAppear {
            text = store.drafts?.composerText(draftKey) ?? ""
            if chat.takeFocus() {
                focused = true
            }
        }
        .onChange(of: store.composerFocus) {
            if chat.takeFocus() {
                focused = true
            }
        }
        .onChange(of: text) {
            store.drafts?.setComposerText(draftKey, text)
            store.drafts?.setComposerImages(draftKey, store.images.images(draftKey).count)
            reportDraft()
        }
        .onChange(of: store.drafts?.composerText(draftKey)) { _, new in
            if let new, new != text {
                text = new
            }
        }
        .onChange(of: store.images.images(draftKey).count) {
            store.drafts?.setComposerImages(draftKey, store.images.images(draftKey).count)
        }
        .onChange(of: modelKey) {
            pick = nil
        }
        .onChange(of: overlays.leadFocus) {
            if chat.kind == .lead {
                focused = true
            }
        }
        .onChange(of: focused) {
            if !focused {
                store.drafts?.flush()
            }
        }
        .onDisappear {
            if inRightPane {
                store.paneFocus.setDraft(draftKey, false)
            }
            store.drafts?.flush()
        }
    }

    private var editor: some View {
        TextEditor(text: $text)
            .font(Fonts.font(13))
            .scrollContentBackground(.hidden)
            .writingToolsBehavior(.disabled)
            .focused($focused)
            .focusedValue(\.composer, ComposerTarget(paste: { pb in store.images.add(draftKey, pasteboard: pb, overlays) }))
            .onKeyPress(.return, phases: .down) { press in
                let marked = (NSApp.keyWindow?.firstResponder as? NSTextView)?.hasMarkedText() ?? false
                if press.modifiers.contains(.shift) || marked {
                    return .ignored
                }
                submit(modelState())
                return .handled
            }
            .padding(.horizontal, 5)
            .padding(.vertical, 6)
            .frame(height: 104)
            .background(Palette.raised, in: RoundedRectangle(cornerRadius: 6))
            .overlay(alignment: .topLeading) {
                if text.isEmpty {
                    Text(chat.placeholder)
                        .foregroundStyle(Palette.muted)
                        .padding(.horizontal, 10)
                        .padding(.vertical, 6)
                        .allowsHitTesting(false)
                }
            }
            .overlay(RoundedRectangle(cornerRadius: 6).stroke(targeted ? Palette.accent : inRightPane ? accent : focused ? Palette.accent : Palette.border, lineWidth: inRightPane && focused ? 2 : 1))
            .overlay {
                Color.clear
                    .contentShape(Rectangle())
                    .allowsHitTesting(false)
                    .onDrop(of: [.image, .fileURL], isTargeted: $targeted) { providers in
                        store.images.add(draftKey, providers: providers, overlays)
                        return true
                    }
            }
    }

    private func reportDraft() {
        if inRightPane {
            store.paneFocus.setDraft(draftKey, !trimmed(text).isEmpty)
        }
    }

    private func submit(_ m: ModelState) {
        let value = text
        let images = store.images
        if trimmed(value).isEmpty && images.images(draftKey).isEmpty {
            return
        }
        guard let drafts = store.drafts else {
            return
        }
        let imgs = images.take(draftKey)
        text = ""
        let sent = drafts.startSend(draftKey, value, imgs.count)
        let cancel = chat.team ? {} : chat.builder.showPending(value, imgs)
        Task {
            do {
                let ids = try await uploadImages(store.api, imgs)
                try await chat.send(sessionId: m.live?.id, text: value, images: ids, model: m.selected != m.current ? m.selected : nil, firstModel: m.selected)
                modelChoices[m.modelKey] = nil
                pick = nil
                drafts.endSend(sent, true)
            } catch {
                cancel()
                drafts.endSend(sent, false)
                text = drafts.composerText(draftKey)
                images.restore(draftKey, imgs)
                overlays.showAlert(errorText(error))
            }
        }
    }
}