byre/mac/Sources/Byre/Images.swift

import AppKit
import OpenAPIRuntime
import SwiftUI

enum ImageSource: Hashable {
    case blob(String)
    case worktree(taskId: String, path: String)
    case data(Data)
}

private let dataImage = Re(#"^data:image/(png|jpeg|gif|webp);"#, caseless: true)
private let fileURL = Re(#"^file://"#, caseless: true)
private let anyScheme = Re(#"^[a-z][a-z0-9+.-]*:"#, caseless: true)

func worktreeImageSource(_ path: String, _ taskId: String?) -> ImageSource? {
    if dataImage.test(path) {
        guard let comma = path.firstIndex(of: ","), path[..<comma].lowercased().hasSuffix(";base64"), let data = Data(base64Encoded: String(path[path.index(after: comma)...])) else {
            return nil
        }
        return .data(data)
    }
    var file = path
    if fileURL.test(file) {
        guard let decoded = String(file.dropFirst("file://".count)).removingPercentEncoding else {
            return nil
        }
        file = decoded
    } else if anyScheme.test(file) {
        return nil
    }
    guard let taskId else {
        return nil
    }
    return .worktree(taskId: taskId, path: file)
}

@MainActor
final class ImageCache {
    static let shared = ImageCache()
    private var done: [ImageSource: NSImage] = [:]
    private var loading: [ImageSource: Task<NSImage, any Error>] = [:]

    func cached(_ source: ImageSource) -> NSImage? {
        done[source]
    }

    func load(_ source: ImageSource, _ api: API) async throws -> NSImage {
        if let img = done[source] {
            return img
        }
        if let task = loading[source] {
            return try await task.value
        }
        let task = Task { @MainActor in
            let data = try await Self.fetch(source, api)
            guard let img = NSImage(data: data) else {
                throw ApiError(status: 0, message: "not an image")
            }
            return img
        }
        loading[source] = task
        defer {
            loading[source] = nil
        }
        let img = try await task.value
        done[source] = img
        return img
    }

    private static func fetch(_ source: ImageSource, _ api: API) async throws -> Data {
        switch source {
        case .data(let data):
            return data
        case .blob(let id):
            let body = switch try await api.client.getBlob(path: .init(blobId: id)).ok.body {
            case .png(let b), .jpeg(let b), .imageGif(let b), .imageWebp(let b): b
            }
            return try await Data(collecting: body, upTo: 32 << 20)
        case .worktree(let taskId, let path):
            let body = switch try await api.client.getTaskImage(path: .init(taskId: taskId), query: .init(path: path)).ok.body {
            case .png(let b), .jpeg(let b), .imageGif(let b), .imageWebp(let b), .imageSvgXml(let b): b
            }
            return try await Data(collecting: body, upTo: 32 << 20)
        }
    }
}

struct ByreImage: View {
    let source: ImageSource
    let src: String
    var maxWidth: CGFloat = .infinity
    var maxHeight: CGFloat = 480
    @Environment(Store.self) private var store
    @Environment(Overlays.self) private var overlays
    @ViewState private var image: NSImage?
    @ViewState private var failed = false

    var body: some View {
        Group {
            if let image = image ?? ImageCache.shared.cached(source) {
                Image(nsImage: image)
                    .resizable()
                    .scaledToFit()
                    .frame(maxWidth: min(maxWidth, image.size.width), maxHeight: min(maxHeight, image.size.height), alignment: .leading)
                    .clipShape(RoundedRectangle(cornerRadius: 6))
                    .overlay(RoundedRectangle(cornerRadius: 6).stroke(Palette.border))
                    .onTapGesture {
                        overlays.image = image
                    }
                    .help("Open full size")
            } else if failed {
                Text(verbatim: "image not available: \(src)")
                    .foregroundStyle(Palette.muted)
            } else {
                Color.clear.frame(width: 1, height: 1)
            }
        }
        .task(id: source) {
            if image != nil {
                return
            }
            do {
                image = try await ImageCache.shared.load(source, store.api)
            } catch {
                log.error("load image \(src, privacy: .public): \(errorText(error), privacy: .public)")
                failed = true
            }
        }
    }
}

struct MessageImages: View {
    let ids: [String]

    var body: some View {
        FlowLayout(spacing: 6) {
            ForEach(ids, id: \.self) { id in
                ByreImage(source: .blob(id), src: "/api/blobs/\(id)", maxWidth: 240, maxHeight: 180)
            }
        }
        .padding(.vertical, 4)
    }
}