//
//  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)
    case gif(Data)
    
    var type: AttachmentType {
        switch self {
        case let .asset(asset):
            return asset.attachmentType!
        case .image(_):
            return .image
        case .video(_):
            return .video
        case .drawing(_):
            return .image
        case .gif(_):
            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(features: InstanceFeatures, skipAllConversion: Bool = false, completion: @escaping (Result<(Data, UTType), 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)!, .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
                    }
                    
                    guard !skipAllConversion else {
                        completion(.success((data, UTType(dataUTI)!)))
                        return
                    }
                    
                    let utType: UTType
                    let image = CIImage(data: data)!
                    let needsColorSpaceConversion = features.needsWideColorGamutHack && image.colorSpace?.name != CGColorSpace.sRGB
    
                    // neither Mastodon nor Pleroma handles HEIC well, so convert to JPEG
                    // they also do a bad job converting wide color gamut images (they seem to just drop the profile, letting the wide-gamut values be reinterprete as sRGB)
                    // if that changes in the future, we'll need to pass the InstanceFeatures in here somehow and gate the conversion
                    if needsColorSpaceConversion || dataUTI == "public.heic" {
                        let context = CIContext()
                        let colorSpace = needsColorSpaceConversion || image.colorSpace != nil ? CGColorSpace(name: CGColorSpace.sRGB)! : image.colorSpace!
                        if dataUTI == "public.png" {
                            data = context.pngRepresentation(of: image, format: .ARGB8, colorSpace: colorSpace)!
                            utType = .png
                        } else {
                            data = context.jpegRepresentation(of: image, colorSpace: colorSpace)!
                            utType = .jpeg
                        }
                    } else {
                        utType = UTType(dataUTI)!
                    }
    
                    completion(.success((data, utType)))
                }
            } 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
                    if let exportSession = exportSession {
                        CompositionAttachmentData.exportVideoData(session: exportSession, completion: completion)
                    } else if let error = info?[PHImageErrorKey] as? Error {
                        completion(.failure(.videoExport(error)))
                    } else {
                        completion(.failure(.noVideoExportSession))
                    }
                }
            } 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 {
                completion(.failure(.noVideoExportSession))
                return
            }
            CompositionAttachmentData.exportVideoData(session: session, completion: completion)
            
        case let .drawing(drawing):
            let image = drawing.imageInLightMode(from: drawing.bounds, scale: 1)
            completion(.success((image.pngData()!, .png)))
        case let .gif(data):
            completion(.success((data, .gif)))
        }
    }
    
    private static func exportVideoData(session: AVAssetExportSession, completion: @escaping (Result<(Data, UTType), 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(.videoExport(session.error!)))
                return
            }
            do {
                let data = try Data(contentsOf: session.outputURL!)
                completion(.success((data, .mpeg4Movie)))
            } catch {
                completion(.failure(.videoExport(error)))
            }
        }
    }
    
    enum AttachmentType {
        case image, video
    }
    
    enum Error: Swift.Error, LocalizedError {
        case missingData
        case videoExport(Swift.Error)
        case noVideoExportSession
        
        var localizedDescription: String {
            switch self {
            case .missingData:
                return "Missing Data"
            case .videoExport(let error):
                return "Exporting video: \(error)"
            case .noVideoExportSession:
                return "Couldn't create video export session"
            }
        }
    }
}

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)
        case .gif(_):
            throw EncodingError.invalidValue(self, EncodingError.Context(codingPath: [], debugDescription: "gif CompositionAttachments cannot be encoded"))
        }
    }
    
    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, asset, or drawing")
        }
    }
    
    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
        }
    }
}