import CoreMedia
import CoreVideo
import Foundation

struct PPGFrameColorMetrics: Sendable {
  let timestampSeconds: Double
  let redMean: Double
  let greenMean: Double
  let blueMean: Double
  let redVariance: Double
  let greenVariance: Double
  let luminanceMean: Double
  let saturationRatio: Double
}

enum PPGFrameMetricsExtractorError: Error {
  case invalidTimestamp
  case unsupportedPixelFormat(OSType)
  case pixelBufferLockFailed(CVReturn)
  case missingPlane
  case missingBaseAddress
  case emptyRegionOfInterest
}

/// Extracts compact color statistics from a central ROI without retaining or copying a frame.
final class PPGFrameMetricsExtractor {
  private let gridSize: Int
  private let roiFraction: Double

  init(gridSize: Int = 32, roiFraction: Double = 0.55) {
    self.gridSize = max(8, gridSize)
    self.roiFraction = min(max(roiFraction, 0.2), 0.9)
  }

  func extract(from sampleBuffer: CMSampleBuffer) throws -> PPGFrameColorMetrics {
    let timestamp = CMTimeGetSeconds(CMSampleBufferGetPresentationTimeStamp(sampleBuffer))
    guard timestamp.isFinite, timestamp >= 0 else {
      throw PPGFrameMetricsExtractorError.invalidTimestamp
    }
    guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
      throw PPGFrameMetricsExtractorError.missingBaseAddress
    }

    let lockFlags = CVPixelBufferLockFlags.readOnly
    let lockStatus = CVPixelBufferLockBaseAddress(pixelBuffer, lockFlags)
    guard lockStatus == kCVReturnSuccess else {
      throw PPGFrameMetricsExtractorError.pixelBufferLockFailed(lockStatus)
    }
    defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, lockFlags) }

    let pixelFormat = CVPixelBufferGetPixelFormatType(pixelBuffer)
    switch pixelFormat {
    case kCVPixelFormatType_420YpCbCr8BiPlanarFullRange:
      return try extractBiPlanarYUV(
        from: pixelBuffer,
        timestamp: timestamp,
        isFullRange: true
      )
    case kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange:
      return try extractBiPlanarYUV(
        from: pixelBuffer,
        timestamp: timestamp,
        isFullRange: false
      )
    case kCVPixelFormatType_32BGRA:
      return try extractBGRA(from: pixelBuffer, timestamp: timestamp)
    default:
      throw PPGFrameMetricsExtractorError.unsupportedPixelFormat(pixelFormat)
    }
  }

  private func extractBiPlanarYUV(
    from pixelBuffer: CVPixelBuffer,
    timestamp: Double,
    isFullRange: Bool
  ) throws -> PPGFrameColorMetrics {
    guard CVPixelBufferGetPlaneCount(pixelBuffer) >= 2 else {
      throw PPGFrameMetricsExtractorError.missingPlane
    }
    guard let yBase = CVPixelBufferGetBaseAddressOfPlane(pixelBuffer, 0),
          let uvBase = CVPixelBufferGetBaseAddressOfPlane(pixelBuffer, 1)
    else {
      throw PPGFrameMetricsExtractorError.missingBaseAddress
    }

    let width = CVPixelBufferGetWidthOfPlane(pixelBuffer, 0)
    let height = CVPixelBufferGetHeightOfPlane(pixelBuffer, 0)
    let yStride = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, 0)
    let uvStride = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, 1)
    let yBytes = yBase.assumingMemoryBound(to: UInt8.self)
    let uvBytes = uvBase.assumingMemoryBound(to: UInt8.self)
    let roi = regionOfInterest(width: width, height: height)
    let stepX = max(roi.width / gridSize, 1)
    let stepY = max(roi.height / gridSize, 1)
    var accumulator = ColorAccumulator()

    for y in stride(from: roi.minY, to: roi.maxY, by: stepY) {
      for x in stride(from: roi.minX, to: roi.maxX, by: stepX) {
        let yValue = Double(yBytes[y * yStride + x])
        let uvOffset = (y / 2) * uvStride + (x / 2) * 2
        let cb = Double(uvBytes[uvOffset])
        let cr = Double(uvBytes[uvOffset + 1])
        let rgb = convertYCbCrToRGB(y: yValue, cb: cb, cr: cr, isFullRange: isFullRange)
        accumulator.add(red: rgb.red, green: rgb.green, blue: rgb.blue, luminance: yValue)
      }
    }

    return try accumulator.result(timestamp: timestamp)
  }

  private func extractBGRA(
    from pixelBuffer: CVPixelBuffer,
    timestamp: Double
  ) throws -> PPGFrameColorMetrics {
    guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else {
      throw PPGFrameMetricsExtractorError.missingBaseAddress
    }
    let width = CVPixelBufferGetWidth(pixelBuffer)
    let height = CVPixelBufferGetHeight(pixelBuffer)
    let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
    let bytes = baseAddress.assumingMemoryBound(to: UInt8.self)
    let roi = regionOfInterest(width: width, height: height)
    let stepX = max(roi.width / gridSize, 1)
    let stepY = max(roi.height / gridSize, 1)
    var accumulator = ColorAccumulator()

    for y in stride(from: roi.minY, to: roi.maxY, by: stepY) {
      for x in stride(from: roi.minX, to: roi.maxX, by: stepX) {
        let offset = y * bytesPerRow + x * 4
        let blue = Double(bytes[offset])
        let green = Double(bytes[offset + 1])
        let red = Double(bytes[offset + 2])
        let luminance = 0.2126 * red + 0.7152 * green + 0.0722 * blue
        accumulator.add(red: red, green: green, blue: blue, luminance: luminance)
      }
    }

    return try accumulator.result(timestamp: timestamp)
  }

  private func regionOfInterest(width: Int, height: Int) -> (
    minX: Int,
    maxX: Int,
    minY: Int,
    maxY: Int,
    width: Int,
    height: Int
  ) {
    let roiWidth = max(Int(Double(width) * roiFraction), 1)
    let roiHeight = max(Int(Double(height) * roiFraction), 1)
    let minX = max((width - roiWidth) / 2, 0)
    let minY = max((height - roiHeight) / 2, 0)
    return (minX, minX + roiWidth, minY, minY + roiHeight, roiWidth, roiHeight)
  }

  private func convertYCbCrToRGB(
    y: Double,
    cb: Double,
    cr: Double,
    isFullRange: Bool
  ) -> (red: Double, green: Double, blue: Double) {
    let red: Double
    let green: Double
    let blue: Double

    if isFullRange {
      let shiftedCb = cb - 128
      let shiftedCr = cr - 128
      red = y + 1.402 * shiftedCr
      green = y - 0.344_136 * shiftedCb - 0.714_136 * shiftedCr
      blue = y + 1.772 * shiftedCb
    } else {
      let shiftedY = max(y - 16, 0)
      let shiftedCb = cb - 128
      let shiftedCr = cr - 128
      red = 1.164_383 * shiftedY + 1.596_027 * shiftedCr
      green = 1.164_383 * shiftedY - 0.391_762 * shiftedCb - 0.812_968 * shiftedCr
      blue = 1.164_383 * shiftedY + 2.017_232 * shiftedCb
    }

    return (clampByte(red), clampByte(green), clampByte(blue))
  }

  private func clampByte(_ value: Double) -> Double {
    min(max(value, 0), 255)
  }
}

