import AVFoundation
import CoreMediaIO
import ImageIO
import UniformTypeIdentifiers

func writePNG(bgra pixels: inout [UInt8], width: Int, height: Int, to url: URL) {
    let bitmapInfo = CGBitmapInfo.byteOrder32Little.rawValue | CGImageAlphaInfo.noneSkipFirst.rawValue
    guard let context = CGContext(data: &pixels, width: width, height: height, bitsPerComponent: 8, bytesPerRow: width * 4,
                                  space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: bitmapInfo),
          let image = context.makeImage(),
          let destination = CGImageDestinationCreateWithURL(url as CFURL, UTType.png.identifier as CFString, 1, nil)
    else { return }
    CGImageDestinationAddImage(destination, image, nil)
    CGImageDestinationFinalize(destination)
}

/// Writes frames to an H.264 .mp4 file.
final class Recorder {
    let url: URL
    private let writer: AVAssetWriter
    private let input: AVAssetWriterInput
    private let adaptor: AVAssetWriterInputPixelBufferAdaptor
    private let width: Int, height: Int
    private var startTime: Double?

    init(url: URL, width: Int, height: Int) throws {
        self.url = url
        self.width = width
        self.height = height
        writer = try AVAssetWriter(outputURL: url, fileType: .mp4)
        input = AVAssetWriterInput(mediaType: .video, outputSettings: [
            AVVideoCodecKey: AVVideoCodecType.h264,
            AVVideoWidthKey: width,
            AVVideoHeightKey: height,
        ])
        input.expectsMediaDataInRealTime = true
        adaptor = AVAssetWriterInputPixelBufferAdaptor(assetWriterInput: input, sourcePixelBufferAttributes: [
            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
            kCVPixelBufferWidthKey as String: width,
            kCVPixelBufferHeightKey as String: height,
        ])
        writer.add(input)
        guard writer.startWriting() else { throw writer.error ?? CocoaError(.fileWriteUnknown) }
        writer.startSession(atSourceTime: .zero)
    }

    func append(bgra pixels: [UInt8], at time: Double) {
        guard input.isReadyForMoreMediaData, let pool = adaptor.pixelBufferPool else { return }
        var created: CVPixelBuffer?
        CVPixelBufferPoolCreatePixelBuffer(nil, pool, &created)
        guard let buffer = created else { return }

        CVPixelBufferLockBaseAddress(buffer, [])
        if let base = CVPixelBufferGetBaseAddress(buffer) {
            let rowBytes = CVPixelBufferGetBytesPerRow(buffer)
            pixels.withUnsafeBytes { source in
                for y in 0..<height {
                    memcpy(base + y * rowBytes, source.baseAddress! + y * width * 4, width * 4)
                }
            }
        }
        CVPixelBufferUnlockBaseAddress(buffer, [])

        if startTime == nil { startTime = time }
        adaptor.append(buffer, withPresentationTime: CMTime(seconds: time - startTime!, preferredTimescale: 600))
    }

    func finish() {
        input.markAsFinished()
        let done = DispatchSemaphore(value: 0)
        writer.finishWriting { done.signal() }
        done.wait()
    }
}

struct VirtualCameraError: Error, CustomStringConvertible {
    let description: String
}

/// Feeds frames into OBS's virtual camera, so Meet, Zoom and anything else that takes a webcam can use them.
/// OBS itself must not be running its own virtual camera at the same time.
final class VirtualCamera {
    // fixed identifier of the camera OBS installs (from obs-studio's mac-virtualcam plugin)
    private static let obsDeviceUID = "7626645E-4425-469E-9D8B-97E0FA59AC75"

    private let deviceID: CMIODeviceID
    private let streamID: CMIOStreamID
    private let queue: CMSimpleQueue
    private let format: CMFormatDescription
    private let pool: CVPixelBufferPool

