Tusker/Tusker/Models/CompositionAttachmentData.swift

//
//  CompositionAttachmentData.swift
//  Tusker
//
//  Created by Shadowfacts on 1/1/20.
//  Copyright © 2020 Shadowfacts. All rights reserved.
//

import UIKit
import Photos
import UniformTypeIdentifiers
import PencilKit

enum CompositionAttachmentData {
    case asset(PHAsset)
    case image(UIImage)
    case video(URL)
    case drawing(PKDrawing)
    
    var type: AttachmentType {
        switch self {
        case let .asset(asset):
            return asset.attachmentType!
        case .image(_):
            return .image
        case .video(_):
            return .video
        case .drawing(_):
            return .image
        }
    }
    
    var isAsset: Bool {
        switch self {
        case .asset(_):
            return true
        default:
            return false
        }
    }
    
    var canSaveToDraft: Bool {
        switch self {
        case .video(_):
            return false
        default:
            return true
        }
    }
    
    func getData(completion: @escaping (Result<(Data, String), Error>) -> Void) {
        switch self {
        case let .image(image):
            // Export as JPEG instead of PNG, otherweise photos straight from the camera are too large
            // for Mastodon in its default configuration (max of 10MB).
            // The quality of 0.8 was chosen completely arbitrarily, it may need to be tuned in the future.
            completion(.success((image.jpegData(compressionQuality: 0.8)!, "image/jpeg")))
        case let .asset(asset):
            if asset.mediaType == .image {
                let options = PHImageRequestOptions()
                options.version = .current
                options.deliveryMode = .highQualityFormat
                options.resizeMode = .none
                options.isNetworkAccessAllowed = true
                PHImageManager.default().requestImageDataAndOrientation(for: asset, options: options) { (data, dataUTI, orientation, info) in
                    guard var data = data, let dataUTI = dataUTI else {
                        completion(.failure(.missingData))
                        return
                    }
    
                    let mimeType: String
                    if dataUTI == "public.heic" {
                        // neither Mastodon nor Pleroma handles HEIC well, so convert to JPEG
                        let image = CIImage(data: data)!
                        let context = CIContext()
                        let colorSpace = image.colorSpace ?? CGColorSpace(name: CGColorSpace.sRGB)!
                        data = context.jpegRepresentation(of: image, colorSpace: colorSpace, options: [:])!
                        mimeType = "image/jpeg"
                    } else {
                        mimeType = UTType(dataUTI)!.preferredMIMEType!
                    }
    
                    completion(.success((data, mimeType)))
                }
            } else if asset.mediaType == .video {
                let options = PHVideoRequestOptions()
                options.deliveryMode = .automatic
                options.isNetworkAccessAllowed = true
                options.version = .current
                PHImageManager.default().requestExportSession(forVideo: asset, options: options, exportPreset: AVAssetExportPresetHighestQuality) { (exportSession, info) in
                    guard let exportSession = exportSession else { fatalError("failed to create export session") }
                    CompositionAttachmentData.exportVideoData(session: exportSession, completion: completion)
                }
            } else {
                fatalError("assetType must be either image or video")
            }
        case let .video(url):
            let asset = AVURLAsset(url: url)
            guard let session = AVAssetExportSession(asset: asset, presetName: AVAssetExportPresetHighestQuality) else {
                fatalError("failed to create export session")
            }
            CompositionAttachmentData.exportVideoData(session: session, completion: completion)
            
        case let .drawing(drawing):
            let image = drawing.imageInLightMode(from: drawing.bounds, scale: 1)
            completion(.success((image.pngData()!, "image/png")))
        }
    }
    
    private static func exportVideoData(session: AVAssetExportSession, completion: @escaping (Result<(Data, String), Error>) -> Void) {
        session.outputFileType = .mp4
        session.outputURL = FileManager.default.temporaryDirectory.appendingPathComponent("exported_video_\(UUID())").appendingPathExtension("mp4")
        session.exportAsynchronously {
            guard session.status == .completed else {
                completion(.failure(.export(session.error!)))
                return
            }
            do {
                let data = try Data(contentsOf: session.outputURL!)
                completion(.success((data, "video/mp4")))
            } catch {
                completion(.failure(.export(error)))
            }
        }
    }
    
    enum AttachmentType {
        case image, video
    }
    
    enum Error: Swift.Error {
        case missingData
        case export(Swift.Error)
    }
}

extension PHAsset {
    var attachmentType: CompositionAttachmentData.AttachmentType? {
        switch self.mediaType {
        case .image:
            return .image
        case .video:
            return .video
        default:
            return nil
        }
    }
}

extension CompositionAttachmentData: Codable {
    func encode(to encoder: Encoder) throws {
        var container = encoder.container(keyedBy: CodingKeys.self)
        
        switch self {
        case let .asset(asset):
            try container.encode("asset", forKey: .type)
            try container.encode(asset.localIdentifier, forKey: .assetIdentifier)
        case let .image(image):
            try container.encode("image", forKey: .type)
            try container.encode(image.pngData()!, forKey: .imageData)
        case .video(_):
            throw EncodingError.invalidValue(self, EncodingError.Context(codingPath: [], debugDescription: "video CompositionAttachments cannot be encoded"))
        case let .drawing(drawing):
            try container.encode("drawing", forKey: .type)
            let drawingData = drawing.dataRepresentation()
            try container.encode(drawingData, forKey: .drawing)
        }
    }
    
    init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        
        switch try container.decode(String.self, forKey: .type) {
        case "asset":
            let identifier = try container.decode(String.self, forKey: .assetIdentifier)
            guard let asset = PHAsset.fetchAssets(withLocalIdentifiers: [identifier], options: nil).firstObject else {
                throw DecodingError.dataCorruptedError(forKey: .assetIdentifier, in: container, debugDescription: "Could not fetch asset with local identifier")
            }
            self = .asset(asset)
        case "image":
            guard let image = UIImage(data: try container.decode(Data.self, forKey: .imageData)) else {
                throw DecodingError.dataCorruptedError(forKey: .imageData, in: container, debugDescription: "Could not decode UIImage from image data")
            }
            self = .image(image)
        case "drawing":
            let drawingData = try container.decode(Data.self, forKey: .drawing)
            let drawing = try PKDrawing(data: drawingData)
            self = .drawing(drawing)
        default:
            throw DecodingError.dataCorruptedError(forKey: .type, in: container, debugDescription: "CompositionAttachment type must be one of 'image' or 'asset'")
        }
    }
    
    enum CodingKeys: CodingKey {
        case type
        case imageData
        /// The local identifier of the PHAsset for this attachment
        case assetIdentifier
        /// The PKDrawing object for this attachment.
        case drawing
    }
}

extension CompositionAttachmentData: Equatable {
    static func ==(lhs: CompositionAttachmentData, rhs: CompositionAttachmentData) -> Bool {
        switch (lhs, rhs) {
        case let (.asset(a), .asset(b)):
            return a.localIdentifier == b.localIdentifier
        case let (.image(a), .image(b)):
            return a == b
        case let (.video(a), .video(b)):
            return a == b
        case let (.drawing(a), .drawing(b)):
            return a == b
        default:
            return false
        }
    }
}