//
// Copyright 2024 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//

import Foundation
import SDWebImage
public import UIKit

extension UIImage {

    public static func from(
        _ attachment: AttachmentStream,
    ) throws -> UIImage {
        return try .fromEncryptedFile(
            at: attachment.fileURL,
            attachmentKey: AttachmentKey(combinedKey: attachment.attachment.encryptionKey),
            plaintextLength: attachment.unencryptedByteCount,
            mimeType: attachment.mimeType,
        )
    }

    public static func from(
        _ attachmentThumbnail: AttachmentBackupThumbnail,
    ) throws -> UIImage {
        return try .fromEncryptedFile(
            at: attachmentThumbnail.fileURL,
            attachmentKey: AttachmentKey(combinedKey: attachmentThumbnail.attachment.encryptionKey),
            plaintextLength: nil,
            mimeType: MimeType.imageWebp.rawValue,
        )
    }

    /// If no plaintext length is provided, the file is assumed to only use pkcs7 padding.
    static func fromEncryptedFile(
        at fileURL: URL,
        attachmentKey: AttachmentKey,
        plaintextLength: UInt32?,
        mimeType: String,
    ) throws -> UIImage {
        if
            mimeType.caseInsensitiveCompare(MimeType.imageJpeg.rawValue) == .orderedSame,
            /// We can use a CGDataProvider. UIImage tends to load the whole thing into memory _anyway_,
            /// but this at least makes it possible for it to choose not to.
            let jpegImage = try? CGDataProvider.loadFromEncryptedFile(
                at: fileURL,
                attachmentKey: attachmentKey,
                plaintextLength: plaintextLength,
                block: { dataProvider in
                    let (cgImage, orientation) = try dataProvider.toJpegCGImage()
                    return UIImage(cgImage: cgImage, scale: 1, orientation: orientation)
                },
            )
        {
            return jpegImage
        }
        if
            mimeType.caseInsensitiveCompare(MimeType.imagePng.rawValue) == .orderedSame,
            /// We can use a CGDataProvider. UIImage tends to load the whole thing into memory _anyway_,
            /// but this at least makes it possible for it to choose not to.
            let pngImage = try? CGDataProvider.loadFromEncryptedFile(
                at: fileURL,
                attachmentKey: attachmentKey,
                plaintextLength: plaintextLength,
                block: { dataProvider in
                    return UIImage(cgImage: try dataProvider.toPngCGImage())
                },
            )
        {
            return pngImage
        }

        // hmac and digest are validated at download time; no need to revalidate every read.
        let data = try Cryptography.decryptFileWithoutValidating(
            at: fileURL,
            metadata: DecryptionMetadata(
                key: attachmentKey,
                plaintextLength: plaintextLength.map(UInt64.init(safeCast:)),
            ),
        )
        let image: UIImage?
        if mimeType.caseInsensitiveCompare(MimeType.imageWebp.rawValue) == .orderedSame {
            /// Use SDAnimatedImage for webp.
            image = SDAnimatedImage(data: data)
        } else {
            image = UIImage(data: data)
        }

        guard let image else {
            throw OWSAssertionError("Failed to load image")
        }
        return image
    }
}

extension CGDataProvider {
    /// If no plaintext length is provided, the file is assumed to only use pkcs7 padding.
    fileprivate static func loadFromEncryptedFile<T>(
        at fileURL: URL,
        attachmentKey: AttachmentKey,
        plaintextLength: UInt32?,
        block: (CGDataProvider) throws -> T,
    ) throws -> T {
        let fileHandle: EncryptedFileHandle
        if let plaintextLength {
            fileHandle = try Cryptography.encryptedAttachmentFileHandle(
                at: fileURL,
                plaintextLength: UInt64(safeCast: plaintextLength),
                attachmentKey: attachmentKey,
            )
        } else {
            fileHandle = try Cryptography.encryptedFileHandle(
                at: fileURL,
                attachmentKey: attachmentKey,
            )
        }
        let dataProvider = CGDataProvider.from(fileHandle: fileHandle, fileSize: Int64(fileHandle.plaintextLength))
        guard let dataProvider else {
            throw OWSAssertionError("couldn't initialize encrypted data provider")
        }
        return try block(dataProvider)
    }
}

extension CGDataProvider {

    enum ParsingError: Error {
        case failedToParsePng
        case failedToParseJpg
    }

    fileprivate func toPngCGImage() throws -> CGImage {
        guard
            let cgImage = CGImage(
                pngDataProviderSource: self,
                decode: nil,
                shouldInterpolate: true,
                intent: .defaultIntent,
            )
        else {
            throw ParsingError.failedToParsePng
        }
        return cgImage
    }

    fileprivate func toJpegCGImage() throws -> (CGImage, UIImage.Orientation) {
        let orientation: UIImage.Orientation = {
            guard let imageSource = CGImageSourceCreateWithDataProvider(self, nil) else {
                return nil
            }
            // Get image orientation
            let options: [CFString: Any] = [
                kCGImageSourceShouldAllowFloat: true,
            ]
            let properties = CGImageSourceCopyPropertiesAtIndex(
                imageSource,
                0,
                options as CFDictionary,
            ) as? [CFString: Any]
            guard
                let raw = properties?[kCGImagePropertyOrientation] as? Int,
                let raw = UInt32(exactly: raw)
            else {
                return nil
            }
            return CGImagePropertyOrientation(rawValue: raw)?.uiImageOrientation
        }() ?? .up

        guard
            let cgImage = CGImage(
                jpegDataProviderSource: self,
                decode: nil,
                shouldInterpolate: true,
                intent: .defaultIntent,
            )
        else {
            throw ParsingError.failedToParseJpg
        }
        return (cgImage, orientation)
    }
}

extension CGImagePropertyOrientation {

    var uiImageOrientation: UIImage.Orientation {
        switch self {
        case .up:
            return .up
        case .down:
            return .down
        case .left:
            return .left
        case .right:
            return .right
        case .upMirrored:
            return .upMirrored
        case .downMirrored:
            return .downMirrored
        case .leftMirrored:
            return .leftMirrored
        case .rightMirrored:
            return .rightMirrored
        }
    }
}
