//
//  BBMetalVideoSource.swift
//  BBMetalImage
//
//  Created by Kaibo Lu on 4/24/19.
//  Copyright © 2019 Kaibo Lu. All rights reserved.
//

import AVFoundation

public typealias BBMetalVideoSourceProgress = (CMTime) -> Void
public typealias BBMetalVideoSourceCompletion = (Bool) -> Void

/// Video source reading video frame and providing Metal texture
public class BBMetalVideoSource {
    /// Image consumers
    public var consumers: [BBMetalImageConsumer] { return transformFilter.consumers }
    
    private var transformFilter: BBMetalTransformFilter!
    
    private let url: URL
    private let lock: DispatchSemaphore
    private let playPauseLock: DispatchSemaphore
    
    private var asset: AVAsset!
    private var assetReader: AVAssetReader!
    private var videoOutput: AVAssetReaderTrackOutput!
    
    private var audioOutput: AVAssetReaderTrackOutput!
    private var lastAudioBuffer: CMSampleBuffer?
    
    private let queue = DispatchQueue(label: "BBMetalImageVideoSource.queue")
    
    /// A block to call before processing each video sample buffer
    public var preprocessVideo: ((CMSampleBuffer) -> Void)? {
        get {
            lock.wait()
            let p = _preprocessVideo
            lock.signal()
            return p
        }
        set {
            lock.wait()
            _preprocessVideo = newValue
            lock.signal()
        }
    }
    private var _preprocessVideo: ((CMSampleBuffer) -> Void)?
    
    /// Audio consumer processing audio sample buffer.
    /// Set this property to nil (default value) if not processing audio.
    /// Set this property to a given audio consumer if processing audio.
    public var audioConsumer: BBMetalAudioConsumer? {
        get {
            lock.wait()
            let a = _audioConsumer
            lock.signal()
            return a
        }
        set {
            lock.wait()
            _audioConsumer = newValue
            lock.signal()
        }
    }
    private var _audioConsumer: BBMetalAudioConsumer?
    
    /// Whether to process video with the actual rate. False by default, meaning the processing speed is faster than the actual video rate.
    public var playWithVideoRate: Bool {
        get {
            lock.wait()
            let p = _playWithVideoRate
            lock.signal()
            return p
        }
        set {
            lock.wait()
            _playWithVideoRate = newValue
            lock.signal()
        }
    }
    private var _playWithVideoRate: Bool
    
    private var lastSampleFrameTime: CMTime!
    private var lastActualPlayTime: Double!
    
    /// Whether to run benchmark or not.
    /// Running benchmark records frame duration.
    /// False by default.
    public var benchmark: Bool {
        get {
            lock.wait()
            let b = _benchmark
            lock.signal()
            return b
        }
        set {
            lock.wait()
            _benchmark = newValue
            lock.signal()
        }
    }
    private var _benchmark: Bool
    
    /// Average frame duration, or 0 if not valid value.
    /// To get valid value, set `benchmark` to true.
    public var averageFrameDuration: Double {
        lock.wait()
        let d = processedFrameCount > 0 ? totalProcessFrameTime / Double(processedFrameCount) : 0
        lock.signal()
        return d
    }
    
    private var processedFrameCount: Int
    private var totalProcessFrameTime: Double
    
    private var textureCache: CVMetalTextureCache!
    
    public init?(url: URL) {
        self.url = url
        lock = DispatchSemaphore(value: 1)
        playPauseLock = DispatchSemaphore(value: 0)
        _playWithVideoRate = false
        _benchmark = false
        processedFrameCount = 0
        totalProcessFrameTime = 0
        
        #if !targetEnvironment(simulator)
        if CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, BBMetalDevice.sharedDevice, nil, &textureCache) != kCVReturnSuccess ||
            textureCache == nil {
            return nil
        }
        #endif
        
