import Foundation
import Vision
import ArgumentParser
import AppKit
import AVFoundation

/// Full comparison: Apple Vision Framework vs InsightFace for face processing
@main
struct FaceCompareTest: ParsableCommand {
    @Argument(help: "Video path or image path")
    var inputPath: String

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

    @Option(name: .long, help: "Maximum frames to process")
    var maxFrames: Int = 20

    mutating func run() throws {
        let startTime = Date()
        print("=== Apple Vision Framework Face Processing ===")
        #if arch(arm64)
        print("HW: Apple Silicon ✅")
        #endif

        let url = URL(fileURLWithPath: inputPath)
        let ext = (inputPath as NSString).pathExtension.lowercased()

        if ["mp4", "mov", "mkv", "avi"].contains(ext) {
            try processVideo(url: url)
        } else {
            try processImage(url: url)
        }

        print("Time: \(String(format: "%.2f", Date().timeIntervalSince(startTime)))s")
    }

    func processVideo(url: URL) throws {
        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)
        print("Video: \(duration)s @ \(fps)fps = \(totalFrames) frames")

        // Extract frames with ffmpeg at sample interval
        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent("face_compare_\(UUID().uuidString)")
        let framesDir = tempDir.appendingPathComponent("frames")
        try FileManager.default.createDirectory(at: framesDir, withIntermediateDirectories: true)
        defer { try? FileManager.default.removeItem(at: tempDir) }

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

        let allFiles = (try? FileManager.default.contentsOfDirectory(atPath: framesDir.path)) ?? []
        let frameFiles = allFiles.filter { $0.hasSuffix(".jpg") }.sorted().prefix(maxFrames)

        var totalFaces = 0
        var framesWithFaces = 0
        var frameCount = 0

        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 frameNum = Int(fname.replacingOccurrences(of: "frame_", with: "").replacingOccurrences(of: ".jpg", with: "")) ?? 0
            let timestamp = Double(frameNum) / fps

            // Run all face detection requests
            let faceResult = detectFaces(cgImage: cgImage)
            if faceResult.count > 0 {
                totalFaces += faceResult.count
                framesWithFaces += 1
                print("  Frame \(frameNum) (\(String(format: "%.1f", timestamp))s): \(faceResult.count) faces")
                for (i, f) in faceResult.enumerated() {
                    print("    [\(i)] bbox=(\(String(format: "%.0f", f.x)),\(String(format: "%.0f", f.y))) size=\(String(format: "%.0f", f.w))x\(String(format: "%.0f", f.h)) conf=\(String(format: "%.3f", f.conf)) quality=\(String(format: "%.3f", f.quality)) landmarks=\(f.landmarks) embedding=\(f.hasEmbedding ? "✅" : "❌")")
                }
            }
            frameCount += 1
        }

        print("\n=== Summary ===")
        print("Frames processed: \(frameCount)")
        print("Frames with faces: \(framesWithFaces)")
        print("Total faces detected: \(totalFaces)")

        // Compare with existing InsightFace JSON if available
        let uuid = extractUUID(from: url.lastPathComponent)
        if uuid != "" {
            let faceJsonPath = "/Users/accusys/momentry/output_dev/\(uuid).face.json"
            if FileManager.default.fileExists(atPath: faceJsonPath) {
                print("\n=== Comparison with InsightFace (\(uuid).face.json) ===")
                if let data = try? Data(contentsOf: URL(fileURLWithPath: faceJsonPath)),
                   let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
                   let frames = json["frames"] as? [[String: Any]] {
                    let insightFaces = frames.filter { ($0["faces"] as? [Any])?.count ?? 0 > 0 }.count
                    let totalInsightFaces = frames.reduce(0) { $0 + (($1["faces"] as? [Any])?.count ?? 0) }
                    print("  InsightFace frames with faces: \(insightFaces)")
                    print("  InsightFace total faces: \(totalInsightFaces)")
                }
            }
        }
    }

    func processImage(url: URL) throws {
        guard let imgData = try? Data(contentsOf: url),
              let img = NSImage(data: imgData),
              let cgImage = img.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
            print("Cannot load image"); return
        }
        print("Image: \(cgImage.width)x\(cgImage.height)")
        let result = detectFaces(cgImage: cgImage)
        print("Vision faces: \(result.count)")
        for (i, f) in result.enumerated() {
            print("  [\(i)] bbox=(\(String(format: "%.0f", f.x)),\(String(format: "%.0f", f.y))) size=\(String(format: "%.0f", f.w))x\(String(format: "%.0f", f.h)) conf=\(String(format: "%.3f", f.conf)) quality=\(String(format: "%.3f", f.quality)) landmarks=\(f.landmarks) embedding=\(f.hasEmbedding ? "✅" : "❌")")
        }
    }

    struct FaceResult {
        let x, y, w, h: Float
        let conf: Float
        let quality: Float
        let landmarks: Int
        let hasEmbedding: Bool
    }

    func detectFaces(cgImage: CGImage) -> [FaceResult] {
        var results: [FaceResult] = []
        let handler = VNImageRequestHandler(cgImage: cgImage, options: [:])

        // 1. Face Detection
        let detectReq = VNDetectFaceRectanglesRequest()
        try? handler.perform([detectReq])

        // 2. Face Landmarks (run separately for each detected face)
        let landmarkReq = VNDetectFaceLandmarksRequest()

        // 3. Face Capture Quality
        let qualityReq = VNDetectFaceCaptureQualityRequest()

        // Run all requests
        try? handler.perform([landmarkReq, qualityReq])

        guard let detections = detectReq.results else { return [] }

        let qualityResults = qualityReq.results ?? []

        for (i, face) in detections.enumerated() {
            let bb = face.boundingBox
            let w = Float(cgImage.width)
            let h = Float(cgImage.height)
            let x = Float(bb.origin.x) * w
            let y = Float(bb.origin.y) * h
            let fw = Float(bb.size.width) * w
            let fh = Float(bb.size.height) * h

            // Get landmarks count
            var lmCount = 0
            if let lmResults = landmarkReq.results, i < lmResults.count {
                let lms = lmResults[i].landmarks
                if let left = lms?.leftEye { lmCount += left.pointCount }
                if let right = lms?.rightEye { lmCount += right.pointCount }
                if let nose = lms?.nose { lmCount += nose.pointCount }
            }

            // Get quality score
            var quality: Float = 0
            if i < qualityResults.count {
                if #available(macOS 14, *) {
                    quality = (qualityResults[i].value(forKey: "faceCaptureQuality") as? Double).map { Float($0) } ?? 0
                }
            }

            // Check for faceprint (embedding) via KVC
            var hasEmbedding = false
            if #available(macOS 14, *) {
                if let fp = face.value(forKey: "faceprint") as? NSObject {
                    hasEmbedding = (fp.value(forKey: "data") as? Data) != nil
                }
            }

            results.append(FaceResult(x: x, y: y, w: fw, h: fh, conf: face.confidence, quality: quality, landmarks: lmCount, hasEmbedding: hasEmbedding))
        }

        return results
    }

    func extractUUID(from filename: String) -> String {
        // Try to extract 32-char hex UUID from filename
        let pattern = try? NSRegularExpression(pattern: "[a-f0-9]{32}")
        if let match = pattern?.firstMatch(in: filename, range: NSRange(location: 0, length: filename.count)) {
            return (filename as NSString).substring(with: match.range)
        }
        return ""
    }
}
