markbaseengine/Sources/MarkBaseServer/VideoProcessor.swift

import Accelerate
import AVFoundation
import CoreVideo
import Foundation

// ─────────────────────────────────────────────────────────────
// Video Processor — Extract frames + audio from video files
// ─────────────────────────────────────────────────────────────

public struct VideoFrame {
    public let index: Int
    public let timestamp: CMTime
    public let pixelBuffer: CVPixelBuffer
}

public struct VideoData {
    public let frames: [VideoFrame]
    public let audioSamples: [Float]
    public let sampleRate: Int
    public let duration: CMTime
    public let naturalSize: CGSize
    public let estimatedFrameRate: Float
}

public enum VideoProcessor {
    public struct Config {
        public let maxFrames: Int
        public let targetFPS: Float
        public let audioSampleRate: Int
        public let sceneThreshold: Double  // 0…1; 0 = fixed FPS, >0 = scene detection

        public init(maxFrames: Int = 64, targetFPS: Float = 2.0,
                    audioSampleRate: Int = 16000, sceneThreshold: Double = 0.15) {
            self.maxFrames = maxFrames
            self.targetFPS = targetFPS
            self.audioSampleRate = audioSampleRate
            self.sceneThreshold = sceneThreshold
        }
    }

    /// Compute luminance histogram difference between two pixel buffers (0…1, higher = more different).
    public static func sceneDiff(_ a: CVPixelBuffer, _ b: CVPixelBuffer) -> Double {
        CVPixelBufferLockBaseAddress(a, .readOnly)
        CVPixelBufferLockBaseAddress(b, .readOnly)
        defer {
            CVPixelBufferUnlockBaseAddress(a, .readOnly)
            CVPixelBufferUnlockBaseAddress(b, .readOnly)
        }

        let w = min(CVPixelBufferGetWidth(a), CVPixelBufferGetWidth(b))
        let h = min(CVPixelBufferGetHeight(a), CVPixelBufferGetHeight(b))
        let rowA = CVPixelBufferGetBytesPerRow(a)
        let rowB = CVPixelBufferGetBytesPerRow(b)
        guard let baseA = CVPixelBufferGetBaseAddress(a),
              let baseB = CVPixelBufferGetBaseAddress(b) else { return 0 }

        let ptrA = baseA.assumingMemoryBound(to: UInt8.self)
        let ptrB = baseB.assumingMemoryBound(to: UInt8.self)

        var histA = [Double](repeating: 0, count: 256)
        var histB = [Double](repeating: 0, count: 256)
        let total = w * h

        for y in 0..<h {
            for x in 0..<w {
                // Luminance ≈ green channel (BGRA → index 1)
                histA[Int(ptrA[y * rowA + x * 4 + 1])] += 1
                histB[Int(ptrB[y * rowB + x * 4 + 1])] += 1
            }
        }

        // Normalise and compute intersection
        var diff: Double = 0
        for i in 0..<256 {
            let normA = histA[i] / Double(total)
            let normB = histB[i] / Double(total)
            diff += abs(normA - normB)
        }
        return diff / 2  // 0..1
    }

    public static func process(url: URL, config: Config = Config()) async throws -> VideoData {
        let asset = AVURLAsset(url: url)
        let duration = try await asset.load(.duration)

        let videoTracks = try await asset.loadTracks(withMediaType: .video)
        let audioTracks = try await asset.loadTracks(withMediaType: .audio)

        guard let videoTrack = videoTracks.first else {
            throw MarkBaseError.invalidParameter(parameter: "url", message: "No video track found")
        }

        let naturalSize = try await videoTrack.load(.naturalSize)
        let nominalFrameRate = try await videoTrack.load(.nominalFrameRate)

        // Read audio samples concurrently
        let audioSamples: [Float]
        if let audioTrack = audioTracks.first {
            audioSamples = try readAudioTrack(audioTrack, sampleRate: config.audioSampleRate)
        } else {
            audioSamples = []
        }

        // Read video frames (scene detection or fixed FPS)
        let frames: [VideoFrame]
        if config.sceneThreshold > 0 {
            frames = try await readVideoTrackSceneDetect(
                videoTrack,
                duration: duration,
                threshold: config.sceneThreshold,
                maxFrames: config.maxFrames
            )
            print("  Scene detection: \(frames.count) keyframes")
        } else {
            frames = try await readVideoTrack(
                videoTrack,
                duration: duration,
                nominalFrameRate: nominalFrameRate,
                targetFPS: config.targetFPS,
                maxFrames: config.maxFrames
            )
        }

        return VideoData(
            frames: frames,
            audioSamples: audioSamples,
            sampleRate: config.audioSampleRate,
            duration: duration,
            naturalSize: naturalSize,
            estimatedFrameRate: nominalFrameRate
        )
    }

