import Foundation
import Vision
import ArgumentParser
import AVFoundation

/// Swift Face Processor - Apple Vision Framework for face detection + pose
/// Uses AVAssetImageGenerator for reliable frame access (no AVAssetReader corruption).
@main
struct SwiftFace: ParsableCommand {
    @Argument(help: "Video file path")
    var inputPath: String

    @Argument(help: "Output JSON path")
    var outputPath: String

    @Option(name: .long, help: "Sample interval (frames, default=30)")
    var sampleInterval: Int = 30

    @Option(name: .long, help: "UUID for logging")
    var uuid: String = ""

    mutating func run() throws {
        let startTime = Date()
        print("[SwiftFace] Vision-based face detection: \(inputPath)")

        let url = URL(fileURLWithPath: inputPath)
        let asset = AVAsset(url: url)

        guard let videoTrack = asset.tracks(withMediaType: .video).first else {
            print("[SwiftFace] No video track found")
            return
        }

        let fps = videoTrack.nominalFrameRate
        let duration = CMTimeGetSeconds(asset.duration)
        let totalFrames = Int(duration * Double(fps))
        let width = Int(videoTrack.naturalSize.width)
        let height = Int(videoTrack.naturalSize.height)
        print("[SwiftFace] Video: \(width)x\(height), \(String(format: "%.1f", fps))fps, \(totalFrames) frames")

        let generator = AVAssetImageGenerator(asset: asset)
        generator.requestedTimeToleranceBefore = .zero
        generator.requestedTimeToleranceAfter = .zero
        generator.appliesPreferredTrackTransform = true

        var allFrames: [[String: Any]] = []
        var processedCount = 0
        var checkedCount = 0

        let frameInterval = TimeInterval(sampleInterval) / Double(fps)

        // Process in batches of 1000 frames to avoid memory pressure
        let batchSize = 1000
        let totalSamples = totalFrames / sampleInterval

        for batchStart in stride(from: 0, to: totalSamples, by: batchSize) {
            let batchEnd = min(batchStart + batchSize, totalSamples)
            var times: [NSValue] = []

            for i in batchStart..<batchEnd {
                let seconds = Double(i) * frameInterval
                let cmTime = CMTime(seconds: seconds, preferredTimescale: 1000)
                times.append(NSValue(time: cmTime))
            }

            let semaphore = DispatchSemaphore(value: 0)
            var batchError: Error? = nil

            generator.generateCGImagesAsynchronously(forTimes: times) { requestedTime, cgImage, actualTime, result, error in
                defer { semaphore.signal() }

                checkedCount += 1

                guard result == .succeeded, let cgImage = cgImage else {
                    if let error = error {
                        fputs("[SwiftFace] Frame error at \(CMTimeGetSeconds(requestedTime)): \(error.localizedDescription)\n", stderr)
                    }
                    return
                }

                // Convert CGImage to CVPixelBuffer for Vision
                var pixelBuffer: CVPixelBuffer?
                let attrs: [CFString: Any] = [
                    kCVPixelBufferCGImageCompatibilityKey: true,
                    kCVPixelBufferCGBitmapContextCompatibilityKey: true,
                    kCVPixelBufferWidthKey: cgImage.width,
                    kCVPixelBufferHeightKey: cgImage.height,
                ]
                CVPixelBufferCreate(kCFAllocatorDefault, cgImage.width, cgImage.height,
                    kCVPixelFormatType_32BGRA, attrs as CFDictionary, &pixelBuffer)

                guard let pb = pixelBuffer else { return }

                CVPixelBufferLockBaseAddress(pb, [])
                defer { CVPixelBufferUnlockBaseAddress(pb, []) }

                let context = CGContext(data: CVPixelBufferGetBaseAddress(pb),
                    width: cgImage.width, height: cgImage.height,
                    bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(pb),
                    space: CGColorSpaceCreateDeviceRGB(),
                    bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue | CGBitmapInfo.byteOrder32Little.rawValue)!
                context.draw(cgImage, in: CGRect(x: 0, y: 0, width: cgImage.width, height: cgImage.height))

                let handler = VNImageRequestHandler(cvPixelBuffer: pb, options: [:])
                let detectReq = VNDetectFaceRectanglesRequest()
                let lmReq = VNDetectFaceLandmarksRequest()

                do {
                    try handler.perform([detectReq, lmReq])
                } catch {
                    return
                }

                let faceObservations = detectReq.results ?? []
                let landmarkObservations = lmReq.results ?? []
                guard !faceObservations.isEmpty || !landmarkObservations.isEmpty else {
                    return
                }

                let seconds = CMTimeGetSeconds(actualTime)
                let frameNumber = Int(seconds * Double(fps))
                var frameFaces: [[String: Any]] = []

                // Use actual CGImage size (may differ from naturalSize after transform)
                let imgW = CGFloat(cgImage.width)
                let imgH = CGFloat(cgImage.height)

                // Process landmark observations FIRST (each has bbox + landmarks, self-consistent)
                // Quality filtering
                let MIN_CONFIDENCE = 0.6
                let MIN_SIZE = 20
                
                for lmObs in landmarkObservations {
                    // Confidence filter
                    let lmConf = Double(lmObs.confidence)
                    if lmConf < MIN_CONFIDENCE { continue }
                    
                    let bb = lmObs.boundingBox
                    let faceW = Int(bb.size.width * imgW)
                    let faceH = Int(bb.size.height * imgH)
                    
                    // Size filter
                    if faceW < MIN_SIZE || faceH < MIN_SIZE { continue }
                    
                    let faceX = Int(bb.origin.x * imgW)
                    let faceY = Int((1.0 - bb.origin.y - bb.size.height) * imgH)

                    var faceData: [String: Any] = [
                        "bbox": ["x": max(0, faceX), "y": max(0, faceY),
                                 "width": faceW, "height": faceH],
                        "confidence": Double(lmObs.confidence),
                    ]

                    // Pose from landmark observation
                    if let yaw = lmObs.yaw?.doubleValue,
                       let roll = lmObs.roll?.doubleValue {
                        var poseInfo: [String: Any] = ["roll": roll, "yaw": yaw]
                        if let pitch = lmObs.pitch?.doubleValue {
                            poseInfo["pitch"] = pitch
                        }
                        faceData["pose"] = poseInfo
                    }

                    // Landmarks with Y-flip (macOS image coords: bottom-left -> top-left)
                    if let lms = lmObs.landmarks {
                        let imgSize = CGSize(width: imgW, height: imgH)
                        let leftEye = lms.leftEye?.pointsInImage(imageSize: imgSize) ?? []
                        let rightEye = lms.rightEye?.pointsInImage(imageSize: imgSize) ?? []
                        let nose = lms.nose?.pointsInImage(imageSize: imgSize) ?? []

                        if !leftEye.isEmpty || !rightEye.isEmpty || !nose.isEmpty {
                            var lm: [String: [[Double]]] = [:]
                            if !leftEye.isEmpty {
                                lm["left_eye"] = leftEye.map { [Double($0.x), Double(imgH - $0.y)] }
                            }
                            if !rightEye.isEmpty {
                                lm["right_eye"] = rightEye.map { [Double($0.x), Double(imgH - $0.y)] }
                            }
                            if !nose.isEmpty {
                                lm["nose"] = nose.map { [Double($0.x), Double(imgH - $0.y)] }
                            }
                            faceData["landmarks"] = lm
                        }

                        let outer = lms.outerLips?.pointsInImage(imageSize: imgSize) ?? []
                        let inner = lms.innerLips?.pointsInImage(imageSize: imgSize) ?? []
                        if !outer.isEmpty || !inner.isEmpty {
                            faceData["lips"] = [
                                "outer_lips": outer.map { [Double($0.x), Double(imgH - $0.y)] },
                                "inner_lips": inner.map { [Double($0.x), Double(imgH - $0.y)] }
                            ]
                        }
                    }

                    frameFaces.append(faceData)
                }

                // Output face rect observations that the landmark detector missed.
                // Match against ALL landmark observations via IoU to avoid duplicates.
                for faceObs in faceObservations {
                    let fBB = faceObs.boundingBox
                    var matched = false
                    for lmObs in landmarkObservations {
                        let lBB = lmObs.boundingBox
                        let ix = max(fBB.origin.x, lBB.origin.x)
                        let iy = max(fBB.origin.y, lBB.origin.y)
                        let iw = min(fBB.maxX, lBB.maxX) - ix
                        let ih = min(fBB.maxY, lBB.maxY) - iy
                        if iw <= 0 || ih <= 0 { continue }
                        let intersection = iw * ih
                        let union = fBB.width * fBB.height + lBB.width * lBB.height - intersection
                        if intersection / union > 0.3 {
                            matched = true
                            break
                        }
                    }
                    if matched { continue }
                    
                    // Quality filtering for unmatched face rects
                    let MIN_CONFIDENCE = 0.6
                    let MIN_SIZE = 20
                    
                    let faceConf = Double(faceObs.faceCaptureQuality ?? faceObs.confidence)
                    if faceConf < MIN_CONFIDENCE { continue }
                    
                    let faceW = Int(fBB.size.width * imgW)
                    let faceH = Int(fBB.size.height * imgH)
                    if faceW < MIN_SIZE || faceH < MIN_SIZE { continue }
                    
                    // Unmatched face rect: output without landmarks
                    let faceX = Int(fBB.origin.x * imgW)
                    let faceY = Int((1.0 - fBB.origin.y - fBB.size.height) * imgH)

                    var faceData: [String: Any] = [
                        "bbox": ["x": max(0, faceX), "y": max(0, faceY),
                                 "width": faceW, "height": faceH],
                        "confidence": Double(faceObs.faceCaptureQuality ?? faceObs.confidence),
                    ]
                    if let yaw = faceObs.yaw?.doubleValue,
                       let roll = faceObs.roll?.doubleValue {
                        var poseInfo: [String: Any] = ["roll": roll, "yaw": yaw]
                        if let pitch = faceObs.pitch?.doubleValue {
                            poseInfo["pitch"] = pitch
                        }
                        faceData["pose"] = poseInfo
                    }
                    frameFaces.append(faceData)
                }

                if !frameFaces.isEmpty {
                    allFrames.append([
                        "frame": frameNumber,
                        "timestamp": seconds,
                        "faces": frameFaces,
                    ])
                    processedCount += 1
                }
            }

            // Wait for batch to complete
            for _ in batchStart..<batchEnd {
                semaphore.wait()
            }

            let elapsed = Date().timeIntervalSince(startTime)
            let pct = Int(Double(min(batchEnd, totalSamples)) / Double(totalSamples) * 100)
            print("[SwiftFace] \(processedCount) frames with faces, \(pct)% complete, \(Int(elapsed))s elapsed")
            fflush(stdout)
        }

        generator.cancelAllCGImageGeneration()

        let output: [String: Any] = [
            "frame_count": allFrames.count,
            "fps": Double(fps),
            "frames": allFrames,
        ]

        guard let jsonData = try? JSONSerialization.data(withJSONObject: output, options: []),
              let jsonString = String(data: jsonData, encoding: .utf8) else {
            print("[SwiftFace] Failed to serialize JSON")
            return
        }

        let outputURL = URL(fileURLWithPath: outputPath)
        try jsonString.write(to: outputURL, atomically: false, encoding: .utf8)

        let elapsed = Date().timeIntervalSince(startTime)
        print("[SwiftFace] Done: \(allFrames.count) frames, \(String(format: "%.1f", elapsed))s → \(outputPath)")
    }
}
