import CoreGraphics
import Foundation

// ---------------------------------------------------------------
// The iPhone's screen, from the broadcast extension to the call
// ---------------------------------------------------------------
//
// Sharing the whole iPhone screen is a ReplayKit broadcast, and ReplayKit
// hands the frames to an extension — a separate process, capped at about
// fifty megabytes, that cannot hold the call's peer connections and could
// not afford WebRTC if it did. So the extension does the least it can:
// shrinks each frame and passes it to the app, which is already in the
// call, over a Unix socket in the App Group container both of them can
// reach. The app listens for as long as a call is up; the extension
// connecting *is* the broadcast starting, and either side closing is it
// ending.
//
// This module is everything the two sides share, and nothing else: no
// WebRTC, no UIKit, no ReplayKit. The extension links this and only this.
//
// One frame on the wire is a 32-byte header and then the picture, NV12
// packed tight: the luma plane, `width` bytes a row, then the interleaved
// chroma plane at half height, also `width` bytes a row. Little-endian.
//
//   0  magic        u32  "CSF1"
//   4  version      u16  1
//   6  orientation  u16  CGImagePropertyOrientation, as ReplayKit says it
//   8  width        u32  even, 2...4096
//   12 height       u32  even, 2...4096
//   16 timestamp    u64  nanoseconds, the extension's clock
//   24 length       u32  width * height * 3 / 2
//   28 flags        u32  bit 0: video range (16-235) rather than full

public struct BroadcastFrameHeader: Equatable, Sendable {
  public static let size = 32
  static let magic: UInt32 = 0x3146_5343  // "CSF1" read little-endian
  static let version: UInt16 = 1
  static let largest = 4096

  public var width: Int
  public var height: Int
  public var orientation: UInt16
  public var timestampNs: UInt64
  /// Whether the luma runs 16-235 rather than 0-255. ReplayKit has sent
  /// both over the years, and a picture decoded with the wrong one is
  /// washed out or crushed.
  public var videoRange: Bool

  public init(
    width: Int, height: Int, orientation: UInt16, timestampNs: UInt64, videoRange: Bool = false
  ) {
    self.width = width
    self.height = height
    self.orientation = orientation
    self.timestampNs = timestampNs
    self.videoRange = videoRange
  }

  public var payloadLength: Int { width * height * 3 / 2 }

  public func write(to bytes: UnsafeMutableRawPointer) {
    bytes.storeBytes(of: Self.magic.littleEndian, toByteOffset: 0, as: UInt32.self)
    bytes.storeBytes(of: Self.version.littleEndian, toByteOffset: 4, as: UInt16.self)
    bytes.storeBytes(of: orientation.littleEndian, toByteOffset: 6, as: UInt16.self)
    bytes.storeBytes(of: UInt32(width).littleEndian, toByteOffset: 8, as: UInt32.self)
    bytes.storeBytes(of: UInt32(height).littleEndian, toByteOffset: 12, as: UInt32.self)
    bytes.storeBytes(of: timestampNs.littleEndian, toByteOffset: 16, as: UInt64.self)
    bytes.storeBytes(of: UInt32(payloadLength).littleEndian, toByteOffset: 24, as: UInt32.self)
    bytes.storeBytes(of: UInt32(videoRange ? 1 : 0).littleEndian, toByteOffset: 28, as: UInt32.self)
  }

  /// Reads a header, or nil for anything that is not one this side can
  /// use. The socket is in a container only this app's own processes can
  /// open, so a bad header is a bug or a version skew rather than an
  /// attack — but it is still refused rather than trusted, because the
  /// length in it is what the reader allocates.
  public static func read(_ bytes: UnsafeRawPointer) -> BroadcastFrameHeader? {
    func u32(_ at: Int) -> UInt32 { UInt32(littleEndian: bytes.loadUnaligned(fromByteOffset: at, as: UInt32.self)) }
    func u16(_ at: Int) -> UInt16 { UInt16(littleEndian: bytes.loadUnaligned(fromByteOffset: at, as: UInt16.self)) }
    guard u32(0) == magic, u16(4) == version else { return nil }
    let width = Int(u32(8))
    let height = Int(u32(12))
    guard (2...largest).contains(width), (2...largest).contains(height),
      width % 2 == 0, height % 2 == 0
    else { return nil }
    let header = BroadcastFrameHeader(
      width: width, height: height, orientation: u16(6),
      timestampNs: UInt64(littleEndian: bytes.loadUnaligned(fromByteOffset: 16, as: UInt64.self)),
      videoRange: u32(28) & 1 == 1)
    guard Int(u32(24)) == header.payloadLength else { return nil }
    return header
  }
}

