import CallsBroadcastKit
import CoreVideo
import Foundation
import os
import LiveKitWebRTC

/// The iPhone's whole screen, arriving from the broadcast extension over
/// its socket (see CallsBroadcastKit), delivered into the call's screen
/// source.
///
/// There is no picker to wait on by the time one of these exists: the
/// person chose to broadcast at the system's own sheet, and the extension
/// connecting is the proof. So `start` returns at once, and `ended` is the
/// extension going away — the broadcast stopped from Control Center or the
/// red status bar. `stop` closes the connection, which is how stopping in
/// the app reaches the system: the extension's next write fails and it
/// ends the broadcast.
@MainActor public final class BroadcastScreenCapture: ScreenCapture {
  private let connection: BroadcastConnection
  private var stopped = false
  private var feed: StillScreenFeed?
  /// Called after the share has ended on its own, once the call has heard.
  /// The app uses it to leave a call it is only still in for the screen.
  public var afterEnded: (@MainActor () -> Void)?

  public init(_ connection: BroadcastConnection) { self.connection = connection }

  public func start(source: RTCVideoSource, ended: @escaping @MainActor () -> Void) async throws {
    let feed = StillScreenFeed(source: source)
    self.feed = feed
    connection.onFrame = { pixels, orientation in feed.frame(pixels, rotation: broadcastRotation(orientation)) }
    connection.onClose = { [weak self] in
      Task { @MainActor [weak self] in
        guard let self, !self.stopped else { return }
        self.stopped = true
        ended()
        self.afterEnded?()
      }
    }
    connection.resume()
  }

  public var counts: ScreenCaptureCounts { feed?.counts ?? ScreenCaptureCounts() }

  public func stop() {
    stopped = true
    feed?.stop()
    feed = nil
    connection.onClose = nil
    connection.onFrame = nil
    connection.close()
  }
}

/// ReplayKit's orientation for a frame, as the rotation WebRTC should
/// apply to show it upright. ReplayKit says which way the *image* is
/// turned, as a CGImagePropertyOrientation: `left` (8) is a phone held in
/// landscape with its top to the left, whose frames need turning a quarter
/// clockwise to stand up; `right` (6) the other way. Anything else — and a
/// phone upright — is sent as it came.
func broadcastRotation(_ orientation: UInt16) -> RTCVideoRotation {
  switch orientation {
  case 3, 4: return ._180
  case 5, 8: return ._90
  case 6, 7: return ._270
  default: return ._0
  }
}

/// The broadcast's frames into the call's source, and the last of them
/// again once a second while the screen is still (2026-09-30). ReplayKit
/// sends nothing at all for a screen that is not changing, and a receiver
/// that sees no frame for four seconds reports the screen stalled (spec
/// §15.21) — which, for a still iPhone screen, was a false alarm the
/// sharer then tried to mend. The Mac's capture has repeated its last
/// frame since 2026-09-28 (MacScreenCapture); this is the same rule.
final class StillScreenFeed: @unchecked Sendable {
  static let repeatAfter: UInt64 = 1_000_000_000

  private let source: RTCVideoSource
  private let capturer: RTCVideoCapturer
  private let queue = DispatchQueue(label: "calls.broadcast.feed", qos: .userInteractive)
  private var timer: DispatchSourceTimer?
  // Only touched on `queue`.
  private var last: (CVPixelBuffer, RTCVideoRotation)?
  private var lastSent: UInt64 = 0
  private let tally = OSAllocatedUnfairLock(initialState: ScreenCaptureCounts(heartbeats: nil, delivered: 0))
  var counts: ScreenCaptureCounts { tally.withLock { $0 } }

  init(source: RTCVideoSource) {
    self.source = source
    capturer = RTCVideoCapturer(delegate: source)
    let timer = DispatchSource.makeTimerSource(queue: queue)
    timer.schedule(deadline: .now() + .milliseconds(500), repeating: .milliseconds(500), leeway: .milliseconds(50))
    timer.setEventHandler { [weak self] in self?.repeatIfIdle() }
    timer.resume()
    self.timer = timer
  }

  /// One frame from the extension. Stamped on arrival, on this device's
  /// clock, rather than with the extension's: the two are separate
  /// processes, and the encoder only needs the frames in order and spaced
  /// as they came.
  func frame(_ pixels: CVPixelBuffer, rotation: RTCVideoRotation) {
    queue.async {
      self.last = (pixels, rotation)
      self.deliver(pixels, rotation)
    }
  }

  func stop() {
    timer?.cancel()
    timer = nil
    queue.async { self.last = nil }
  }

  private func repeatIfIdle() {
    guard let (pixels, rotation) = last else { return }
    let now = DispatchTime.now().uptimeNanoseconds
    guard now >= lastSent, now - lastSent >= Self.repeatAfter else { return }
    deliver(pixels, rotation)
  }

  private func deliver(_ pixels: CVPixelBuffer, _ rotation: RTCVideoRotation) {
    let ns = max(DispatchTime.now().uptimeNanoseconds, lastSent + 1)
    lastSent = ns
    tally.withLock { $0.delivered += 1 }
    source.capturer(
      capturer,
      didCapture: RTCVideoFrame(buffer: RTCCVPixelBuffer(pixelBuffer: pixels), rotation: rotation, timeStampNs: Int64(ns)))
  }
}