PDF_Office/PDF Office/PDF Master/Class/PDFTools/OCRNew/Model/KMGOCROperation.swift

//
//  KMGOCROperation.swift
//  PDF Reader Pro
//
//  Created by liujiajie on 2023/11/15.
//

import Foundation

@objc(KMGOCROperationDelegate)
protocol KMGOCROperationDelegate: AnyObject {
    @objc optional func GOCROperation(_ operation: KMGOCROperation, cancelOCRImageAtIndex index: Int)
    @objc optional func GOCROperation(_ operation: KMGOCROperation, startOCRImageAtIndex index: Int)
    @objc optional func GOCROperation(_ operation: KMGOCROperation, finishOCRImageAtIndex index: Int, results: Array<Any>)
    @objc optional func GOCROperation(_ operation: KMGOCROperation, failureOCRImageAtIndex index: Int, error: Error?)

}

let KMGOC_API_URL = "https://vision.googleapis.com/v1/images:annotate"

#if VERSION_DMG
let KMGOC_API_KEY = "AIzaSyBhSRohpngAu8pSgFDXPytslNDHgGm7uDs"
#else
let KMGOC_API_KEY = "AIzaSyCJuqJ9YvtkFKMl1mW3Yq-av3mmI9ScbRY"
#endif


@objcMembers class KMGOCROperation: Operation{
    @objc var operationDelegate: KMGOCROperationDelegate?
    var hasCanceled: Bool = false
    var hasFinished: Bool = false
    var hasExcuting: Bool = false
    var selectedLanguages: Array<Any>?
    override var isExecuting: Bool{
        return self.hasExcuting
    }
    override var isFinished: Bool{
        return self.hasFinished
    }
    override var isCancelled: Bool{
        return self.hasCanceled
    }

    var fileName: String = ""
    var orcImage: NSImage?
    var task: URLSessionDataTask?
    var imageIndex: Int = 0
    
    init(recognitionImg:NSImage, imgIndex:Int) {
        super.init()
        self.imageIndex = imgIndex
        self.fileName = fileNameWithDate()
        self.orcImage = recognitionImg
        self.queuePriority = .normal
        self.name = self.fileName
        self.hasExcuting = false
        self.hasFinished = false
    }
    func fileNameWithDate() -> String {
        let formatter = DateFormatter()
        formatter.dateFormat = "YYYY-MM-dd-hh-mm-ss-SSS"
        let dateString = formatter.string(from: Date())
        let fileName = "\(dateString) \(imageIndex)"
        return fileName
    }
    
    override func start() {
        if p_checkCancelled() { return }
        willChangeValue(forKey: "isExecuting")
        self.hasExcuting = true
        didChangeValue(forKey: "isExecuting")
        
//        Thread.detachNewThreadSelector(#selector(main), toTarget: self, with: nil)
//    }
//    @objc override func main() {
//        do {
            if p_checkCancelled() { return }
            recognitionImage(orcImage ?? NSImage())
//            while isExecuting {
//                if p_checkCancelled() {
//                    return
//                }
//            }
//        } catch let e {
//            Swift.debugPrint("Exception %@", e)
//        }
    }
    override func cancel() { 
        super.cancel()
        if task != nil {
            task?.cancel()
            task = nil
        }
        
        self.operationDelegate?.GOCROperation?(self, cancelOCRImageAtIndex: self.imageIndex)
        if isExecuting {
            willChangeValue(forKey: "isFinished")
            hasFinished = true
            didChangeValue(forKey: "isFinished")
        } else {
            willChangeValue(forKey: "isCancelled")
            hasCanceled = true
            didChangeValue(forKey: "isCancelled")
        }
    }
    func p_done() {
        self.willChangeValue(forKey: "isFinished")
        self.hasFinished = true
        self.didChangeValue(forKey: "isFinished")
    }
    func p_checkCancelled() -> Bool { 
        if isCancelled {
            willChangeValue(forKey: "isFinished")
            willChangeValue(forKey: "isExecuting")
            hasExcuting = false
            hasFinished = true
            didChangeValue(forKey: "isExecuting")
            didChangeValue(forKey: "isFinished")
            return true
        }
        return false
    }
    func recognitionImage(_ image: NSImage) {
        self.operationDelegate?.GOCROperation?(self, startOCRImageAtIndex: self.imageIndex)
        let binaryImageData = base64EncodeImage(image)
        
        if binaryImageData == nil {
            self.operationDelegate?.GOCROperation?(self, failureOCRImageAtIndex: self.imageIndex, error: nil)
            return
        }
        let urlString = "\(KMGOC_API_URL)?key=\(KMGOC_API_KEY)"
        var request = URLRequest(url: URL(string: urlString)!)
        request.httpMethod = "POST"
        request.addValue("application/json", forHTTPHeaderField: "Content-Type")
        
        let imageDictionary = ["content": binaryImageData]
        let featuresArray = [["type": "TEXT_DETECTION", "maxResults": 10]]
        var paramsDictionary = ["requests": [["image": imageDictionary, "features": featuresArray]]] as [String: Any]
        if selectedLanguages != nil && selectedLanguages?.count ?? 0 > 0 {
            let imageContextDictionary = ["languageHints": selectedLanguages]
            paramsDictionary = ["requests": [["image": imageDictionary, "features": featuresArray, "imageContext": imageContextDictionary]]] as [String: Any]
        }
        
        let requestData = try? JSONSerialization.data(withJSONObject: paramsDictionary, options: [])
        request.httpBody = requestData
        
        let URLSession = URLSession.shared
        task = URLSession.dataTask(with: request) { (data, response, error) in
            if (error as NSError?)?.code == NSURLErrorCancelled {
                return
            }
            var results: [Any]? = nil
            if error == nil {
                let dictionary = try? JSONSerialization.jsonObject(with: data!, options: []) as? [String: Any]
                results = self.responseDataResults(dictionary! as NSDictionary)
            }
            
            DispatchQueue.main.async {
                if error != nil || results == nil {
                    self.operationDelegate?.GOCROperation?(self, failureOCRImageAtIndex: self.imageIndex, error: error)
                } else {
                    self.operationDelegate?.GOCROperation?(self, finishOCRImageAtIndex: self.imageIndex, results: results ?? [])
                }
            }
            if error != nil || results == nil {
                self.cancel()
            } else {
                self.p_done()
            }
        }
        task?.resume()
    }
    