/// Where the socket lives: one well-known name in the App Group container,
/// short because a Unix socket's path has to fit in 104 bytes, and the
/// container's own path on a device already takes about ninety.
public enum BroadcastPath {
  public static let name = "calls.sock"

  /// The App Group this build shares with its extension, from the
  /// `CallsAppGroup` key both Info.plists carry. One build setting feeds
  /// the key and the entitlement, so a build with a different bundle
  /// identifier changes the group in one place.
  public static var appGroup: String? {
    Bundle.main.object(forInfoDictionaryKey: "CallsAppGroup") as? String
  }

  public static func socket(appGroup: String) -> String? {
    FileManager.default.containerURL(forSecurityApplicationGroupIdentifier: appGroup)?
      .appendingPathComponent(name).path
  }
}

enum BroadcastSocket {
  struct Failure: Error, CustomStringConvertible {
    let description: String
  }

  /// A Unix-socket address for `path`, or nil if the path does not fit.
  static func address(_ path: String) -> sockaddr_un? {
    var address = sockaddr_un()
    address.sun_family = sa_family_t(AF_UNIX)
    let bytes = Array(path.utf8)
    let capacity = MemoryLayout.size(ofValue: address.sun_path)
    guard bytes.count < capacity else { return nil }
    withUnsafeMutableBytes(of: &address.sun_path) { raw in
      raw.copyBytes(from: bytes)
      raw[bytes.count] = 0
    }
    address.sun_len = UInt8(MemoryLayout<sockaddr_un>.size)
    return address
  }

  /// A socket that will not raise SIGPIPE when the far side has gone: a
  /// write to a closed peer answers EPIPE instead, which is the ending
  /// both sides are listening for, rather than killing the process.
  static func make() throws -> Int32 {
    let fd = Darwin.socket(AF_UNIX, SOCK_STREAM, 0)
    guard fd >= 0 else { throw Failure(description: "socket: \(errno)") }
    var on: Int32 = 1
    setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &on, socklen_t(MemoryLayout<Int32>.size))
    return fd
  }

  static func listen(path: String) throws -> Int32 {
    guard var address = address(path) else { throw Failure(description: "socket path too long") }
    unlink(path)
    let fd = try make()
    let bound = withUnsafePointer(to: &address) {
      $0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
        bind(fd, $0, socklen_t(MemoryLayout<sockaddr_un>.size))
      }
    }
    guard bound == 0, Darwin.listen(fd, 1) == 0 else {
      let code = errno
      Darwin.close(fd)
      throw Failure(description: "listen: \(code)")
    }
    return fd
  }

  static func connect(path: String) throws -> Int32 {
    guard var address = address(path) else { throw Failure(description: "socket path too long") }
    let fd = try make()
    let connected = withUnsafePointer(to: &address) {
      $0.withMemoryRebound(to: sockaddr.self, capacity: 1) {
        Darwin.connect(fd, $0, socklen_t(MemoryLayout<sockaddr_un>.size))
      }
    }
    guard connected == 0 else {
      let code = errno
      Darwin.close(fd)
      throw Failure(description: "connect: \(code)")
    }
    return fd
  }

  static func nonBlocking(_ fd: Int32) {
    _ = fcntl(fd, F_SETFL, fcntl(fd, F_GETFL) | O_NONBLOCK)
  }
}

/// A frame size that fits within `longest` on its long side, keeps its
/// shape, and is even in both directions, which is what 4:2:0 chroma and
/// the encoders need. The Mac's screen capture and the iPhone's broadcast
/// both size their frames with it.
public func fitted(width: CGFloat, height: CGFloat, longest: CGFloat) -> (width: Int, height: Int) {
  guard width > 0, height > 0 else { return (2, 2) }
  let scale = min(1, longest / max(width, height))
  func even(_ value: CGFloat) -> Int { max(2, Int((value * scale).rounded(.down)) & ~1) }
  return (even(width), even(height))
}