another.im-ios/AnotherIM/AppData/Store/AttachmentsStore.swift

import Combine
import Foundation
import GRDB
import Photos
import SwiftUI

@MainActor
final class AttachmentsStore: ObservableObject {
    @Published private(set) var cameraAccessGranted = false
    @Published private(set) var galleryAccessGranted = false
    @Published private(set) var galleryItems: [GalleryItem] = []

    private let client: Client
    private let roster: Roster
    private var secured: Bool = false

    private var messagesCancellable: AnyCancellable?
    private var chatCancellable: AnyCancellable?
    private var processing: Set<String> = []

    init(roster: Roster, client: Client) {
        self.client = client
        self.roster = roster
        subscribe()
    }
}

// MARK: - Camera and Gallery access
extension AttachmentsStore {
    func checkCameraAuthorization() async {
        let status = AVCaptureDevice.authorizationStatus(for: .video)
        var isAuthorized = status == .authorized
        if status == .notDetermined {
            isAuthorized = await AVCaptureDevice.requestAccess(for: .video)
        }
        cameraAccessGranted = isAuthorized
    }

    func checkGalleryAuthorization() async {
        let status = PHPhotoLibrary.authorizationStatus()
        var isAuthorized = status == .authorized
        if status == .notDetermined {
            let req = await PHPhotoLibrary.requestAuthorization(for: .readWrite)
            isAuthorized = (req == .authorized) || (req == .limited)
        }
        galleryAccessGranted = isAuthorized
        if isAuthorized {
            await fetchGalleryItems()
        }
    }

    private func fetchGalleryItems() async {
        guard galleryAccessGranted else { return }
        galleryItems = await GalleryItem.fetchAll()
    }
}

// MARK: - Save outgoing attachments for future uploadings
extension AttachmentsStore {
    func sendMedia(_ items: [GalleryItem]) {
        Task {
            for item in items {
                Task {
                    var message = Message.blank
                    message.from = roster.bareJid
                    message.to = roster.contactBareJid
                    message.secure = secured

                    switch item.type {
                    case .photo:
                        guard let asset = PHAsset.fetchAssets(withLocalIdentifiers: [item.id], options: nil).firstObject else { return }
                        guard let photo = try? await PHImageManager.default().getPhoto(for: asset) else { return }
                        guard let data = photo.jpegData(compressionQuality: 1.0) else { return }
                        let localName = "\(message.id)_\(UUID().uuidString).jpg"
                        let localUrl = FolderWrapper.shared.fileFolder.appendingPathComponent(localName)
                        try? data.write(to: localUrl)
                        message.contentType = .attachment(
                            Attachment(
                                type: .image,
                                localName: localName,
                                thumbnailName: nil,
                                remotePath: nil
                            )
                        )
                        try? await message.save()

                    case .video:
                        guard let asset = PHAsset.fetchAssets(withLocalIdentifiers: [item.id], options: nil).firstObject else { return }
                        guard let video = try? await PHImageManager.default().getVideo(for: asset) else { return }
                        // swiftlint:disable:next force_cast
                        let assetURL = video as! AVURLAsset
                        let url = assetURL.url
                        let localName = "\(message.id)_\(UUID().uuidString).mov"
                        let localUrl = FolderWrapper.shared.fileFolder.appendingPathComponent(localName)
                        try? FileManager.default.copyItem(at: url, to: localUrl)
                        message.contentType = .attachment(
                            Attachment(
                                type: .video,
                                localName: localName,
                                thumbnailName: nil,
                                remotePath: nil
                            )
                        )
                        try? await message.save()
                    }
                }
            }
        }
    }

