//
//  RecordingSession.swift
//  VisionCamera
//
//  Created by Marc Rousavy on 01.05.21.
//  Copyright © 2021 mrousavy. All rights reserved.
//

import AVFoundation
import Foundation

// MARK: - BufferType

enum BufferType {
  case audio
  case video
}

// MARK: - RecordingSession

class RecordingSession {
  private let assetWriter: AVAssetWriter
  private var audioWriter: AVAssetWriterInput?
  private var videoWriter: AVAssetWriterInput?
  private let completionHandler: (RecordingSession, AVAssetWriter.Status, Error?) -> Void

  private var initialTimestamp: CMTime?
  private var latestTimestamp: CMTime?
  private var hasStartedWritingSession = false
  private var hasWrittenFirstVideoFrame = false
  private var isFinishing = false

  /**
   Gets the file URL of the recorded video.
   */
  var url: URL {
    return assetWriter.outputURL
  }

  /**
   Get the duration (in seconds) of the recorded video.
   */
  var duration: Double {
    guard let latestTimestamp = latestTimestamp,
          let initialTimestamp = initialTimestamp else {
      return 0.0
    }
    return (latestTimestamp - initialTimestamp).seconds
  }

  init(url: URL,
       fileType: AVFileType,
       completion: @escaping (RecordingSession, AVAssetWriter.Status, Error?) -> Void) throws {
    completionHandler = completion

    do {
      assetWriter = try AVAssetWriter(outputURL: url, fileType: fileType)
      assetWriter.shouldOptimizeForNetworkUse = false
    } catch let error as NSError {
      throw CameraError.capture(.createRecorderError(message: error.description))
    }
  }

  deinit {
    if assetWriter.status == .writing {
      ReactLogger.log(level: .info, message: "Cancelling AssetWriter...")
      assetWriter.cancelWriting()
    }
  }

  /**
   Initializes an AssetWriter for video frames (CMSampleBuffers).
   */
  func initializeVideoWriter(withSettings settings: [String: Any]) {
    guard !settings.isEmpty else {
      ReactLogger.log(level: .error, message: "Tried to initialize Video Writer with empty settings!")
      return
    }
    guard videoWriter == nil else {
      ReactLogger.log(level: .error, message: "Tried to add Video Writer twice!")
      return
    }

    ReactLogger.log(level: .info, message: "Initializing Video AssetWriter with settings: \(settings.description)")
    videoWriter = AVAssetWriterInput(mediaType: .video, outputSettings: settings)
    videoWriter!.expectsMediaDataInRealTime = true
    assetWriter.add(videoWriter!)
    ReactLogger.log(level: .info, message: "Initialized Video AssetWriter.")
  }

  /**
   Initializes an AssetWriter for audio frames (CMSampleBuffers).
   */
  func initializeAudioWriter(withSettings settings: [String: Any]?, format: CMFormatDescription) {
    guard audioWriter == nil else {
      ReactLogger.log(level: .error, message: "Tried to add Audio Writer twice!")
      return
    }

    if let settings = settings {
      ReactLogger.log(level: .info, message: "Initializing Audio AssetWriter with settings: \(settings.description)")
    } else {
      ReactLogger.log(level: .info, message: "Initializing Audio AssetWriter default settings...")
    }
    audioWriter = AVAssetWriterInput(mediaType: .audio, outputSettings: settings, sourceFormatHint: format)
    audioWriter!.expectsMediaDataInRealTime = true
    assetWriter.add(audioWriter!)
    ReactLogger.log(level: .info, message: "Initialized Audio AssetWriter.")
  }

  /**
   Start the Asset Writer(s). If the AssetWriter failed to start, an error will be thrown.
   */
  func startAssetWriter() throws {
    ReactLogger.log(level: .info, message: "Starting Asset Writer(s)...")

    let success = assetWriter.startWriting()
    if success {
      ReactLogger.log(level: .info, message: "Asset Writer(s) started!")
    } else {
      ReactLogger.log(level: .error, message: "Failed to start Asset Writer(s)!")
      throw CameraError.capture(.createRecorderError(message: "Failed to start Asset Writer(s)!"))
    }
  }

  /**
   Appends a new CMSampleBuffer to the Asset Writer. Use bufferType to specify if this is a video or audio frame.
   */
  func appendBuffer(_ buffer: CMSampleBuffer, type bufferType: BufferType) {
    guard assetWriter.status == .writing else {
      ReactLogger.log(level: .error, message: "Frame arrived, but AssetWriter status is \(assetWriter.status.descriptor)!")
      return
    }
    if !CMSampleBufferDataIsReady(buffer) {
      ReactLogger.log(level: .error, message: "Frame arrived, but sample buffer is not ready!")
      return
    }

    switch bufferType {
    case .video:
      guard let videoWriter = videoWriter else {
        ReactLogger.log(level: .error, message: "Video Frame arrived but VideoWriter was nil!")
        return
      }
      if !videoWriter.isReadyForMoreMediaData {
        ReactLogger.log(level: .warning,
                        message: "The Video AVAssetWriterInput was not ready for more data! Is your frame rate too high?")
        return
      }
      let timestamp = CMSampleBufferGetPresentationTimeStamp(buffer)
      // Start the writing session before we write the first video frame
      if !hasStartedWritingSession {
        initialTimestamp = timestamp
        assetWriter.startSession(atSourceTime: timestamp)
        ReactLogger.log(level: .info, message: "Started RecordingSession at \(timestamp.seconds) seconds.")
        hasStartedWritingSession = true
      }
      // Write Video Buffer!
      videoWriter.append(buffer)
      // Update state
      latestTimestamp = timestamp
      if !hasWrittenFirstVideoFrame {
        hasWrittenFirstVideoFrame = true
      }
    case .audio:
      guard let audioWriter = audioWriter else {
        ReactLogger.log(level: .error, message: "Audio Frame arrived but AudioWriter was nil!")
        return
      }
      if !audioWriter.isReadyForMoreMediaData {
        return
      }
      if !hasWrittenFirstVideoFrame || !hasStartedWritingSession {
        // first video frame has not been written yet, so skip this audio frame.
        return
      }
      // Write Audio Sample!
      audioWriter.append(buffer)
    }

    // If we failed to write the frames, stop the Recording
    if assetWriter.status == .failed {
      ReactLogger.log(level: .error,
                      message: "AssetWriter failed to write buffer! Error: \(assetWriter.error?.localizedDescription ?? "none")")
      finish()
    }
  }

  /**
   Marks the AssetWriters as finished and stops writing frames. The callback will be invoked either with an error or the status "success".
   */
  func finish() {
    ReactLogger.log(level: .info, message: "Finishing Recording with AssetWriter status \"\(assetWriter.status.descriptor)\"...")

    if isFinishing {
      ReactLogger.log(level: .warning, message: "Tried calling finish() twice while AssetWriter is still writing!")
      return
    }

    if !hasWrittenFirstVideoFrame {
      let error = NSError(domain: "capture/aborted",
                          code: 1,
                          userInfo: [NSLocalizedDescriptionKey: "Stopped Recording Session too early, no frames have been recorded!"])
      completionHandler(self, .failed, error)
    } else if assetWriter.status == .writing {
      isFinishing = true
      videoWriter?.markAsFinished()
      audioWriter?.markAsFinished()
      assetWriter.finishWriting {
        self.isFinishing = false
        self.completionHandler(self, self.assetWriter.status, self.assetWriter.error)
      }
    } else {
      completionHandler(self, assetWriter.status, assetWriter.error)
    }
  }
}