#!/usr/bin/env swift
import Foundation
import Vision
import AVFoundation
import ArgumentParser

/// Full-movie body pose scanner: compute head-to-body ratio for every frame
/// with face detections. Outputs JSONL (one object per frame).
@main
struct BodyPoseScanner: ParsableCommand {
    @Argument(help: "Video file path")
    var videoPath: String

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

    @Option(help: "Frames to scan (comma-separated, e.g. '840,900,960') or 'all' to scan everything")
    var frames: String = "all"

    @Option(help: "Sample interval (every N frames, for 'all' mode)")
    var interval: Int = 60

    func run() throws {
        let url = URL(fileURLWithPath: videoPath)
        let asset = AVAsset(url: url)
        guard let reader = try? AVAssetReader(asset: asset) else {
            print("[BodyPose] Cannot open video"); return
        }

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

        let fps = videoTrack.nominalFrameRate
        let totalFrames = Int(videoTrack.timeRange.duration.seconds * Double(fps))

        let readerOutput = AVAssetReaderTrackOutput(track: videoTrack, outputSettings: [
            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange
        ])
        readerOutput.alwaysCopiesSampleData = false
        reader.add(readerOutput)
        reader.startReading()

        // Parse target frames
        var targetFrames = Set<Int>()
        if frames == "all" {
            targetFrames = Set(stride(from: 0, to: totalFrames, by: interval))
        } else {
            targetFrames = Set(frames.split(separator: ",").compactMap { Int($0.trimmingCharacters(in: .whitespaces)) })
        }

        var frameCount = 0
        var results: [[String: Any]] = []
        let bodyRequest = VNDetectHumanBodyPoseRequest()

        guard let fh = FileHandle(forWritingAtPath: outputPath) else {
            print("[BodyPose] Cannot create output"); return
        }

        while let sampleBuffer = readerOutput.copyNextSampleBuffer() {
            defer { frameCount += 1 }
            guard targetFrames.contains(frameCount) else { continue }
            targetFrames.remove(frameCount)

            guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else { continue }

            let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer, options: [:])
            try? handler.perform([bodyRequest])

            guard let poses = bodyRequest.results, !poses.isEmpty else { continue }

            let imgW = CGFloat(CVPixelBufferGetWidth(pixelBuffer))
            let imgH = CGFloat(CVPixelBufferGetHeight(pixelBuffer))

            for obs in poses {
                guard let pts = try? obs.recognizedPoints(.all) else { continue }

                var joints: [String: CGFloat] = [:]
                for (name, pt) in pts where pt.confidence > 0.3 {
                    // Convert Vision (bottom-left origin) to pixel (top-left origin)
                    joints[String(describing: name)] = pt.location.x * imgW
                    joints[String(describing: name) + "_y"] = imgH - pt.location.y * imgH
                }

                // Get head top
                let headY = joints["head_joint_y"] ?? joints["neck_1_joint_y"] ?? 
                            joints["neck_2_joint_y"] ?? joints["right_eye_joint_y"] ?? 0

                // Get lowest visible body point (foot > ankle > knee > hip)
                var bodyBottom = CGFloat.greatestFiniteMagnitude
                for jn in ["right_ankle_joint_y", "left_ankle_joint_y",
                           "right_knee_joint_y", "left_knee_joint_y",
                           "right_hip_joint_y", "left_hip_joint_y"] {
                    if let v = joints[jn], v > 0 {
                        bodyBottom = min(bodyBottom, v)
                    }
                }

                let bodyH = bodyBottom == .greatestFiniteMagnitude ? 0 : abs(headY - bodyBottom)
                let headH = abs(headY - (joints["neck_1_joint_y"] ?? headY))

                let h2b = bodyH > 0 ? headH / bodyH : 0

                let row: [String: Any] = [
                    "frame": frameCount, "timestamp": Double(frameCount)/Double(fps),
                    "head_top_y": headY, "body_bottom_y": bodyBottom,
                    "body_h_px": bodyH, "head_h_px": headH,
                    "h2b_ratio": Double(String(format: "%.3f", h2b)) ?? 0,
                    "has_full_body": bodyH > 0 && headH > 0,
                    "joints": joints.mapValues { Double($0) }
                ]

                var jsonData = try! JSONSerialization.data(withJSONObject: row)
                jsonData.append(10) // newline
                fh.write(jsonData)
            }

            if targetFrames.isEmpty { break }
        }
        reader.cancelReading()
        fh.closeFile()
        print("[BodyPose] Done: \(results.count) frames → \(outputPath)")
    }
}