    init(width: Int, height: Int) throws {
        func address(_ selector: Int) -> CMIOObjectPropertyAddress {
            CMIOObjectPropertyAddress(mSelector: CMIOObjectPropertySelector(selector),
                                      mScope: CMIOObjectPropertyScope(kCMIOObjectPropertyScopeGlobal),
                                      mElement: CMIOObjectPropertyElement(kCMIOObjectPropertyElementMain))
        }
        func objects(of object: CMIOObjectID, _ selector: Int) -> [CMIOObjectID] {
            var property = address(selector)
            var size: UInt32 = 0, used: UInt32 = 0
            CMIOObjectGetPropertyDataSize(object, &property, 0, nil, &size)
            var ids = [CMIOObjectID](repeating: 0, count: Int(size) / MemoryLayout<CMIOObjectID>.size)
            CMIOObjectGetPropertyData(object, &property, 0, nil, size, &used, &ids)
            return ids
        }
        func uid(of device: CMIOObjectID) -> String? {
            var property = address(kCMIODevicePropertyDeviceUID)
            var value: Unmanaged<CFString>?
            var used: UInt32 = 0
            CMIOObjectGetPropertyData(device, &property, 0, nil, UInt32(MemoryLayout<Unmanaged<CFString>?>.size), &used, &value)
            return value?.takeRetainedValue() as String?
        }

        let devices = objects(of: CMIOObjectID(kCMIOObjectSystemObject), kCMIOHardwarePropertyDevices)
        guard let device = devices.first(where: { uid(of: $0) == VirtualCamera.obsDeviceUID }) else {
            throw VirtualCameraError(description: """
                OBS Virtual Camera not found. One-time setup: enable OBS under System Settings > General > \
                Login Items & Extensions > Camera Extensions, then in OBS click Start Virtual Camera, Stop Virtual Camera, and quit OBS.
                """)
        }
        // the camera has two streams: [0] is what apps watch, [1] is where frames are fed in
        let streams = objects(of: device, kCMIODevicePropertyStreams)
        guard streams.count >= 2 else { throw VirtualCameraError(description: "OBS Virtual Camera has no input stream.") }
        deviceID = device
        streamID = streams[1]

        var createdQueue: Unmanaged<CMSimpleQueue>?
        CMIOStreamCopyBufferQueue(streamID, { _, _, _ in }, nil, &createdQueue)
        var createdFormat: CMFormatDescription?
        CMVideoFormatDescriptionCreate(allocator: nil, codecType: kCVPixelFormatType_422YpCbCr8,
                                       width: Int32(width), height: Int32(height), extensions: nil,
                                       formatDescriptionOut: &createdFormat)
        var createdPool: CVPixelBufferPool?
        CVPixelBufferPoolCreate(nil, nil, [
            kCVPixelBufferPixelFormatTypeKey: kCVPixelFormatType_422YpCbCr8,
            kCVPixelBufferWidthKey: width,
            kCVPixelBufferHeightKey: height,
            kCVPixelBufferIOSurfacePropertiesKey: [:] as CFDictionary,
        ] as CFDictionary, &createdPool)
        guard let createdQueue, let createdFormat, let createdPool else {
            throw VirtualCameraError(description: "Could not set up the OBS Virtual Camera stream.")
        }
        queue = createdQueue.takeRetainedValue()
        format = createdFormat
        pool = createdPool
        guard CMIODeviceStartStream(deviceID, streamID) == noErr else {
            throw VirtualCameraError(description: "Could not start the OBS Virtual Camera stream (is OBS using it?).")
        }
    }

    deinit {
        CMIODeviceStopStream(deviceID, streamID)
    }

    /// Sends one frame. The picture is flipped back left-to-right: the preview is a mirror,
    /// but the people on the call should see the real thing.
    func send(bgra pixels: [UInt8], width: Int, height: Int) {
        var created: CVPixelBuffer?
        CVPixelBufferPoolCreatePixelBuffer(nil, pool, &created)
        guard let buffer = created else { return }

        CVPixelBufferLockBaseAddress(buffer, [])
        if let base = CVPixelBufferGetBaseAddress(buffer)?.assumingMemoryBound(to: UInt8.self) {
            let rowBytes = CVPixelBufferGetBytesPerRow(buffer)
            pixels.withUnsafeBufferPointer { source in
                // BGRA -> UYVY (BT.601 video range): two pixels share one U and one V
                for y in 0..<height {
                    let row = base + y * rowBytes
                    for pair in 0..<(width / 2) {
                        let first = (y * width + (width - 1 - pair * 2)) * 4
                        let second = first - 4
                        let b0 = Int(source[first]), g0 = Int(source[first + 1]), r0 = Int(source[first + 2])
                        let b1 = Int(source[second]), g1 = Int(source[second + 1]), r1 = Int(source[second + 2])
                        row[pair * 4] = UInt8(clamping: ((-38 * r0 - 74 * g0 + 112 * b0 + 128) >> 8) + 128)
                        row[pair * 4 + 1] = UInt8(clamping: ((66 * r0 + 129 * g0 + 25 * b0 + 128) >> 8) + 16)
                        row[pair * 4 + 2] = UInt8(clamping: ((112 * r0 - 94 * g0 - 18 * b0 + 128) >> 8) + 128)
                        row[pair * 4 + 3] = UInt8(clamping: ((66 * r1 + 129 * g1 + 25 * b1 + 128) >> 8) + 16)
                    }
                }
            }
        }
        CVPixelBufferUnlockBaseAddress(buffer, [])

        var timing = CMSampleTimingInfo(
            duration: .invalid,
            presentationTimeStamp: CMTime(value: CMTimeValue(clock_gettime_nsec_np(CLOCK_UPTIME_RAW)), timescale: 1_000_000_000),
            decodeTimeStamp: .invalid)
        var sample: CMSampleBuffer?
        CMSampleBufferCreateForImageBuffer(allocator: nil, imageBuffer: buffer, dataReady: true, makeDataReadyCallback: nil,
                                           refcon: nil, formatDescription: format, sampleTiming: &timing, sampleBufferOut: &sample)
        if let sample {
            // the queue takes ownership of the sample
            if CMSimpleQueueEnqueue(queue, element: Unmanaged.passRetained(sample).toOpaque()) != noErr {
                Unmanaged.passUnretained(sample).release()
            }
        }
    }
}