// personmatte.swift — Apple Vision person-matte extractor for masked luxury grading.
//
//   Usage: personmatte <video> <outW> <outH> [fast|balanced|accurate]
//
// Decodes the video SEQUENTIALLY (AVAssetReader — never per-frame seeks), runs ONE
// reused VNGeneratePersonSegmentationRequest (Apple's documented pattern: the request
// keeps temporal state across frames, which stabilises the matte), scales each matte
// to <outW>x<outH> and writes raw 8-bit gray frames to stdout (outW*outH bytes per
// frame, full range 0-255, white = person). Pipe into:
//   ffmpeg -f rawvideo -pix_fmt gray -s WxH -r FPS -i - -c:v libx264 -crf 6 \
//          -pix_fmt gray -color_range pc matte.mp4
// Prints "personmatte: frames=N coverage=F" to stderr on completion.
//
//   Build: swiftc -O -o personmatte personmatte.swift

import AVFoundation
import CoreImage
import Foundation
import Vision

func fail(_ msg: String) -> Never {
    FileHandle.standardError.write(Data(("personmatte: " + msg + "\n").utf8))
    exit(1)
}

let args = CommandLine.arguments
guard args.count >= 4, let outW = Int(args[2]), let outH = Int(args[3]), outW > 0, outH > 0 else {
    fail("usage: personmatte <video> <outW> <outH> [fast|balanced|accurate]")
}
let url = URL(fileURLWithPath: args[1])
let quality = args.count > 4 ? args[4] : "accurate"

let asset = AVURLAsset(url: url)
// loadTracks(withMediaType:) is the non-deprecated async API — bridge to sync.
var loadedTrack: AVAssetTrack?
let sem = DispatchSemaphore(value: 0)
Task {
    loadedTrack = try? await asset.loadTracks(withMediaType: .video).first
    sem.signal()
}
sem.wait()
guard let track = loadedTrack else { fail("no video track in \(args[1])") }

guard let reader = try? AVAssetReader(asset: asset) else { fail("cannot open AVAssetReader") }
let output = AVAssetReaderTrackOutput(
    track: track,
    outputSettings: [kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA])
output.alwaysCopiesSampleData = false
reader.add(output)
guard reader.startReading() else { fail("reader failed: \(String(describing: reader.error))") }

// ONE request instance reused across every frame (temporal smoothing).
let request = VNGeneratePersonSegmentationRequest()
switch quality {
case "fast": request.qualityLevel = .fast
case "balanced": request.qualityLevel = .balanced
default: request.qualityLevel = .accurate
}
request.outputPixelFormat = kCVPixelFormatType_OneComponent8

// No color management — matte values pass through untouched.
let ciContext = CIContext(options: [.workingColorSpace: NSNull(), .outputColorSpace: NSNull()])
var scaledPB: CVPixelBuffer?
CVPixelBufferCreate(kCFAllocatorDefault, outW, outH, kCVPixelFormatType_OneComponent8, nil, &scaledPB)
guard let outBuf = scaledPB else { fail("cannot allocate \(outW)x\(outH) matte buffer") }

let out = FileHandle.standardOutput
let black = Data(count: outW * outH)
var frames = 0
var coverage = 0.0

while let sample = output.copyNextSampleBuffer() {
    guard let frame = CMSampleBufferGetImageBuffer(sample) else { continue }
    do { try VNImageRequestHandler(cvPixelBuffer: frame, options: [:]).perform([request]) }
    catch { fail("vision inference failed at frame \(frames): \(error)") }

    var data: Data
    if let mask = request.results?.first?.pixelBuffer {
        // Vision's matte resolution depends on qualityLevel — scale to output dims.
        let mw = CVPixelBufferGetWidth(mask), mh = CVPixelBufferGetHeight(mask)
        let scaled = CIImage(cvPixelBuffer: mask).transformed(
            by: CGAffineTransform(scaleX: CGFloat(outW) / CGFloat(mw),
                                  y: CGFloat(outH) / CGFloat(mh)))
        ciContext.render(scaled, to: outBuf,
                         bounds: CGRect(x: 0, y: 0, width: outW, height: outH), colorSpace: nil)
        CVPixelBufferLockBaseAddress(outBuf, .readOnly)
        let base = CVPixelBufferGetBaseAddress(outBuf)!
        let stride = CVPixelBufferGetBytesPerRow(outBuf)
        data = Data(capacity: outW * outH)
        for y in 0..<outH { data.append(Data(bytes: base + y * stride, count: outW)) }
        CVPixelBufferUnlockBaseAddress(outBuf, .readOnly)
    } else {
        data = black  // no person found this frame
    }

    // Sampled mean for the no-person fallback signal (~every 251st pixel).
    var sum = 0, n = 0
    data.withUnsafeBytes { (p: UnsafeRawBufferPointer) in
        var i = 0
        while i < p.count { sum += Int(p[i]); n += 1; i += 251 }
    }
    coverage += Double(sum) / (255.0 * Double(max(n, 1)))
    out.write(data)
    frames += 1
}
if reader.status == .failed { fail("decode failed: \(String(describing: reader.error))") }
let mean = frames > 0 ? coverage / Double(frames) : 0
FileHandle.standardError.write(
    Data("personmatte: frames=\(frames) coverage=\(String(format: "%.4f", mean))\n".utf8))