import Foundation
import AVFoundation
import ArgumentParser
import Accelerate

/// POC: Swift-based scene cut detection using AVFoundation histogram analysis
/// Compared against Python PySceneDetect ContentDetector (threshold=27)
@main
struct SwiftCutTest: ParsableCommand {
    @Argument(help: "Video file path")
    var inputPath: String

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

    @Option(name: .long, help: "Detection threshold (higher= fewer cuts, default 0.3)")
    var threshold: Double = 0.3

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

    mutating func run() throws {
        let startTime = Date()
        print("=== Swift Scene Cut Detection POC ===")
        #if arch(arm64)
        print("HW: Apple Silicon ✅ (ANE available)")
        #endif

        let url = URL(fileURLWithPath: inputPath)
        let asset = AVAsset(url: url)
        
        guard let videoTrack = asset.tracks(withMediaType: .video).first else {
            print("Error: No video track found"); return
        }
        
        let duration = asset.duration.seconds
        let fps = videoTrack.nominalFrameRate
        let totalFrames = Int(duration * Double(fps))
        print("Video: \(inputPath)")
        print("Duration: \(String(format: "%.1f", duration))s")
        print("FPS: \(String(format: "%.1f", fps))")
        print("Total frames: \(totalFrames)")
        print("Threshold: \(String(format: "%.2f", threshold))")
        print("Sample interval: \(sampleInterval)")

        // Read frame histogram data using AVAssetReader
        guard let reader = try? AVAssetReader(asset: asset) else {
            print("Error: Cannot create asset reader"); return
        }

        let settings: [String: Any] = [
            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
            kCVPixelBufferWidthKey as String: 320,  // downscale for speed
            kCVPixelBufferHeightKey as String: 180,
        ]
        let trackOutput = AVAssetReaderTrackOutput(track: videoTrack, outputSettings: settings)
        reader.add(trackOutput)
        reader.startReading()

        var frameIndex = 0
        var prevHistogram: [Float]?
        var scenes: [(start: Double, end: Double)] = []
        var sceneStart: Double = 0
        var diffs: [(frame: Int, diff: Float)] = []

        let frameStep = sampleInterval
        var lastPrint = 0

        while reader.status == .reading {
            guard let sampleBuffer = trackOutput.copyNextSampleBuffer() else { break }
            guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
                CMSampleBufferInvalidate(sampleBuffer); continue
            }

            if frameIndex % frameStep == 0 {
                let timestamp = CMTimeGetSeconds(CMSampleBufferGetPresentationTimeStamp(sampleBuffer))
                
                // Compute histogram
                let histogram = computeLuminanceHistogram(pixelBuffer: pixelBuffer)
                
                if let prev = prevHistogram {
                    let diff = histogramDifference(prev, histogram)
                    
                    if diff > Float(threshold) {
                        // Scene cut detected
                        let sceneEnd = timestamp
                        scenes.append((start: sceneStart, end: sceneEnd))
                        sceneStart = timestamp
                        diffs.append((frame: frameIndex, diff: diff))
                        
                        if scenes.count % 50 == 0 {
                            print("  Scenes so far: \(scenes.count) at frame \(frameIndex)/\(totalFrames)")
                        }
                    }
                }
                prevHistogram = histogram
            }

            frameIndex += 1
            CMSampleBufferInvalidate(sampleBuffer)

            // Progress every 5%
            let pct = Int(Float(frameIndex) / Float(totalFrames) * 100)
            if pct >= lastPrint + 5 {
                print("  Progress: \(pct)% (\(frameIndex)/\(totalFrames) frames)")
                lastPrint = pct
            }
        }

        // Add last scene
        if sceneStart < duration {
            scenes.append((start: sceneStart, end: duration))
        }

        let elapsed = Date().timeIntervalSince(startTime)
        print("\n=== Results ===")
        print("Scenes detected: \(scenes.count)")
        print("Time: \(String(format: "%.2f", elapsed))s")
        if totalFrames > 0 {
            let rtf = elapsed / duration
            print("RTF: \(String(format: "%.3f", rtf))x")
        }
        print("Last 5 cuts:")
        for s in scenes.suffix(5) {
            print("  \(String(format: "%.1f", s.start))s - \(String(format: "%.1f", s.end))s")
        }

        // Output JSON if requested
        if let outPath = outputPath {
            let outputDict: [String: Any] = [
                "scenes": scenes.map { ["start_time": $0.start, "end_time": $0.end] },
                "metadata": [
                    "video_path": inputPath,
                    "duration": duration,
                    "fps": fps,
                    "total_frames": totalFrames,
                    "threshold": threshold,
                    "sample_interval": sampleInterval,
                    "processing_time": elapsed,
                    "rtf": elapsed / duration,
                ],
                "diffs": diffs.map { ["frame": $0.frame, "diff": String(format: "%.4f", $0.diff)] }
            ]
            if let jsonData = try? JSONSerialization.data(withJSONObject: outputDict, options: [.prettyPrinted]) {
                try jsonData.write(to: URL(fileURLWithPath: outPath))
                print("Output written to: \(outPath)")
            }
        }
    }

    func computeLuminanceHistogram(pixelBuffer: CVPixelBuffer) -> [Float] {
        CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
        defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly) }

        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)
        let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)
        let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer)!
        
        var histogram = [Float](repeating: 0, count: 64)  // 64 bins for speed
        
        for y in 0..<height {
            let row = baseAddress.advanced(by: y * bytesPerRow)
            let pixels = row.assumingMemoryBound(to: UInt8.self)
            for x in 0..<width {
                let b = pixels[x * 4]
                let g = pixels[x * 4 + 1]
                let r = pixels[x * 4 + 2]
                let luminance = (0.299 * Float(r) + 0.587 * Float(g) + 0.114 * Float(b))
                let bin = min(Int(luminance / 256.0 * 64.0), 63)
                histogram[bin] += 1
            }
        }
        
        // Normalize
        let total = Float(width * height)
        for i in 0..<histogram.count {
            histogram[i] /= total
        }
        
        return histogram
    }

    func histogramDifference(_ a: [Float], _ b: [Float]) -> Float {
        var diff: Float = 0
        for i in 0..<min(a.count, b.count) {
            diff += abs(a[i] - b[i])
        }
        return diff
    }
}