    // ── Video reading ─────────────────────────────────

    private static func readVideoTrack(
        _ track: AVAssetTrack,
        duration: CMTime,
        nominalFrameRate: Float,
        targetFPS: Float,
        maxFrames: Int
    ) async throws -> [VideoFrame] {
        let reader = try AVAssetReader(asset: track.asset!)

        let formatDescriptions = try await track.load(.formatDescriptions)
        let pixelFormat: OSType
        if let firstDesc = formatDescriptions.first {
            pixelFormat = CMFormatDescriptionGetMediaSubType(firstDesc)
        } else {
            pixelFormat = kCVPixelFormatType_32BGRA
        }

        let settings: [String: Any] = [
            kCVPixelBufferPixelFormatTypeKey as String: pixelFormat,
            kCVPixelBufferMetalCompatibilityKey as String: true,
        ]

        let output = AVAssetReaderTrackOutput(track: track, outputSettings: settings)
        reader.add(output)
        reader.startReading()

        defer {
            if reader.status == .reading {
                reader.cancelReading()
            }
        }

        let stepSeconds = CMTime(value: CMTimeValue(1.0 / targetFPS), timescale: CMTimeScale(targetFPS * 100))
            .seconds

        var frames: [VideoFrame] = []
        var lastSampleTime: Double = -stepSeconds

        while reader.status == .reading, frames.count < maxFrames {
            guard let sampleBuffer = output.copyNextSampleBuffer() else { break }

            let presentationTime = CMSampleBufferGetPresentationTimeStamp(sampleBuffer)
            let timeSeconds = presentationTime.seconds

            if timeSeconds - lastSampleTime >= stepSeconds {
                if let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) {
                    frames.append(VideoFrame(
                        index: frames.count,
                        timestamp: presentationTime,
                        pixelBuffer: pixelBuffer
                    ))
                }
                lastSampleTime = timeSeconds
            }
        }

