import Foundation
import Vision
import ArgumentParser
import AppKit
import AVFoundation

/// Benchmark: Apple Vision Framework body pose detection speed
@main
struct PoseBenchmark: ParsableCommand {
    @Argument(help: "Video path or image directory")
    var inputPath: String

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

    mutating func run() throws {
        let start = Date()
        print("=== Vision Body Pose Benchmark ===")
        #if arch(arm64)
        print("HW: Apple Silicon ✅")
        #endif

        let url = URL(fileURLWithPath: inputPath)
        let asset = AVAsset(url: url)
        guard let track = asset.tracks(withMediaType: .video).first else {
            print("No video track"); return
        }
        let duration = asset.duration.seconds
        let fps = Double(track.nominalFrameRate)
        let totalFrames = Int(duration * fps)

        // Extract frames with ffmpeg
        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent("pose_bench_\(UUID().uuidString)")
        let framesDir = tempDir.appendingPathComponent("frames")
        try FileManager.default.createDirectory(at: framesDir, withIntermediateDirectories: true)

        let pattern = framesDir.appendingPathComponent("frame_%05d.jpg").path
        let extract = Process()
        extract.executableURL = URL(fileURLWithPath: "/opt/homebrew/bin/ffmpeg")
        extract.arguments = ["-y", "-v", "quiet", "-i", inputPath,
                             "-vf", "select=not(mod(n\\,\(sampleInterval)))",
                             "-vsync", "vfr", "-q:v", "5", pattern]
        try extract.run()
        extract.waitUntilExit()

        let files = (try? FileManager.default.contentsOfDirectory(atPath: framesDir.path)) ?? []
        let frameFiles = files.filter { $0.hasSuffix(".jpg") }.sorted()
        print("Frames: \(frameFiles.count)")

        // Process all frames in one loop (no subprocess overhead)
        var totalPoses = 0
        var framesWithPose = 0
        let inferenceStart = Date()

        for fname in frameFiles {
            let imgPath = framesDir.appendingPathComponent(fname).path
            guard let imgData = try? Data(contentsOf: URL(fileURLWithPath: imgPath)),
                  let img = NSImage(data: imgData),
                  let cgImage = img.cgImage(forProposedRect: nil, context: nil, hints: nil) else { continue }

            let handler = VNImageRequestHandler(cgImage: cgImage, options: [:])
            let req = VNDetectHumanBodyPoseRequest()
            try? handler.perform([req])

            if let poses = req.results, !poses.isEmpty {
                framesWithPose += 1
                totalPoses += poses.count
            }
        }

        let inferenceTime = Date().timeIntervalSince(inferenceStart)
        let totalTime = Date().timeIntervalSince(start)

        print("\n=== Results ===")
        print("Frames: \(frameFiles.count), with poses: \(framesWithPose)")
        print("Total poses: \(totalPoses)")
        print("Inference: \(String(format: "%.2f", inferenceTime))s")
        print("Per frame: \(String(format: "%.0f", inferenceTime / Double(frameFiles.count) * 1000))ms")
        print("Total: \(String(format: "%.2f", totalTime))s")

        try FileManager.default.removeItem(at: tempDir)
    }
}