    func sendCaptured(_ data: Data, _ type: GalleryMediaType) {
        Task {
            var message = Message.blank
            message.from = roster.bareJid
            message.to = roster.contactBareJid
            message.secure = secured

            let localName: String
            let msgType: AttachmentType
            do {
                (localName, msgType) = try await Task {
                    // local name
                    let fileId = UUID().uuidString
                    let localName: String
                    let msgType: AttachmentType
                    switch type {
                    case .photo:
                        localName = "\(message.id)_\(fileId).jpg"
                        msgType = .image

                    case .video:
                        localName = "\(message.id)_\(fileId).mov"
                        msgType = .video
                    }

                    // save
                    let localUrl = FolderWrapper.shared.fileFolder.appendingPathComponent(localName)
                    try data.write(to: localUrl)
                    return (localName, msgType)
                }.value
            } catch {
                logIt(.error, "Can't save file for uploading: \(error)")
                return
            }

            // save message
            message.contentType = .attachment(
                Attachment(
                    type: msgType,
                    localName: localName,
                    thumbnailName: nil,
                    remotePath: nil
                )
            )
            do {
                try await message.save()
            } catch {
                logIt(.error, "Can't save message: \(error)")
                return
            }
        }
    }

    func sendDocuments(_ data: [Data], _ extensions: [String]) {
        Task {
            for (index, data) in data.enumerated() {
                Task {
                    let newMessageId = UUID().uuidString
                    let fileId = UUID().uuidString
                    let localName = "\(newMessageId)_\(fileId).\(extensions[index])"
                    let localUrl = FolderWrapper.shared.fileFolder.appendingPathComponent(localName)
                    do {
                        try data.write(to: localUrl)
                    } catch {
                        print("FileProcessing: Error writing document: \(error)")
                        return
                    }

                    var message = Message.blank
                    message.from = roster.bareJid
                    message.to = roster.contactBareJid
                    message.secure = secured
                    message.contentType = .attachment(
                        Attachment(
                            type: localName.attachmentType,
                            localName: localName,
                            thumbnailName: nil,
                            remotePath: nil
                        )
                    )
                    do {
                        try await message.save()
                    } catch {
                        print("FileProcessing: Error saving document: \(error)")
                    }
                }
            }
        }
    }
}

// MARK: - Processing attachments
private extension AttachmentsStore {
    func subscribe() {
        messagesCancellable = ValueObservation.tracking(Message
            .filter(
                (Column("to") == roster.bareJid && Column("from") == roster.contactBareJid) ||
                    (Column("from") == roster.bareJid && Column("to") == roster.contactBareJid)
            )
            .order(Column("date").desc)
            .fetchAll
        )
        .publisher(in: Database.shared.dbQueue, scheduling: .immediate)
        .receive(on: DispatchQueue.main)
        .sink { _ in
        } receiveValue: { [weak self] messages in
            let forProcessing = messages
                .filter { $0.status != .error }
                .filter { self?.processing.contains($0.id) == false }
                .filter { $0.contentType.isAttachment }
            for message in forProcessing {
                if case .attachment(let attachment) = message.contentType {
                    let localPath = attachment.localPath
                    if localPath != nil, attachment.remotePath == nil {
                        // Uploading
                        self?.processing.insert(message.id)
                        Task {
                            await self?.uploadAttachment(message)
                        }
                    } else if localPath == nil, attachment.remotePath != nil {
                        // Downloading
                        self?.processing.insert(message.id)
                        Task {
                            await self?.downloadAttachment(message)
                        }
                    } else if localPath != nil, attachment.remotePath != nil, attachment.thumbnailName == nil, attachment.type == .image {
                        // Generate thumbnail
                        self?.processing.insert(message.id)
                        Task {
                            await self?.generateThumbnail(message)
                        }
                    }
                }
            }
        }

        chatCancellable = ValueObservation.tracking(Chat
            .filter(Column("account") == roster.bareJid && Column("participant") == roster.contactBareJid)
            .fetchOne
        )
        .publisher(in: Database.shared.dbQueue, scheduling: .immediate)
        .receive(on: DispatchQueue.main)
        .sink { _ in
        } receiveValue: { [weak self] chat in
            guard let self = self else { return }
            self.secured = chat?.encrypted ?? false
        }
    }
}

// MARK: - Uploadings/Downloadings
extension AttachmentsStore {
    private func uploadAttachment(_ message: Message) async {
        do {
            try await message.setStatus(.pending)
            var message = message
            guard case .attachment(let attachment) = message.contentType else {
                throw AppError.invalidContentType
            }
            guard let localName = attachment.localPath else {
                throw AppError.invalidLocalName
            }
            let remotePath = try await client.uploadFile(localName, needEncrypt: message.secure)
            message.contentType = .attachment(
                Attachment(
                    type: attachment.type,
                    localName: attachment.localName,
                    thumbnailName: nil,
                    remotePath: remotePath
                )
            )
            message.body = remotePath
            message.oobUrl = remotePath
            try await message.save()
            try await client.sendMessage(message)
            processing.remove(message.id)
            try await message.setStatus(.sent)
        } catch {
            processing.remove(message.id)
            try? await message.setStatus(.error)
        }
    }