        return frames
    }

    // ── Scene-detection reading ─────────────────────────

    private static func readVideoTrackSceneDetect(
        _ track: AVAssetTrack,
        duration: CMTime,
        threshold: Double,
        maxFrames: Int
    ) async throws -> [VideoFrame] {
        let reader = try AVAssetReader(asset: track.asset!)
        let formatDescriptions = try await track.load(.formatDescriptions)
        let pixelFormat: OSType
        if let firstDesc = formatDescriptions.first {
            pixelFormat = CMFormatDescriptionGetMediaSubType(firstDesc)
        } else {
            pixelFormat = kCVPixelFormatType_32BGRA
        }

        let settings: [String: Any] = [
            kCVPixelBufferPixelFormatTypeKey as String: pixelFormat,
            kCVPixelBufferMetalCompatibilityKey as String: true,
        ]

        let output = AVAssetReaderTrackOutput(track: track, outputSettings: settings)
        reader.add(output)
        reader.startReading()

        defer {
            if reader.status == .reading { reader.cancelReading() }
        }

        var frames: [VideoFrame] = []
        var prevBuffer: CVPixelBuffer?

        while reader.status == .reading, frames.count < maxFrames {
            guard let sampleBuffer = output.copyNextSampleBuffer() else { break }
            guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else { continue }

            if let prev = prevBuffer {
                let diff = sceneDiff(prev, pixelBuffer)
                if diff >= threshold || frames.isEmpty {
                    frames.append(VideoFrame(
                        index: frames.count,
                        timestamp: CMSampleBufferGetPresentationTimeStamp(sampleBuffer),
                        pixelBuffer: pixelBuffer
                    ))
                }
            } else {
                frames.append(VideoFrame(
                    index: frames.count,
                    timestamp: CMSampleBufferGetPresentationTimeStamp(sampleBuffer),
                    pixelBuffer: pixelBuffer
                ))
            }
            prevBuffer = pixelBuffer
        }

        return frames
    }

    // ── Audio reading ─────────────────────────────────

    private static func readAudioTrack(_ track: AVAssetTrack, sampleRate: Int) throws -> [Float] {
        let reader = try AVAssetReader(asset: track.asset!)

        let settings: [String: Any] = [
            AVFormatIDKey: kAudioFormatLinearPCM,
            AVLinearPCMIsFloatKey: true,
            AVLinearPCMBitDepthKey: 32,
            AVNumberOfChannelsKey: 1,
            AVSampleRateKey: sampleRate,
        ]

        let output = AVAssetReaderTrackOutput(track: track, outputSettings: settings)
        reader.add(output)
        reader.startReading()

        defer {
            if reader.status == .reading {
                reader.cancelReading()
            }
        }

        var samples: [Float] = []
        while reader.status == .reading {
            guard let sampleBuffer = output.copyNextSampleBuffer() else { break }
            guard let blockBuffer = CMSampleBufferGetDataBuffer(sampleBuffer) else { continue }

            let length = CMBlockBufferGetDataLength(blockBuffer)
            var data = [Float](repeating: 0, count: length / MemoryLayout<Float>.stride)
            CMBlockBufferCopyDataBytes(blockBuffer, atOffset: 0, dataLength: length, destination: &data)
            samples.append(contentsOf: data)
        }

        return samples
    }

    // ── Pixel buffer → float array ────────────────────

    public static func pixelBufferToFloats(_ pixelBuffer: CVPixelBuffer) -> [Float] {
        CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
        defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly) }

        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)
        let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)

        guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else { return [] }

        var floats = [Float](repeating: 0, count: width * height * 3) // RGB
        let pixelBuffer = baseAddress.assumingMemoryBound(to: UInt8.self)

        for y in 0..<height {
            for x in 0..<width {
                let offset = y * bytesPerRow + x * 4
                let b = Float(pixelBuffer[offset]) / 255.0
                let g = Float(pixelBuffer[offset + 1]) / 255.0
                let r = Float(pixelBuffer[offset + 2]) / 255.0
                let pixelOffset = (y * width + x) * 3
                floats[pixelOffset] = r
                floats[pixelOffset + 1] = g
                floats[pixelOffset + 2] = b
            }
        }

        return floats
    }

    /// Resize pixel buffer to target size using vImage
    public static func resizePixelBuffer(
        _ pixelBuffer: CVPixelBuffer,
        targetWidth: Int,
        targetHeight: Int
    ) -> CVPixelBuffer? {
        let srcWidth = CVPixelBufferGetWidth(pixelBuffer)
        let srcHeight = CVPixelBufferGetHeight(pixelBuffer)

        var srcBuffer = vImage_Buffer()
        CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
        srcBuffer.data = CVPixelBufferGetBaseAddress(pixelBuffer)
        srcBuffer.width = vImagePixelCount(srcWidth)
        srcBuffer.height = vImagePixelCount(srcHeight)
        srcBuffer.rowBytes = CVPixelBufferGetBytesPerRow(pixelBuffer)

        var destPixelBuffer: CVPixelBuffer?
        let attrs: [String: Any] = [
            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
            kCVPixelBufferMetalCompatibilityKey as String: true,
        ]
        CVPixelBufferCreate(
            kCFAllocatorDefault,
            targetWidth, targetHeight,
            kCVPixelFormatType_32BGRA,
            attrs as CFDictionary,
            &destPixelBuffer
        )

        guard let dest = destPixelBuffer else {
            CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly)
            return nil
        }

        CVPixelBufferLockBaseAddress(dest, [])
        var destBuffer = vImage_Buffer()
        destBuffer.data = CVPixelBufferGetBaseAddress(dest)
        destBuffer.width = vImagePixelCount(targetWidth)
        destBuffer.height = vImagePixelCount(targetHeight)
        destBuffer.rowBytes = CVPixelBufferGetBytesPerRow(dest)

        let scale = vImageScale_ARGB8888(&srcBuffer, &destBuffer, nil, vImage_Flags(kvImageNoFlags))

        CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly)
        CVPixelBufferUnlockBaseAddress(dest, [])

        return scale == kvImageNoError ? dest : nil
    }

    /// Frame to patch embeddings (simple 16×16 grid)
    public static func frameToPatchEmbeddings(
        _ pixelBuffer: CVPixelBuffer,
        patchSize: Int = 16
    ) -> (embeddings: [Float], numPatches: Int, patchDim: Int) {
        let rgb = pixelBufferToFloats(pixelBuffer)
        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)

        let numPatchesH = height / patchSize
        let numPatchesW = width / patchSize
        let numPatches = numPatchesH * numPatchesW
        let patchDim = patchSize * patchSize * 3

        var embeddings = [Float](repeating: 0, count: numPatches * patchDim)

        for ph in 0..<numPatchesH {
            for pw in 0..<numPatchesW {
                let patchIdx = ph * numPatchesW + pw
                for py in 0..<patchSize {
                    for px in 0..<patchSize {
                        let srcY = ph * patchSize + py
                        let srcX = pw * patchSize + px
                        let srcIdx = (srcY * width + srcX) * 3
                        let dstIdx = (patchIdx * patchDim) + (py * patchSize + px) * 3
                        if srcIdx + 2 < rgb.count, dstIdx + 2 < embeddings.count {
                            embeddings[dstIdx] = rgb[srcIdx]
                            embeddings[dstIdx + 1] = rgb[srcIdx + 1]
                            embeddings[dstIdx + 2] = rgb[srcIdx + 2]
                        }
                    }
                }
            }
        }

        return (embeddings, numPatches, patchDim)
    }
}