import AVFoundation
import CoreGraphics
import CoreImage
import CoreMedia
import Foundation
import UIKit

extension CMSampleBuffer {
    var cgImage: CGImage? {
        let pixelBuffer: CVPixelBuffer? = CMSampleBufferGetImageBuffer(self)

        guard let imagePixelBuffer = pixelBuffer else {
            return nil
        }

        return CIImage(cvPixelBuffer: imagePixelBuffer).cgImage
    }
}

extension CIImage {
    var cgImage: CGImage? {
        let ciContext = CIContext()

        guard let cgImage = ciContext.createCGImage(self, from: self.extent) else {
            return nil
        }

        return cgImage
    }
}

extension UIImage {
    func rotate(to orientation: UIImage.Orientation) -> UIImage? {
        UIGraphicsBeginImageContext(size)
        draw(in: CGRect(origin: .zero, size: size))
        let normalizedImage = UIGraphicsGetImageFromCurrentImageContext()
        UIGraphicsEndImageContext()
        return normalizedImage?.withOrientation(orientation)
    }

    func withOrientation(_ orientation: UIImage.Orientation) -> UIImage? {
        guard let cgImage = self.cgImage else { return nil }
        return UIImage(cgImage: cgImage, scale: scale, orientation: orientation)
    }
}

extension CGImage {
    func rotated(byDegrees degrees: CGFloat) -> CGImage? {
        let radians = degrees * .pi / 180
        let imageRect = CGRect(origin: .zero, size: CGSize(width: self.width, height: self.height))
        var rotatedRect = imageRect.applying(CGAffineTransform(rotationAngle: radians))
        rotatedRect.origin = .zero

        let context = CGContext(
            data: nil,
            width: Int(rotatedRect.width),
            height: Int(rotatedRect.height),
            bitsPerComponent: self.bitsPerComponent,
            bytesPerRow: 0,
            space: self.colorSpace ?? CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: self.bitmapInfo.rawValue
        )

        guard let ctx = context else { return nil }

        // Move origin to center for rotation
        ctx.translateBy(x: rotatedRect.width / 2, y: rotatedRect.height / 2)
        ctx.rotate(by: radians)

        // Use this to flip image if needed
        // ctx.scaleBy(x: 1.0, y: 1.0)

        let drawRect = CGRect(
            x: -CGFloat(self.width) / 2,
            y: -CGFloat(self.height) / 2,
            width: CGFloat(self.width),
            height: CGFloat(self.height)
        )

        ctx.draw(self, in: drawRect)

        return ctx.makeImage()
    }

    static func rotate(pixelBuffer: CVPixelBuffer, angle: CGFloat) -> CVPixelBuffer? {
        let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
        let transform = CGAffineTransform(rotationAngle: angle)

        let rotatedImage = ciImage.transformed(by: transform)
        let context = CIContext()

        var outputBuffer: CVPixelBuffer?
        let attrs = [
            kCVPixelBufferCGImageCompatibilityKey: true,
            kCVPixelBufferCGBitmapContextCompatibilityKey: true,
        ] as CFDictionary

        let width = Int(rotatedImage.extent.width)
        let height = Int(rotatedImage.extent.height)

        let status = CVPixelBufferCreate(
            nil,
            width,
            height,
            CVPixelBufferGetPixelFormatType(pixelBuffer),
            attrs,
            &outputBuffer
        )

        guard status == kCVReturnSuccess, let rotatedBuffer = outputBuffer else {
            return nil
        }

        context.render(rotatedImage, to: rotatedBuffer)
        return rotatedBuffer
    }
}