import Foundation
import Vision
import ArgumentParser
import AppKit
import AVFoundation

/// Swift Pose Processor - replaces YOLOv8 Pose / MediaPipe Pose
/// Uses VNDetectHumanBodyPoseRequest with ANE acceleration
///
/// Output format (compatible with PoseResult Rust struct):
/// {
///   "frame_count": N, "fps": 30.0,
///   "frames": [
///     { "frame": 0, "timestamp": 0.0, "persons": [
///       { "keypoints": [{"name":"nose","x":100,"y":200,"confidence":0.95}],
///         "bbox": {"x":0,"y":0,"width":100,"height":200}
///       }
///     ]}
///   ]
/// }
@main
struct SwiftPose: 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("[SwiftPose] Starting: \(inputPath)")

        let url = URL(fileURLWithPath: inputPath)
        let asset = AVAsset(url: url)
        guard let track = asset.tracks(withMediaType: .video).first else {
            print("[SwiftPose] Error: No video track"); return
        }
        let duration = asset.duration.seconds
        let fps = Double(track.nominalFrameRate)
        print("[SwiftPose] Duration: \(String(format: "%.1f", duration))s, FPS: \(String(format: "%.1f", fps))")

        // Extract frames (use shared cache or ffmpeg)
        let tempDir: URL
        let framesDir: URL
        if let cacheDir = ProcessInfo.processInfo.environment["MOMENTRY_FRAME_DIR"] {
            framesDir = URL(fileURLWithPath: cacheDir)
            tempDir = framesDir
            print("[SwiftPose] Using shared frame cache: \(cacheDir)")
        } else {
            tempDir = FileManager.default.temporaryDirectory.appendingPathComponent("swift_pose_\(UUID().uuidString)")
            framesDir = tempDir.appendingPathComponent("frames")
            try FileManager.default.createDirectory(at: framesDir, withIntermediateDirectories: true)

            let pattern = framesDir.appendingPathComponent("frame_%05d.jpg").path
            print("[SwiftPose] Extracting frames...")
            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", "15", pattern]
            try extract.run()
            extract.waitUntilExit()
        }

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

        let jointNames: [VNHumanBodyPoseObservation.JointName] = [
            .nose, .leftEye, .rightEye, .leftEar, .rightEar,
            .neck, .root,
            .leftShoulder, .rightShoulder,
            .leftElbow, .rightElbow,
            .leftWrist, .rightWrist,
            .leftHip, .rightHip,
            .leftKnee, .rightKnee,
            .leftAnkle, .rightAnkle,
        ]

        var poseFrames: [[String: Any]] = []
        var lastProgress = 0

        for (i, fname) in frameFiles.enumerated() {
            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 frameNum = Int(fname.replacingOccurrences(of: "frame_", with: "").replacingOccurrences(of: ".jpg", with: "")) ?? (i * sampleInterval)
            let timestamp = Double(frameNum) / fps
            let w = cgImage.width
            let h = cgImage.height

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

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

            var persons: [[String: Any]] = []
            for pose in poses {
                var keypoints: [[String: Any]] = []
                var minX = CGFloat.greatestFiniteMagnitude
                var minY = CGFloat.greatestFiniteMagnitude
                var maxX: CGFloat = 0
                var maxY: CGFloat = 0

                for joint in jointNames {
                    if let point = try? pose.recognizedPoint(joint) {
                        let desc = String(describing: joint.rawValue)
                        var rawName = desc
                            .replacingOccurrences(of: "VNRecognizedPointKey(_rawValue: ", with: "")
                            .replacingOccurrences(of: ")", with: "")
                            .trimmingCharacters(in: .whitespaces)
                        // Map Vision Framework joint names to standard names
                        let nameMap: [String: String] = [
                            "head_joint": "nose",
                            "left_eye_joint": "left_eye",
                            "right_eye_joint": "right_eye",
                            "left_ear_joint": "left_ear",
                            "right_ear_joint": "right_ear",
                            "neck_1_joint": "neck",
                            "left_shoulder_1_joint": "left_shoulder",
                            "right_shoulder_1_joint": "right_shoulder",
                            "left_elbow_1_joint": "left_elbow",
                            "right_elbow_1_joint": "right_elbow",
                            "left_hand_joint": "left_wrist",
                            "right_hand_joint": "right_wrist",
                            "left_hip_1_joint": "left_hip",
                            "right_hip_1_joint": "right_hip",
                            "left_knee_1_joint": "left_knee",
                            "right_knee_1_joint": "right_knee",
                            "left_ankle_1_joint": "left_ankle",
                            "right_ankle_1_joint": "right_ankle",
                            "center_hip_joint": "root",
                            "left_forearm_joint": "left_elbow",
                            "right_forearm_joint": "right_elbow",
                            "left_upLeg_joint": "left_hip",
                            "right_upLeg_joint": "right_hip",
                            "left_leg_joint": "left_knee",
                            "right_leg_joint": "right_knee",
                            "left_foot_joint": "left_ankle",
                            "right_foot_joint": "right_ankle",
                        ]
                        if let mapped = nameMap[rawName] {
                            rawName = mapped
                        }
                        let px = point.location.x * CGFloat(w)
                        let py = CGFloat(h) - point.location.y * CGFloat(h)
                        keypoints.append([
                            "name": rawName.isEmpty ? "\(joint)" : rawName,
                            "x": px,
                            "y": py,
                            "confidence": point.confidence,
                        ])
                        if point.confidence > 0.1 {
                            minX = min(minX, px)
                            minY = min(minY, py)
                            maxX = max(maxX, px)
                            maxY = max(maxY, py)
                        }
                    }
                }

                var bbox: [String: Any] = [
                    "x": 0, "y": 0, "width": 0, "height": 0
                ]
                if maxX > minX {
                    bbox = [
                        "x": Int(minX),
                        "y": Int(minY),
                        "width": Int(maxX - minX),
                        "height": Int(maxY - minY),
                    ]
                }

                persons.append(["keypoints": keypoints, "bbox": bbox])
            }

            if !persons.isEmpty {
                poseFrames.append([
                    "frame": frameNum,
                    "timestamp": timestamp,
                    "persons": persons,
                ])
            }

            let pct = Int(Float(i) / Float(frameFiles.count) * 100)
            if pct >= lastProgress + 10 {
                print("[SwiftPose] Progress: \(pct)% (\(i)/\(frameFiles.count), \(persons.count) poses)")
                lastProgress = pct
            }
        }

        // Write output
        let outputDict: [String: Any] = [
            "frame_count": poseFrames.count,
            "fps": fps,
            "frames": poseFrames,
        ]
        if let jsonData = try? JSONSerialization.data(withJSONObject: outputDict, options: [.prettyPrinted]) {
            try jsonData.write(to: URL(fileURLWithPath: outputPath))
        }

        let elapsed = Date().timeIntervalSince(startTime)
        print("[SwiftPose] Complete: \(poseFrames.count) frames, \(String(format: "%.1f", elapsed))s")
        if duration > 0 {
            print("[SwiftPose] RTF: \(String(format: "%.3f", elapsed / duration))x")
        }

        // Clean up temp dir if we created it
        if ProcessInfo.processInfo.environment["MOMENTRY_FRAME_DIR"] == nil {
            try? FileManager.default.removeItem(at: tempDir)
        }
    }
}
