import Foundation

typealias NewFileCallback = (fileURL: NSURL) -> Void

class ScreenshotDetector: NSObject, NSMetadataQueryDelegate {
    
    let query = NSMetadataQuery()
    
    var newFileCallback: NewFileCallback?

    override init() {
        super.init()
        
        let center = NSNotificationCenter.defaultCenter()
        center.addObserver(self, selector: Selector("queryUpdated:"), name: NSMetadataQueryDidStartGatheringNotification, object: query)
        center.addObserver(self, selector: Selector("queryUpdated:"), name: NSMetadataQueryDidUpdateNotification, object: query)
        center.addObserver(self, selector: Selector("queryUpdated:"), name: NSMetadataQueryDidFinishGatheringNotification, object: query)
        
        query.delegate = self
        query.predicate = NSPredicate(format: "kMDItemIsScreenCapture = 1")
        query.startQuery()
    }
    
    deinit {
        query.stopQuery()
    }
    
    func queryUpdated(notification: NSNotification) {
        if let userInfo = notification.userInfo {
            for v in userInfo.values {
                let items = v as! [NSMetadataItem]
                if items.count > 0 {
                    let item = items[0]
                    if let filename = item.valueForAttribute("kMDItemFSName") as? String {
                        let filenameWithPath = NSString(string: "~/Desktop/" + filename).stringByExpandingTildeInPath
                        let url = NSURL(fileURLWithPath: filenameWithPath, isDirectory: false)
                        if let cb = self.newFileCallback {
                            cb(fileURL: url)
                        }
                    }
                }
            }
        }
    }
}

////////////////////////////
// Usage                  //
////////////////////////////

// Because this object is observing events it should stay
// in memory for the duration of its usage.
// In a view controller for instance it should be allocated
// in the class scope. If its reference is stored inside
// a method it will be deallocated automatically when the
// method returns and consequently stop obsering any events.
let detector = ScreenshotDetector()

// The actual callback can be assigned anywhere there's
// a reference to the instance.
detector.newFileCallback = { fileURL in
  // fileURL is an NSURL instance of the screenshot file.
}