//
//  AttachmentData.swift
//  ComposeUI
//
//  Created by Shadowfacts on 1/1/20.
//  Copyright © 2020 Shadowfacts. All rights reserved.
//

import UIKit
import Photos
import UniformTypeIdentifiers
import PencilKit
import InstanceFeatures

enum AttachmentData {
    case asset(PHAsset)
    case image(Data, originalType: UTType)
    case video(URL)
    case drawing(PKDrawing)
    case gif(Data)
    
    var type: AttachmentType {
        switch self {
        case let .asset(asset):
            return asset.attachmentType!
        case .image(_, originalType: _):
            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(originalData, originalType):
            let data: Data
            let type: UTType
            switch originalType {
            case .png, .jpeg:
                data = originalData
                type =  originalType
            default:
                let image = UIImage(data: originalData)!
                // The quality of 0.8 was chosen completely arbitrarily, it may need to be tuned in the future.
                data = image.jpegData(compressionQuality: 0.8)!
                type = .jpeg
            }
            let processed = processImageData(data, type: type, features: features, skipAllConversion: skipAllConversion)
            completion(.success(processed))
        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 let data = data, let dataUTI = dataUTI else {
                        completion(.failure(.missingData))
                        return
                    }
                    let processed = processImageData(data, type: UTType(dataUTI)!, features: features, skipAllConversion: skipAllConversion)
                    completion(.success(processed))
                }
            } 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 {
                        AttachmentData.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
            }
            AttachmentData.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 func processImageData(_ data: Data, type: UTType, features: InstanceFeatures, skipAllConversion: Bool) -> (Data, UTType) {
        guard !skipAllConversion else {
            return (data, type)
        }
        
        var data = data
        var type = type
        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 || type == .heic {
            let context = CIContext()
            let colorSpace = needsColorSpaceConversion || image.colorSpace != nil ? CGColorSpace(name: CGColorSpace.sRGB)! : image.colorSpace!
            if type == .png {
                data = context.pngRepresentation(of: image, format: .ARGB8, colorSpace: colorSpace)!
            } else {
                data = context.jpegRepresentation(of: image, colorSpace: colorSpace)!
                type = .jpeg
            }
        }
        
        return (data, type)
    }
    
    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: AttachmentData.AttachmentType? {
        switch self.mediaType {
        case .image:
            return .image
        case .video:
            return .video
        default:
            return nil
        }
    }
}

extension AttachmentData: 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(originalData, originalType):
            try container.encode("image", forKey: .type)
            try container.encode(originalType, forKey: .imageType)
            try container.encode(originalData, forKey: .imageData)
        case .video(_):
            throw EncodingError.invalidValue(self, EncodingError.Context(codingPath: encoder.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: encoder.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":
            let data = try container.decode(Data.self, forKey: .imageData)
            if let type = try container.decodeIfPresent(UTType.self, forKey: .imageType) {
                self = .image(data, originalType: type)
            } else {
                guard let image = UIImage(data: data) else {
                    throw DecodingError.dataCorruptedError(forKey: .imageData, in: container, debugDescription: "CompositionAttachment data could not be decoded into UIImage")
                }
                let jpegData = image.jpegData(compressionQuality: 1)!
                self = .image(jpegData, originalType: .jpeg)
            }
        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
        case imageType
        /// The local identifier of the PHAsset for this attachment
        case assetIdentifier
        /// The PKDrawing object for this attachment.
        case drawing
    }
}

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