swift-apple-intelligence-grpc/Sources/AppleIntelligenceApp/Models/ChatMessage.swift

import Foundation
import AppKit

/// Represents an attached image in a chat message
struct ImageAttachment: Identifiable, Equatable {
    let id: UUID
    let data: Data
    let filename: String?
    let thumbnail: NSImage?
    let mimeType: String

    init(data: Data, filename: String? = nil) {
        self.id = UUID()
        self.data = data
        self.filename = filename
        self.thumbnail = Self.generateThumbnail(from: data)
        self.mimeType = Self.detectMimeType(from: data)
    }

    private static func generateThumbnail(from data: Data) -> NSImage? {
        guard let image = NSImage(data: data) else { return nil }

        let maxSize: CGFloat = 100
        let ratio = min(maxSize / image.size.width, maxSize / image.size.height, 1.0)
        let newSize = NSSize(
            width: image.size.width * ratio,
            height: image.size.height * ratio
        )

        let thumbnail = NSImage(size: newSize)
        thumbnail.lockFocus()
        image.draw(
            in: NSRect(origin: .zero, size: newSize),
            from: NSRect(origin: .zero, size: image.size),
            operation: .copy,
            fraction: 1.0
        )
        thumbnail.unlockFocus()
        return thumbnail
    }

    private static func detectMimeType(from data: Data) -> String {
        guard data.count >= 4 else { return "application/octet-stream" }
        let bytes = [UInt8](data.prefix(4))

        if bytes[0] == 0x89 && bytes[1] == 0x50 && bytes[2] == 0x4E && bytes[3] == 0x47 {
            return "image/png"
        } else if bytes[0] == 0xFF && bytes[1] == 0xD8 {
            return "image/jpeg"
        } else if bytes[0] == 0x47 && bytes[1] == 0x49 && bytes[2] == 0x46 {
            return "image/gif"
        }
        return "image/png" // Default to PNG
    }

    static func == (lhs: ImageAttachment, rhs: ImageAttachment) -> Bool {
        lhs.id == rhs.id
    }
}

struct ChatMessage: Identifiable, Equatable {
    let id: UUID
    let role: Role
    var content: String
    let timestamp: Date
    var isStreaming: Bool
    var images: [ImageAttachment]

    enum Role: Equatable {
        case user
        case assistant
    }

    init(role: Role, content: String, isStreaming: Bool = false, images: [ImageAttachment] = []) {
        self.id = UUID()
        self.role = role
        self.content = content
        self.timestamp = Date()
        self.isStreaming = isStreaming
        self.images = images
    }
}