import Foundation
import Vision

protocol SegmentationCoordinatorDelegate: AnyObject {
    func didUpdateSegmentationBuffer(_ buffer: CVPixelBuffer)
}

class SegmentationCoordinator {
    // Communicate with the session and other session objects on this queue.
    private let segmentationQueue = DispatchQueue(
        label: "SegmentationQueue",
        attributes: [],
        autoreleaseFrequency: .workItem
    )

    // The Vision requests and the handler to perform them.
    private let requestHandler = VNSequenceRequestHandler()
    private var facePoseRequest: VNDetectFaceRectanglesRequest?
    private var segmentationRequest = VNGeneratePersonSegmentationRequest()

    private var colors: AngleColors?
    weak var delegate: SegmentationCoordinatorDelegate?

    init() {
        self.intializeRequests()
    }

    private func intializeRequests() {
        // Create a request to detect face rectangles.
        facePoseRequest = VNDetectFaceRectanglesRequest { [weak self] request, _ in
            guard let face = request.results?.first as? VNFaceObservation else { return }
            // Generate RGB color intensity values for the face rectangle angles.
            self?.colors = AngleColors(roll: face.roll, pitch: face.pitch, yaw: face.yaw)
        }

        facePoseRequest?.revision = VNDetectFaceRectanglesRequestRevision3

        // Create a request to segment a person from an image.
        segmentationRequest = VNGeneratePersonSegmentationRequest()
        segmentationRequest.qualityLevel = .balanced
        segmentationRequest.outputPixelFormat = kCVPixelFormatType_OneComponent8
    }

    func processVideoFrame(_ framePixelBuffer: CVPixelBuffer) {
        // Perform the requests on the pixel buffer that contains the video frame.

        segmentationQueue.sync {
            guard let facePoseRequest = facePoseRequest else { return }
            let requests: [VNRequest] = [facePoseRequest, segmentationRequest]
            try? requestHandler.perform(
                requests,
                on: framePixelBuffer,
                orientation: .right
            )

            // Get the pixel buffer that contains the mask image.
            guard
                let maskPixelBuffer =
                segmentationRequest.results?.first?.pixelBuffer
            else { return }

            DispatchQueue.main.async { [weak self] in
                self?.delegate?.didUpdateSegmentationBuffer(maskPixelBuffer)
            }
        }
    }
}