private struct ColorAccumulator {
  private(set) var count = 0
  private var redMean = 0.0
  private var greenMean = 0.0
  private var blueMean = 0.0
  private var luminanceMean = 0.0
  private var redM2 = 0.0
  private var greenM2 = 0.0
  private var saturatedCount = 0

  mutating func add(red: Double, green: Double, blue: Double, luminance: Double) {
    count += 1
    let n = Double(count)

    let redDelta = red - redMean
    redMean += redDelta / n
    redM2 += redDelta * (red - redMean)

    let greenDelta = green - greenMean
    greenMean += greenDelta / n
    greenM2 += greenDelta * (green - greenMean)

    blueMean += (blue - blueMean) / n
    luminanceMean += (luminance - luminanceMean) / n
    // A fingertip lit by the torch commonly clips only the red channel while
    // green still contains a clean pulsatile waveform. Count true optical
    // clipping only when both useful color channels (or luminance) are clipped.
    if (red >= 250 && green >= 250) || luminance >= 250 {
      saturatedCount += 1
    }
  }

  func result(timestamp: Double) throws -> PPGFrameColorMetrics {
    guard count > 0 else {
      throw PPGFrameMetricsExtractorError.emptyRegionOfInterest
    }
    let divisor = Double(max(count - 1, 1))
    return PPGFrameColorMetrics(
      timestampSeconds: timestamp,
      redMean: redMean,
      greenMean: greenMean,
      blueMean: blueMean,
      redVariance: redM2 / divisor,
      greenVariance: greenM2 / divisor,
      luminanceMean: luminanceMean,
      saturationRatio: Double(saturatedCount) / Double(count)
    )
  }
}