    private func downloadAttachment(_ message: Message) async {
        guard case .attachment(let attachment) = message.contentType else {
            return
        }
        guard let remotePath = attachment.remotePath, var remoteUrl = URL(string: remotePath) else {
            return
        }
        do {
            // if attachment encrypted, extract the key
            // and format remote url
            var encryptionKey: String?
            if remoteUrl.scheme == "aesgcm", var components = URLComponents(url: remoteUrl, resolvingAgainstBaseURL: true) {
                encryptionKey = components.fragment
                components.scheme = "https"
                components.fragment = nil
                if let tmpUrl = components.url {
                    remoteUrl = tmpUrl
                }
            }

            // make local name/path
            let localName = "\(message.id)_\(UUID().uuidString).\(remoteUrl.lastPathComponent)"
            let localUrl = FolderWrapper.shared.fileFolder.appendingPathComponent(localName)

            // Download the file
            let (tempUrl, _) = try await URLSession.shared.download(from: remoteUrl)
            try FileManager.default.moveItem(at: tempUrl, to: localUrl)

            if let encryptionKey {
                // Decrypt the file
                guard encryptionKey.count % 2 == 0, encryptionKey.count > 64 else {
                    throw AppError.securityError
                }
                let fragmentData = encryptionKey.map { char -> UInt8 in
                    return UInt8(char.hexDigitValue ?? 0)
                }
                let ivLen = fragmentData.count - (32 * 2)
                var iv = Data()
                var key = Data()

                for index in 0 ..< (ivLen / 2) {
                    iv.append(fragmentData[index * 2] * 16 + fragmentData[index * 2 + 1])
                }
                for index in (ivLen / 2) ..< (fragmentData.count / 2) {
                    key.append(fragmentData[index * 2] * 16 + fragmentData[index * 2 + 1])
                }

                let encodedData = try Data(contentsOf: localUrl)
                var result = Data()

                guard AESGSMEngine.shared.decrypt(iv: iv, key: key, encoded: encodedData, auth: nil, output: &result) else {
                    throw AppError.securityError
                }

                try result.write(to: localUrl)
            }

            var message = message
            message.contentType = .attachment(
                Attachment(
                    type: attachment.type,
                    localName: localName,
                    thumbnailName: attachment.thumbnailName,
                    remotePath: remotePath
                )
            )
            processing.remove(message.id)
            try await message.save()
        } catch {
            logIt(.error, "Can't download attachment: \(error)")
        }
    }

    private func generateThumbnail(_ message: Message) async {
        guard case .attachment(let attachment) = message.contentType else {
            return
        }
        guard attachment.type == .image else {
            return
        }
        guard let localName = attachment.localName, let localPath = attachment.localPath else {
            return
        }
        let thumbnailFileName = "thumb_\(localName)"
        let thumbnailUrl = FolderWrapper.shared.fileFolder.appendingPathComponent(thumbnailFileName)

        //
        if !FileManager.default.fileExists(atPath: thumbnailUrl.path) {
            guard let image = UIImage(contentsOfFile: localPath.path) else {
                return
            }
            let targetSize = CGSize(width: Const.attachmentPreviewSize, height: Const.attachmentPreviewSize)
            guard let thumbnail = try? await image.scaleAndCropImage(targetSize) else {
                return
            }
            guard let data = thumbnail.jpegData(compressionQuality: 0.5) else {
                return
            }
            do {
                try data.write(to: thumbnailUrl)
            } catch {
                return
            }
        }

        //
        var message = message
        message.contentType = .attachment(
            Attachment(
                type: attachment.type,
                localName: attachment.localName,
                thumbnailName: thumbnailFileName,
                remotePath: attachment.remotePath
            )
        )
        processing.remove(message.id)
        try? await message.save()
    }
}