import CoreMedia
import CoreVideo
import XCTest
@testable import HaptikaFrameMetrics

final class PPGFrameMetricsExtractorTests: XCTestCase {
  func testFullRangeYUVExtractionUsesCentralROIAndUnlocksBuffer() throws {
    let pixelBuffer = try makePixelBuffer(
      width: 64,
      height: 64,
      pixelFormat: kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
    )
    try fillBiPlanarYUV(pixelBuffer, y: 100, cb: 128, cr: 128)
    let metrics = try PPGFrameMetricsExtractor(gridSize: 16, roiFraction: 0.5)
      .extract(from: makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 4.25))

    XCTAssertEqual(metrics.timestampSeconds, 4.25, accuracy: 0.000_001)
    XCTAssertEqual(metrics.redMean, 100, accuracy: 0.000_001)
    XCTAssertEqual(metrics.greenMean, 100, accuracy: 0.000_001)
    XCTAssertEqual(metrics.blueMean, 100, accuracy: 0.000_001)
    XCTAssertEqual(metrics.luminanceMean, 100, accuracy: 0.000_001)
    try assertCanRelock(pixelBuffer)
  }

  func testVideoRangeYUVExtractionAppliesVideoRangeConversion() throws {
    let pixelBuffer = try makePixelBuffer(
      width: 48,
      height: 48,
      pixelFormat: kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
    )
    try fillBiPlanarYUV(pixelBuffer, y: 126, cb: 128, cr: 128)
    let metrics = try PPGFrameMetricsExtractor(gridSize: 12, roiFraction: 0.5)
      .extract(from: makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 2))

    let expected = 1.164_383 * 110
    XCTAssertEqual(metrics.redMean, expected, accuracy: 0.001)
    XCTAssertEqual(metrics.greenMean, expected, accuracy: 0.001)
    XCTAssertEqual(metrics.blueMean, expected, accuracy: 0.001)
    try assertCanRelock(pixelBuffer)
  }

  func testBGRAExtractionSamplesOnlyCentralROI() throws {
    let pixelBuffer = try makePixelBuffer(
      width: 64,
      height: 64,
      pixelFormat: kCVPixelFormatType_32BGRA
    )
    try fillBGRA(pixelBuffer) { x, y in
      if (16..<48).contains(x), (16..<48).contains(y) {
        return (blue: 20, green: 80, red: 200, alpha: 255)
      }
      return (blue: 240, green: 10, red: 5, alpha: 255)
    }
    let sampleBuffer = try makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 12.5)

    let metrics = try PPGFrameMetricsExtractor(gridSize: 16, roiFraction: 0.5)
      .extract(from: sampleBuffer)

    XCTAssertEqual(metrics.timestampSeconds, 12.5, accuracy: 0.000_001)
    XCTAssertEqual(metrics.redMean, 200, accuracy: 0.000_001)
    XCTAssertEqual(metrics.greenMean, 80, accuracy: 0.000_001)
    XCTAssertEqual(metrics.blueMean, 20, accuracy: 0.000_001)
    XCTAssertEqual(metrics.redVariance, 0, accuracy: 0.000_001)
    XCTAssertEqual(metrics.greenVariance, 0, accuracy: 0.000_001)
    XCTAssertEqual(
      metrics.luminanceMean,
      0.2126 * 200 + 0.7152 * 80 + 0.0722 * 20,
      accuracy: 0.000_001
    )
    XCTAssertEqual(metrics.saturationRatio, 0, accuracy: 0.000_001)
    try assertCanRelock(pixelBuffer)
  }

  func testBGRARedOnlyClippingKeepsUsableGreenChannel() throws {
    let pixelBuffer = try makePixelBuffer(
      width: 48,
      height: 48,
      pixelFormat: kCVPixelFormatType_32BGRA
    )
    try fillBGRA(pixelBuffer) { x, y in
      let isCentral = (12..<36).contains(x) && (12..<36).contains(y)
      return isCentral
        ? (blue: 40, green: 100, red: 255, alpha: 255)
        : (blue: 0, green: 0, red: 0, alpha: 255)
    }
    let sampleBuffer = try makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 1)

    let metrics = try PPGFrameMetricsExtractor(gridSize: 12, roiFraction: 0.5)
      .extract(from: sampleBuffer)

    XCTAssertEqual(metrics.redMean, 255, accuracy: 0.000_001)
    XCTAssertEqual(metrics.saturationRatio, 0, accuracy: 0.000_001)
    try assertCanRelock(pixelBuffer)
  }

  func testBGRASaturationRequiresBothUsefulChannelsToClip() throws {
    let pixelBuffer = try makePixelBuffer(
      width: 48,
      height: 48,
      pixelFormat: kCVPixelFormatType_32BGRA
    )
    try fillBGRA(pixelBuffer) { _, _ in
      (blue: 80, green: 255, red: 255, alpha: 255)
    }
    let sampleBuffer = try makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 1)

    let metrics = try PPGFrameMetricsExtractor(gridSize: 12, roiFraction: 0.5)
      .extract(from: sampleBuffer)

    XCTAssertEqual(metrics.saturationRatio, 1, accuracy: 0.000_001)
    try assertCanRelock(pixelBuffer)
  }

  func testUnsupportedPixelFormatUnlocksBufferOnThrow() throws {
    let pixelBuffer = try makePixelBuffer(
      width: 32,
      height: 32,
      pixelFormat: kCVPixelFormatType_OneComponent8
    )
    let sampleBuffer = try makeSampleBuffer(pixelBuffer: pixelBuffer, timestampSeconds: 2)

    XCTAssertThrowsError(try PPGFrameMetricsExtractor().extract(from: sampleBuffer)) { error in
      guard case PPGFrameMetricsExtractorError.unsupportedPixelFormat(let format) = error else {
        return XCTFail("Expected unsupportedPixelFormat, got \(error)")
      }
      XCTAssertEqual(format, kCVPixelFormatType_OneComponent8)
    }
    try assertCanRelock(pixelBuffer)
  }

  func testRepeatedExtractionDoesNotLeavePixelBuffersLocked() throws {
    let extractor = PPGFrameMetricsExtractor(gridSize: 8, roiFraction: 0.5)
    for index in 0..<100 {
      autoreleasepool {
        do {
          let pixelBuffer = try makePixelBuffer(
            width: 32,
            height: 32,
            pixelFormat: kCVPixelFormatType_32BGRA
          )
          try fillBGRA(pixelBuffer) { _, _ in (20, 90, 180, 255) }
          _ = try extractor.extract(
            from: makeSampleBuffer(
              pixelBuffer: pixelBuffer,
              timestampSeconds: Double(index) / 30
            )
          )
          try assertCanRelock(pixelBuffer)
        } catch {
          XCTFail("Repeated extraction failed at \(index): \(error)")
        }
      }
    }
  }

  private func makePixelBuffer(
    width: Int,
    height: Int,
    pixelFormat: OSType
  ) throws -> CVPixelBuffer {
    if pixelFormat == kCVPixelFormatType_420YpCbCr8BiPlanarFullRange
        || pixelFormat == kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange {
      return try makeBiPlanarPixelBuffer(width: width, height: height, pixelFormat: pixelFormat)
    }
    var pixelBuffer: CVPixelBuffer?
    let status = CVPixelBufferCreate(
      kCFAllocatorDefault,
      width,
      height,
      pixelFormat,
      nil,
      &pixelBuffer
    )
    guard status == kCVReturnSuccess, let pixelBuffer else {
      throw TestBufferError.pixelBufferCreation(status)
    }
    return pixelBuffer
  }

  private func makeBiPlanarPixelBuffer(
    width: Int,
    height: Int,
    pixelFormat: OSType
  ) throws -> CVPixelBuffer {
    let yBytesPerRow = width
    let uvBytesPerRow = width
    let ySize = yBytesPerRow * height
    let uvHeight = height / 2
    let totalSize = ySize + uvBytesPerRow * uvHeight
    guard let storage = malloc(totalSize) else {
      throw TestBufferError.missingBaseAddress
    }
    var planeAddresses: [UnsafeMutableRawPointer?] = [
      storage,
      storage.advanced(by: ySize),
    ]
    var planeWidths = [width, width / 2]
    var planeHeights = [height, uvHeight]
    var planeBytesPerRow = [yBytesPerRow, uvBytesPerRow]
    var pixelBuffer: CVPixelBuffer?
    let status = CVPixelBufferCreateWithPlanarBytes(
      kCFAllocatorDefault,
      width,
      height,
      pixelFormat,
      storage,
      totalSize,
      2,
      &planeAddresses,
      &planeWidths,
      &planeHeights,
      &planeBytesPerRow,
      { _, dataPointer, _, _, _ in
        free(UnsafeMutableRawPointer(mutating: dataPointer))
      },
      nil,
      nil,
      &pixelBuffer
    )
    guard status == kCVReturnSuccess, let pixelBuffer else {
      free(storage)
      throw TestBufferError.pixelBufferCreation(status)
    }
    return pixelBuffer
  }

  private func fillBiPlanarYUV(
    _ pixelBuffer: CVPixelBuffer,
    y: UInt8,
    cb: UInt8,
    cr: UInt8
  ) throws {
    let status = CVPixelBufferLockBaseAddress(pixelBuffer, [])
    guard status == kCVReturnSuccess else {
      throw TestBufferError.pixelBufferLock(status)
    }
    defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, []) }
    guard CVPixelBufferGetPlaneCount(pixelBuffer) >= 2,
          let yBase = CVPixelBufferGetBaseAddressOfPlane(pixelBuffer, 0),
          let uvBase = CVPixelBufferGetBaseAddressOfPlane(pixelBuffer, 1)
    else {
      throw TestBufferError.missingBaseAddress
    }
    let yHeight = CVPixelBufferGetHeightOfPlane(pixelBuffer, 0)
    let yStride = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, 0)
    let yBytes = yBase.assumingMemoryBound(to: UInt8.self)
    for row in 0..<yHeight {
      yBytes.advanced(by: row * yStride).initialize(repeating: y, count: yStride)
    }
    let uvWidth = CVPixelBufferGetWidthOfPlane(pixelBuffer, 1)
    let uvHeight = CVPixelBufferGetHeightOfPlane(pixelBuffer, 1)
    let uvStride = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, 1)
    let uvBytes = uvBase.assumingMemoryBound(to: UInt8.self)
    for row in 0..<uvHeight {
      for column in 0..<uvWidth {
        let offset = row * uvStride + column * 2
        uvBytes[offset] = cb
        uvBytes[offset + 1] = cr
      }
    }
  }

  private func fillBGRA(
    _ pixelBuffer: CVPixelBuffer,
    pixel: (_ x: Int, _ y: Int) -> (blue: UInt8, green: UInt8, red: UInt8, alpha: UInt8)
  ) throws {
    let lockStatus = CVPixelBufferLockBaseAddress(pixelBuffer, [])
    guard lockStatus == kCVReturnSuccess else {
      throw TestBufferError.pixelBufferLock(lockStatus)
    }
    defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, []) }

    guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else {
      throw TestBufferError.missingBaseAddress
    }
    let width = CVPixelBufferGetWidth(pixelBuffer)
    let height = CVPixelBufferGetHeight(pixelBuffer)
    let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
    let bytes = baseAddress.assumingMemoryBound(to: UInt8.self)

    for y in 0..<height {
      for x in 0..<width {
        let color = pixel(x, y)
        let offset = y * bytesPerRow + x * 4
        bytes[offset] = color.blue
        bytes[offset + 1] = color.green
        bytes[offset + 2] = color.red
        bytes[offset + 3] = color.alpha
      }
    }
  }

  private func makeSampleBuffer(
    pixelBuffer: CVPixelBuffer,
    timestampSeconds: Double
  ) throws -> CMSampleBuffer {
    var formatDescription: CMVideoFormatDescription?
    let formatStatus = CMVideoFormatDescriptionCreateForImageBuffer(
      allocator: kCFAllocatorDefault,
      imageBuffer: pixelBuffer,
      formatDescriptionOut: &formatDescription
    )
    guard formatStatus == noErr, let formatDescription else {
      throw TestBufferError.formatDescriptionCreation(formatStatus)
    }

    var timing = CMSampleTimingInfo(
      duration: .invalid,
      presentationTimeStamp: CMTime(seconds: timestampSeconds, preferredTimescale: 60_000),
      decodeTimeStamp: .invalid
    )
    var sampleBuffer: CMSampleBuffer?
    let sampleStatus = CMSampleBufferCreateReadyWithImageBuffer(
      allocator: kCFAllocatorDefault,
      imageBuffer: pixelBuffer,
      formatDescription: formatDescription,
      sampleTiming: &timing,
      sampleBufferOut: &sampleBuffer
    )
    guard sampleStatus == noErr, let sampleBuffer else {
      throw TestBufferError.sampleBufferCreation(sampleStatus)
    }
    return sampleBuffer
  }

  private func assertCanRelock(_ pixelBuffer: CVPixelBuffer) throws {
    let status = CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
    guard status == kCVReturnSuccess else {
      throw TestBufferError.pixelBufferLock(status)
    }
    CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly)
  }
}

private enum TestBufferError: Error {
  case pixelBufferCreation(CVReturn)
  case pixelBufferLock(CVReturn)
  case missingBaseAddress
  case formatDescriptionCreation(OSStatus)
  case sampleBufferCreation(OSStatus)
}