        transformFilter = BBMetalTransformFilter()
        transformFilter.add(source: self)
    }
    
    /// Starts reading and processing video frame
    ///
    /// - Parameter completion: a closure to call after processing; The parameter of closure is true if succeed processing all video frames, or false if fail to processing all the video frames (due to user cancel or error)
    public func start(progress: BBMetalVideoSourceProgress? = nil, completion: BBMetalVideoSourceCompletion? = nil) {
        lock.wait()
        let isReading = (assetReader != nil)
        lock.signal()
        if isReading {
            print("Should not call \(#function) while asset reader is reading")
            return
        }
        let asset = AVAsset(url: url)
        if let track = asset.tracks(withMediaType: .video).first {
            transformFilter.transform = track.preferredTransform
            queue.async { [weak self] in
                guard let self = self else { return }
                self.lock.wait()
                self.asset = asset
                if self.prepareAssetReader() {
                    self.lock.signal()
                    self.processAsset(progress: progress, completion: completion)
                } else {
                    self.reset()
                    self.lock.signal()
                }
            }
        } else {
            self.safeReset()
        }
    }
    
    private var paused = false
    
    public func pause() {
        guard paused == false else { return }
        paused = true
        DispatchQueue.global().async { [weak self] in
            self?.playPauseLock.wait()
        }
    }
    
    public func resume() {
        guard paused == true else { return }
        paused = false
        DispatchQueue.global().async { [weak self] in
            self?.playPauseLock.signal()
        }
    }
    
    /// Cancels reading and processing video frame
    public func cancel() {
        resume()
        lock.wait()
        if let reader = assetReader,
            reader.status == .reading {
            reader.cancelReading()
            reset()
        }
        lock.signal()
    }
    
    /// Resets benchmark record data
    public func resetBenchmark() {
        lock.wait()
        processedFrameCount = 0
        totalProcessFrameTime = 0
        lock.signal()
    }
    
    private func safeReset() {
        lock.wait()
        reset()
        lock.signal()
    }
    
    private func reset() {
        asset = nil
        assetReader = nil
        videoOutput = nil
        audioOutput = nil
        lastAudioBuffer = nil
    }
    
    private func prepareAssetReader() -> Bool {
        guard let reader = try? AVAssetReader(asset: asset),
            let videoTrack = asset.tracks(withMediaType: .video).first else { return false }
        assetReader = reader
        videoOutput = AVAssetReaderTrackOutput(track: videoTrack, outputSettings: [kCVPixelBufferPixelFormatTypeKey as String : kCVPixelFormatType_32BGRA])
        videoOutput.alwaysCopiesSampleData = false
        if !assetReader.canAdd(videoOutput) { return false }
        assetReader.add(videoOutput)
        
        if _audioConsumer != nil,
            let audioTrack = asset.tracks(withMediaType: .audio).first {
            audioOutput = AVAssetReaderTrackOutput(track: audioTrack, outputSettings: [AVFormatIDKey : kAudioFormatLinearPCM])
            audioOutput.alwaysCopiesSampleData = false
            if !assetReader.canAdd(audioOutput) { return false }
            assetReader.add(audioOutput)
        }
        return true
    }
    
    private var frameDelayCorrectionFactor: Double = 1
    
    private func processAsset(progress: BBMetalVideoSourceProgress?, completion: BBMetalVideoSourceCompletion?) {
        lock.wait()
        guard let reader = assetReader,
            reader.status == .unknown,
            reader.startReading() else {
            reset()
            lock.signal()
            return
        }
        
        // Read and process video buffer
        let useVideoRate = _playWithVideoRate
        var startTime: Double = _benchmark ? CACurrentMediaTime() : 0
        var sleepTime: Double = 0
        while let reader = assetReader,
              reader.status == .reading,
              let sampleBuffer = videoOutput.copyNextSampleBuffer() {
            playPauseLock.signal()
            
            if let preprocessVideo = _preprocessVideo {
                lock.signal()
                preprocessVideo(sampleBuffer)
                lock.wait()
            }
            
            guard let texture = texture(with: sampleBuffer) else { break }
            
            let sampleFrameTime = CMSampleBufferGetOutputPresentationTimeStamp(sampleBuffer)
            
            if useVideoRate {
                //sync video with audio if possible
                if let playerTime = _audioConsumer?.audioPlayerCurrentTime(), playerTime != .zero {
                    let diff = CMTimeGetSeconds(sampleFrameTime) - CMTimeGetSeconds(playerTime)
                    if diff > 0.0 {
                        sleepTime = diff
                        if sleepTime > 1.0 {
                            sleepTime = 0.0
                        }
                        usleep(UInt32(1000000 * sleepTime))
                    } else {
                        sleepTime = 0
                    }
                } else {
                    // calculate and correct sleep time between frames
                    if let lastFrameTime = lastSampleFrameTime,
                       let lastPlayTime = lastActualPlayTime {
                        let detalFrameTime = CMTimeGetSeconds(CMTimeSubtract(sampleFrameTime, lastFrameTime))
                        let detalPlayTime = CACurrentMediaTime() - lastPlayTime
                        if detalFrameTime > detalPlayTime {
                            // calculate delay before next frame
                            sleepTime = (detalFrameTime - detalPlayTime) * frameDelayCorrectionFactor
                            // save falling asleep time
                            let fallingAsleepTimestamp = CACurrentMediaTime()
                            // wait for next frame
                            usleep(UInt32(1000000 * sleepTime))
                            // calculate real sleep time
                            let awakeTimestamp = CACurrentMediaTime()
                            let realSleepTimeinterval = awakeTimestamp - fallingAsleepTimestamp
                            // correct sleep time for next frame
                            if realSleepTimeinterval > 0 {
                                frameDelayCorrectionFactor = max(sleepTime / realSleepTimeinterval, 0.7)
                            } else {
                                frameDelayCorrectionFactor = 1
                            }

                        } else {
                            sleepTime = 0
                        }
                    }
                    lastSampleFrameTime = sampleFrameTime
                    lastActualPlayTime = CACurrentMediaTime()
                }
            }
            
            // Read and process audio buffer
            // Let video buffer go faster than audio buffer
            // Make sure audio and video buffer have similar output presentation timestamp
            var currentAudioBuffer: CMSampleBuffer?
            let currentAudioConsumer = _audioConsumer
            if currentAudioConsumer != nil {
                if let last = lastAudioBuffer,
                   CMTimeCompare(CMSampleBufferGetOutputPresentationTimeStamp(last), sampleFrameTime) <= 0 {
                    // Process audio buffer
                    currentAudioBuffer = last
                    lastAudioBuffer = nil
                    
                } else if lastAudioBuffer == nil,
                          audioOutput != nil,
                          let audioBuffer = audioOutput.copyNextSampleBuffer() {
                    if CMTimeCompare(CMSampleBufferGetOutputPresentationTimeStamp(audioBuffer), sampleFrameTime) <= 0 {
                        // Process audio buffer
                        currentAudioBuffer = audioBuffer
                    } else {
                        // Audio buffer goes faster than video
                        // Process audio buffer later
                        lastAudioBuffer = audioBuffer
                    }
                }
            }
            lock.signal()
            
            autoreleasepool {
                // Transmit video texture
                let output = BBMetalDefaultTexture(metalTexture: texture.metalTexture,
                                                   sampleTime: sampleFrameTime,
                                                   cvMetalTexture: texture.cvMetalTexture)
                transformFilter.newTextureAvailable(output, from: self)
            }
            
            progress?(sampleFrameTime)
            
            // Transmit audio buffer
            if let audioBuffer = currentAudioBuffer { currentAudioConsumer?.newAudioSampleBufferAvailable(audioBuffer) }
            
            // Benchmark
            if startTime != 0 {
                let now = CACurrentMediaTime()
                processedFrameCount += 1
                totalProcessFrameTime += now - startTime - sleepTime
                startTime = now
            }
            
            lock.wait()
            playPauseLock.wait()
        }
        // Read and process the rest audio buffers
        if let consumer = _audioConsumer,
           let audioBuffer = lastAudioBuffer {
            lock.signal()
            consumer.newAudioSampleBufferAvailable(audioBuffer)
            lock.wait()
        }
        while let consumer = _audioConsumer,
              let reader = assetReader,
              reader.status == .reading,
              audioOutput != nil,
              let audioBuffer = audioOutput.copyNextSampleBuffer() {
            lock.signal()
            consumer.newAudioSampleBufferAvailable(audioBuffer)
            lock.wait()
        }
        var finish = false
        if assetReader != nil {
            finish = true
            reset()
        }
        lock.signal()
        
        completion?(finish)
    }
    
    private func texture(with sampleBuffer: CMSampleBuffer) -> BBMetalVideoTextureItem? {
        guard let imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else { return nil }
        let width = CVPixelBufferGetWidth(imageBuffer)
        let height = CVPixelBufferGetHeight(imageBuffer)
        
        #if !targetEnvironment(simulator)
        var cvMetalTextureOut: CVMetalTexture?
        let result = CVMetalTextureCacheCreateTextureFromImage(kCFAllocatorDefault,
                                                               textureCache,
                                                               imageBuffer,
                                                               nil,
                                                               .bgra8Unorm,
                                                               width,
                                                               height,
                                                               0,
                                                               &cvMetalTextureOut)
        if result == kCVReturnSuccess,
            let cvMetalTexture = cvMetalTextureOut,
            let texture = CVMetalTextureGetTexture(cvMetalTexture) {
            return BBMetalVideoTextureItem(metalTexture: texture, cvMetalTexture: cvMetalTexture)
        }
        #endif
        return nil
    }
}

extension BBMetalVideoSource: BBMetalImageSource {
    @discardableResult
    public func add<T: BBMetalImageConsumer>(consumer: T) -> T {
        return transformFilter.add(consumer: consumer)
    }
    
    public func add(consumer: BBMetalImageConsumer, at index: Int) {
        transformFilter.add(consumer: consumer, at: index)
    }
    
    public func remove(consumer: BBMetalImageConsumer) {
        transformFilter.remove(consumer: consumer)
    }
    
    public func removeAllConsumers() {
        transformFilter.removeAllConsumers()
    }
}
