import Accelerate
import CoreVideo
import Foundation

/// The extension's half: connects to the app, and sends it each frame it
/// is given, smaller.
///
/// Everything here is shaped by the extension's memory ceiling. One frame
/// buffer, allocated once and reused; the frame is scaled straight into
/// it from ReplayKit's own pixel buffer, which is released as soon as
/// `send` returns; and at most one frame is ever in flight. A frame that
/// arrives while the last is still being written, or sooner than the rate
/// allows, is dropped rather than queued — a queue is exactly how an
/// extension runs out of memory while the app is busy.
public final class BroadcastSender: @unchecked Sendable {
  private let path: String
  private let longest: CGFloat
  private let interval: UInt64
  private let queue = DispatchQueue(label: "calls.broadcast.send")
  private let lock = NSLock()
  // Guarded by lock.
  private var fd: Int32 = -1
  private var busy = false
  private var last: UInt64 = 0
  private var finished = false
  // Only touched on queue, or before connect.
  private var buffer: UnsafeMutableRawPointer?
  private var capacity = 0

  /// Called once, on the sending queue, when the app has gone: it closed
  /// the socket because the person stopped sharing there, or the call
  /// ended. Not called for `close()`.
  public var onClosed: (@Sendable () -> Void)?

  /// `longest` is the long side frames are sent at, at most; 1280 keeps a
  /// phone's text legible on a laptop and the frame near a megabyte.
  public init(path: String, longest: Int = 1280, framesPerSecond: Int = 15) {
    self.path = path
    self.longest = CGFloat(longest)
    interval = 1_000_000_000 / UInt64(max(1, framesPerSecond))
  }

  deinit { buffer?.deallocate() }

  /// Connects to the app, trying for up to `timeout`: the app may still be
  /// settling into the call as the broadcast's countdown ends.
  public func connect(timeout: TimeInterval) -> Bool {
    let until = Date().addingTimeInterval(timeout)
    repeat {
      if let fd = try? BroadcastSocket.connect(path: path) {
        lock.lock()
        self.fd = fd
        lock.unlock()
        return true
      }
      Thread.sleep(forTimeInterval: 0.1)
    } while Date() < until
    return false
  }

  /// Offers one frame. Returns at once; the write happens on the sending
  /// queue. `pixels` must be 4:2:0 bi-planar, which is what ReplayKit
  /// delivers; anything else is dropped.
  public func send(_ pixels: CVPixelBuffer, orientation: UInt16, timestampNs: UInt64) {
    let format = CVPixelBufferGetPixelFormatType(pixels)
    guard format == kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
      || format == kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
    else { return }
    lock.lock()
    let now = DispatchTime.now().uptimeNanoseconds
    guard fd >= 0, !busy, !finished, now >= last + interval else {
      lock.unlock()
      return
    }
    busy = true
    last = now
    lock.unlock()
    // Scaled here, on ReplayKit's thread, so the pixel buffer is not held
    // past this call; only the copy crosses to the sending queue.
    guard let length = pack(pixels, orientation: orientation, timestampNs: timestampNs) else {
      lock.lock()
      busy = false
      lock.unlock()
      return
    }
    queue.async { [self] in
      let ok = writeAll(length)
      lock.lock()
      busy = false
      let first = !ok && !finished
      if !ok { finished = true }
      lock.unlock()
      if first { onClosed?() }
    }
  }

  public func close() {
    lock.lock()
    finished = true
    let fd = self.fd
    self.fd = -1
    lock.unlock()
    if fd >= 0 { Darwin.close(fd) }
  }

  /// Scales `pixels` into the frame buffer behind a header, and answers the
  /// number of bytes to write.
  private func pack(_ pixels: CVPixelBuffer, orientation: UInt16, timestampNs: UInt64) -> Int? {
    let size = fitted(
      width: CGFloat(CVPixelBufferGetWidth(pixels)), height: CGFloat(CVPixelBufferGetHeight(pixels)),
      longest: longest)
    let header = BroadcastFrameHeader(
      width: size.width, height: size.height, orientation: orientation, timestampNs: timestampNs,
      videoRange: CVPixelBufferGetPixelFormatType(pixels)
        == kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange)
    let length = BroadcastFrameHeader.size + header.payloadLength
    if capacity < length {
      buffer?.deallocate()
      buffer = UnsafeMutableRawPointer.allocate(byteCount: length, alignment: 16)
      capacity = length
    }
    guard let buffer else { return nil }
    header.write(to: buffer)
    CVPixelBufferLockBaseAddress(pixels, .readOnly)
    defer { CVPixelBufferUnlockBaseAddress(pixels, .readOnly) }
    guard let luma = CVPixelBufferGetBaseAddressOfPlane(pixels, 0),
      let chroma = CVPixelBufferGetBaseAddressOfPlane(pixels, 1)
    else { return nil }
    var sourceLuma = vImage_Buffer(
      data: luma, height: vImagePixelCount(CVPixelBufferGetHeightOfPlane(pixels, 0)),
      width: vImagePixelCount(CVPixelBufferGetWidthOfPlane(pixels, 0)),
      rowBytes: CVPixelBufferGetBytesPerRowOfPlane(pixels, 0))
    var sourceChroma = vImage_Buffer(
      data: chroma, height: vImagePixelCount(CVPixelBufferGetHeightOfPlane(pixels, 1)),
      width: vImagePixelCount(CVPixelBufferGetWidthOfPlane(pixels, 1)),
      rowBytes: CVPixelBufferGetBytesPerRowOfPlane(pixels, 1))
    let picture = buffer.advanced(by: BroadcastFrameHeader.size)
    var targetLuma = vImage_Buffer(
      data: picture, height: vImagePixelCount(size.height), width: vImagePixelCount(size.width),
      rowBytes: size.width)
    var targetChroma = vImage_Buffer(
      data: picture.advanced(by: size.width * size.height),
      height: vImagePixelCount(size.height / 2), width: vImagePixelCount(size.width / 2),
      rowBytes: size.width)
    let flags = vImage_Flags(kvImageNoFlags)
    guard vImageScale_Planar8(&sourceLuma, &targetLuma, nil, flags) == kvImageNoError,
      vImageScale_CbCr8(&sourceChroma, &targetChroma, nil, flags) == kvImageNoError
    else { return nil }
    return length
  }

  private func writeAll(_ length: Int) -> Bool {
    lock.lock()
    let fd = self.fd
    lock.unlock()
    guard fd >= 0, let buffer else { return false }
    var sent = 0
    while sent < length {
      let n = Darwin.write(fd, buffer.advanced(by: sent), length - sent)
      if n > 0 {
        sent += n
      } else if n < 0 && errno == EINTR {
        continue
      } else {
        return false
      }
    }
    return true
  }
}