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

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

import Foundation

@objc(KMGOCROperationDelegate)
protocol KMGOCROperationDelegate: AnyObject {
    @objc optional func GOCROperationCancel(_ op: KMGOCROperation, atIndex: Int)
    @objc optional func GOCROperationStart(_ op: KMGOCROperation, atIndex: Int)
    @objc optional func GOCROperationFinish(_ op: KMGOCROperation, atIndex: Int, results: Array<Any>)
    @objc optional func GOCROperationFail(_ op: KMGOCROperation, atIndex: Int, err: 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 selectedLanguages: Array<Any>?
    override var isExecuting: Bool{
        set{
            self.willChangeValue(forKey: "isExecuting")
            self.isExecuting = newValue
            didChangeValue(forKey: "isExecuting")
        }
        get{
            return super.isExecuting
        }
    }
    override var isFinished: Bool{
        set{
            self.willChangeValue(forKey: "isFinished")
            self.isFinished = newValue
            didChangeValue(forKey: "isFinished")
        }
        get{
            return super.isFinished
        }
    }
    override var isCancelled: Bool{
        set{
            self.willChangeValue(forKey: "isCancelled")
            self.isCancelled = newValue
            didChangeValue(forKey: "isCancelled")
        }
        get{
            return super.isCancelled
        }
    }

    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.isExecuting = false
        self.isFinished = 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 }
        isExecuting = true
        
        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?.GOCROperationCancel?(self, atIndex: self.imageIndex)
        
        if isExecuting {
            isExecuting = false
            isFinished = true
        } else {
            isFinished = false
        }
        isCancelled = true
    }
    func p_done() {
        isExecuting = false
        isFinished = true
    }
    func p_checkCancelled() -> Bool { 
        if isCancelled {
            isFinished = true
            return true
        }
        return false
    }
    func recognitionImage(_ image: NSImage) {
        self.operationDelegate?.GOCROperationStart?(self, atIndex: self.imageIndex)
        
        let binaryImageData = base64EncodeImage(image)
        
        if binaryImageData == nil {
            self.operationDelegate?.GOCROperationFail?(self, atIndex: self.imageIndex, err: 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?.GOCROperationFail?(self, atIndex: self.imageIndex, err: error)
                } else {
                    self.operationDelegate?.GOCROperationFinish?(self, atIndex: self.imageIndex, results: results ?? [])
                }
            }
            if error != nil || results == nil {
                self.cancel()
            } else {
                self.p_done()
            }
        }
        task?.resume()
    }
    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]? { 
        guard let responses = dictionary["responses"] as? [Any],
                let responseData = responses.first as? NSDictionary,
              let errorObj = dictionary["error"] as? NSDictionary else { 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
    }
}