import math

from Quartz import (
    CFURLCreateWithFileSystemPath,
    CGBitmapContextCreate,
    CGBitmapContextCreateImage,
    CGColorSpaceCreateWithName,
    CGImageDestinationAddImage,
    CGImageDestinationCreateWithURL,
    CGImageDestinationFinalize,
    CGImageSourceCopyPropertiesAtIndex,
    CGImageSourceCreateImageAtIndex,
    CGImageSourceCreateWithURL,
    CGImageSourceGetCount,
    CIContext,
    CIFilter,
    CIImage,
    CIVector,
    NSAffineTransform,
    kCFURLPOSIXPathStyle,
    kCGColorSpaceGenericRGB,
    kCGImageAlphaPremultipliedLast,
    kCIInputAspectRatioKey,
    kCIInputBackgroundImageKey,
    kCIInputImageKey,
    kCIInputScaleKey,
    kCIOutputImageKey,
)

_REMOVABLE_DISK_PATH = (
    "/System/Library/Extensions/IOStorageFamily.kext/Contents/Resources/Removable.icns"
)


def badge_disk_icon(badge_file, output_file):
    # Load the Removable disk icon
    url = CFURLCreateWithFileSystemPath(
        None, _REMOVABLE_DISK_PATH, kCFURLPOSIXPathStyle, False
    )
    backdrop = CGImageSourceCreateWithURL(url, None)
    backdropCount = CGImageSourceGetCount(backdrop)

    # Load the badge
    url = CFURLCreateWithFileSystemPath(None, badge_file, kCFURLPOSIXPathStyle, False)
    badge = CGImageSourceCreateWithURL(url, None)
    assert badge is not None, "Unable to process image file: %s" % badge_file
    badgeCount = CGImageSourceGetCount(badge)

    # Set up a destination for our target
    url = CFURLCreateWithFileSystemPath(None, output_file, kCFURLPOSIXPathStyle, False)
    target = CGImageDestinationCreateWithURL(url, "com.apple.icns", backdropCount, None)

    # Get the RGB colorspace
    rgbColorSpace = CGColorSpaceCreateWithName(kCGColorSpaceGenericRGB)

    # Scale
    scale = 1.0

    # Perspective transform
    corners = ((0.2, 0.95), (0.8, 0.95), (0.85, 0.35), (0.15, 0.35))

    # Translation
    position = (0.5, 0.5)

    for n in range(backdropCount):
        props = CGImageSourceCopyPropertiesAtIndex(backdrop, n, None)
        width = props["PixelWidth"]
        height = props["PixelHeight"]
        dpi = props["DPIWidth"]
        depth = props["Depth"]

        # Choose the best sized badge image
        bestWidth = None
        # ?? bestHeight = None
        bestBadge = None
        bestDepth = None
        # ?? bestDPI = None
        for m in range(badgeCount):
            badgeProps = CGImageSourceCopyPropertiesAtIndex(badge, m, None)
            badgeWidth = badgeProps["PixelWidth"]
            # ?? badgeHeight = badgeProps['PixelHeight']
            badgeDPI = badgeProps["DPIWidth"]
            badgeDepth = badgeProps["Depth"]

            if bestBadge is None or (
                badgeWidth <= width
                and (
                    bestWidth > width
                    or badgeWidth > bestWidth
                    or (
                        badgeWidth == bestWidth
                        and badgeDPI == dpi
                        and badgeDepth <= depth
                        and (bestDepth is None or badgeDepth > bestDepth)
                    )
                )
            ):
                bestBadge = m
                bestWidth = badgeWidth
                # ?? bestHeight = badgeHeight
                # ?? bestDPI = badgeDPI
                bestDepth = badgeDepth

        badgeImage = CGImageSourceCreateImageAtIndex(badge, bestBadge, None)
        badgeCI = CIImage.imageWithCGImage_(badgeImage)

        backgroundImage = CGImageSourceCreateImageAtIndex(backdrop, n, None)
        backgroundCI = CIImage.imageWithCGImage_(backgroundImage)

        compositor = CIFilter.filterWithName_("CISourceOverCompositing")
        lanczos = CIFilter.filterWithName_("CILanczosScaleTransform")
        perspective = CIFilter.filterWithName_("CIPerspectiveTransform")
        transform = CIFilter.filterWithName_("CIAffineTransform")

        lanczos.setValue_forKey_(badgeCI, kCIInputImageKey)
        lanczos.setValue_forKey_(scale * float(width) / bestWidth, kCIInputScaleKey)
        lanczos.setValue_forKey_(1.0, kCIInputAspectRatioKey)

        topLeft = (width * scale * corners[0][0], width * scale * corners[0][1])
        topRight = (width * scale * corners[1][0], width * scale * corners[1][1])
        bottomRight = (width * scale * corners[2][0], width * scale * corners[2][1])
        bottomLeft = (width * scale * corners[3][0], width * scale * corners[3][1])

        out = lanczos.valueForKey_(kCIOutputImageKey)
        if width >= 16:
            perspective.setValue_forKey_(out, kCIInputImageKey)
            perspective.setValue_forKey_(
                CIVector.vectorWithX_Y_(*topLeft), "inputTopLeft"
            )
            perspective.setValue_forKey_(
                CIVector.vectorWithX_Y_(*topRight), "inputTopRight"
            )
            perspective.setValue_forKey_(
                CIVector.vectorWithX_Y_(*bottomRight), "inputBottomRight"
            )
            perspective.setValue_forKey_(
                CIVector.vectorWithX_Y_(*bottomLeft), "inputBottomLeft"
            )
            out = perspective.valueForKey_(kCIOutputImageKey)

        tfm = NSAffineTransform.transform()
        tfm.translateXBy_yBy_(
            math.floor((position[0] - 0.5 * scale) * width),
            math.floor((position[1] - 0.5 * scale) * height),
        )

        transform.setValue_forKey_(out, kCIInputImageKey)
        transform.setValue_forKey_(tfm, "inputTransform")
        out = transform.valueForKey_(kCIOutputImageKey)

        compositor.setValue_forKey_(out, kCIInputImageKey)
        compositor.setValue_forKey_(backgroundCI, kCIInputBackgroundImageKey)

        result = compositor.valueForKey_(kCIOutputImageKey)

        cgContext = CGBitmapContextCreate(
            None, width, height, 8, 0, rgbColorSpace, kCGImageAlphaPremultipliedLast
        )
        context = CIContext.contextWithCGContext_options_(cgContext, None)

        context.drawImage_inRect_fromRect_(
            result, ((0, 0), (width, height)), ((0, 0), (width, height))
        )

        image = CGBitmapContextCreateImage(cgContext)

        CGImageDestinationAddImage(target, image, props)

    CGImageDestinationFinalize(target)