#include <Cocoa/Cocoa.h>
#include <QPainter>
#include <QtMac>

QImage captureFrameCocoa(const QRect geometry)
{
	const auto cgGeometry = CGRectMake(geometry.x(), geometry.y(), geometry.width(), geometry.height());
	const int maxDisplays = 128; // 128 displays should be enough for everyone. According to Qt devs, anyway.
	CGDirectDisplayID displays[maxDisplays];
	CGDisplayCount displayCount;
	const auto err = CGGetDisplaysWithRect(cgGeometry, maxDisplays, displays, &displayCount);

	if (err && displayCount == 0)
		return {};

	// Find the display with the highest scaling factor and use this factor for image creation
	// We could limit this to the displays covering our geometry, but this would make the decision if we want to paint a retina cursor harder
	const NSArray* screens = [NSScreen screens];
	qreal maxScale         = 1;
	for (int i = 0; i < [screens count]; i++)
	{
		const NSScreen* screen = [screens objectAtIndex:i];
		const qreal scale      = [screen backingScaleFactor];
		maxScale               = std::max(maxScale, scale);
	}

	const auto realWidth  = static_cast<int>(geometry.width() * maxScale) & ~1;
	const auto realHeight = static_cast<int>(geometry.height() * maxScale) & ~1;

	QPixmap pixmap{realWidth, realHeight};
	QPainter painter{&pixmap};
	if (displayCount > 1)
		painter.fillRect(pixmap.rect(), Qt::black); // For non-rectangular screen setups

	for (int i = 0; i < displayCount; i++)
	{
		// All of these are in logical units, not pixels
		const auto cgBounds = CGDisplayBounds(displays[i]);
		const auto qtBounds = QRect{qRound(cgBounds.origin.x), qRound(cgBounds.origin.y), qRound(cgBounds.size.width),
		                            qRound(cgBounds.size.height)};

		const auto cgCaptureRegion =
		    CGRectOffset(CGRectIntersection(cgBounds, cgGeometry), -cgBounds.origin.x, -cgBounds.origin.y);
		const auto qtCaptureRegion = qtBounds.intersected(geometry).translated(-geometry.topLeft());

		const auto cgImage = CGDisplayCreateImageForRect(displays[i], cgCaptureRegion);
		const auto qtImage = QtMac::fromCGImageRef(cgImage);

		painter.drawPixmap(QRect{qtCaptureRegion.topLeft() * maxScale, qtCaptureRegion.size() * maxScale}, qtImage);
		CGImageRelease(cgImage);
	}
	return pixmap.toImage();
}
