from CoreGraphics import *
import sys
import os
import subprocess
import tempfile
from Carbon.File import FSPathMakeRef
from Carbon.Folder import FSFindFolder
from Carbon.Folders import kUserDomain, kDesktopFolderType
from Carbon.CF import CFURLCreateFromFileSystemRepresentation
from LaunchServices.Launch import LSSetExtensionHiddenForRef

def check(retval):
    if retval != 0:
        sys.exit(retval)

def dumpPNG(pdf, pageNumber, outFileBase):
    if pageNumber is None:
        outFile = outFileBase + ".png"
        pageNumber = 1
    else:
        outFile = outFileBase + "-%03d.png" % pageNumber

    w = pdf.getTrimBox(pageNumber).getWidth()
    h = pdf.getTrimBox(pageNumber).getHeight()

    ctx = CGBitmapContextCreateWithColor(int(w), int(h), CGColorSpaceCreateDeviceRGB(), (0,0,0,0))
    ctx.drawPDFDocument(pdf.getTrimBox(pageNumber), pdf, pageNumber)
    ctx.writeToFile(outFile, kCGImageFormatPNG)
    return outFile

tempDir = tempfile.mkdtemp()
pdfFile = os.path.join(tempDir, 'screenshot.pdf')

check(subprocess.call(['/usr/sbin/screencapture', '-iW', pdfFile]))

pdf = CGPDFDocumentCreateWithProvider(CGDataProviderCreateWithFilename(pdfFile))
numberOfPages = pdf.getNumberOfPages()

desktopDir = FSFindFolder(kUserDomain, kDesktopFolderType, False).as_pathname()
desktopContents = os.listdir(desktopDir)
index = 1
outFile = None
while outFile is None:
    outFile = 'Screenshot %d' % index
    for filePath in desktopContents:
        if filePath.startswith(outFile):
            index += 1
            outFile = None
            break
outFile = os.path.join(desktopDir, outFile)

writtenPaths = []
if numberOfPages == 1:
    writtenPaths = [dumpPNG(pdf, None, outFile)]
else:
    for pageNumber in xrange(1, numberOfPages+1):
        writtenPaths.append(dumpPNG(pdf, pageNumber, outFile))

os.remove(pdfFile)
os.rmdir(tempDir)

for path in writtenPaths:
    LSSetExtensionHiddenForRef(FSPathMakeRef(path)[0], True)

# Carbon.File module doesn't implement FNNotify
subprocess.call(['/usr/bin/osascript', '-e', 'tell app "Finder" to update desktop'])

# LaunchServices module doesn't implement LSOpenFrom*Spec
subprocess.call([os.path.join(os.getenv('RESOURCEPATH'), 'launch'), '-i', 'tv.kungfoo.1001'] + writtenPaths)