    func base64EncodeImage(_ image: NSImage) -> String? {
        // 获取 TIFF 数据
        guard let data = image.tiffRepresentation else { return nil }
        
        // 创建位图表示
        guard let imageRep = NSBitmapImageRep(data: data) else { return nil }
        imageRep.size = image.size
        
        // 转换为 PNG 数据
        guard let imageData = imageRep.representation(using: .png, properties: [:]) else { return nil }
        
        // 如果数据大小超过限制，尝试压缩
        let finalData: Data
        if imageData.count > 4194304 {
            guard let compressedData = compressImageData(imageData, toMaxFileSize: 4194304) else {
                return nil // 压缩失败时返回 nil
            }
            finalData = compressedData
        } else {
            finalData = imageData
        }
        
        // 返回 Base64 编码字符串
        return finalData.base64EncodedString(options: .lineLength64Characters)
    }

    func compressImageData(_ imageData: Data, toMaxFileSize maxFileSize: Int) -> Data? {
        // 初始压缩比
        var compression: CGFloat = 0.9
        let maxCompression: CGFloat = 0.1
        
        // 循环压缩
        var currentData = imageData
        while currentData.count > maxFileSize && compression > maxCompression {
            compression -= 0.1
            
            // 创建位图表示
            guard let imageRep = NSBitmapImageRep(data: currentData) else {
                return nil // 解码失败时返回 nil
            }
            
            // 尝试压缩为 JPEG 数据
            guard let compressedData = imageRep.representation(
                using: .jpeg,
                properties: [.compressionFactor: NSNumber(value: Float(compression))]
            ) else {
                return nil // 压缩失败时返回 nil
            }
            
            currentData = compressedData
        }
        
        return currentData.count <= maxFileSize ? currentData : nil
    }
//
//    func base64EncodeImage(_ image: NSImage) -> String? {
//        guard let data = image.tiffRepresentation else { return nil }
//        let imageRep = NSBitmapImageRep(data: data)!
//        imageRep.size = image.size
//        let imageData = imageRep.representation(using: .png, properties: [:])
//        // Resize the image if it exceeds the 4MB API limit
//        if imageData?.count ?? 0 > 4194304 {
//            let compressedData = compressImageData(imageData!, toMaxFileSize: 4194304)
//            if let data = compressedData {
//                return data.base64EncodedString(options: .endLineWithCarriageReturn)
//            }
//        }
//        
//        if let data = imageData {
//            if #available(macOS 10.9, *) {
//                return data.base64EncodedString(options: .endLineWithCarriageReturn)
//            } else {
//                return data.base64EncodedString(options: [])
//            }
//        }
//        
//        return nil
//    }
//    
//    func compressImageData(_ imageData: Data, toMaxFileSize maxFileSize: Int) -> Data? {
//        var compression: CGFloat = 0.9
//        let maxCompression: CGFloat = 0.1
//        var compressImageData = imageData
//        while compressImageData.count > maxFileSize && compression > maxCompression {
//            compression -= 0.1
//            let imageRep = NSBitmapImageRep(data: compressImageData)!
//            compressImageData = imageRep.representation(using: .jpeg, properties: [NSBitmapImageRep.PropertyKey.compressionFactor: NSNumber(value: Float(compression))])!
//        }
//        return compressImageData
//    }
    
    func responseDataResults(_ dictionary: NSDictionary) -> [Any]? {
        let responses = dictionary["responses"] as? [Any]
        if responses == nil {
            return nil
        }
        let responseData = responses?.first as? NSDictionary
        if let errorObj = dictionary["error"] as? NSDictionary {
            return nil
        }
        var results: [Any]? = nil
        if let textAnnotations = responseData?["textAnnotations"] as? [Any] {
            results = [Any]()
            for annotation in textAnnotations {
                var textBounds = CGRect.zero
                let dic = (annotation as? NSDictionary)?["boundingPoly"]
                if let vertices = (dic as? NSDictionary)?["vertices"] as? [Any] {
                    var minX: CGFloat = 0, minY: CGFloat = 0, maxX: CGFloat = 0, maxY: CGFloat = 0
                    for i in 0..<vertices.count {
                        if let vertex = vertices[i] as? NSDictionary,
                           let x = vertex["x"] as? CGFloat,
                           let y = vertex["y"] as? CGFloat {
                            minX = i == 0 ? x : min(x, minX)
                            minY = i == 0 ? y : min(y, minY)
                            maxX = i == 0 ? x : max(x, maxX)
                            maxY = i == 0 ? y : max(y, maxY)
                        }
                    }
                    textBounds = CGRect(x: minX, y: minY, width: maxX-minX, height: maxY-minY)
                }
                
                let result = KMGOCRResult()
                result.text = (annotation as? NSDictionary)?["description"] as? String ?? ""
                result.locale = (annotation as? NSDictionary)?["locale"] as? String ?? ""
                result.textBounds = textBounds
                results?.append(result)
            }
        }
        return results
    }
}