import Foundation
import Vision
import ArgumentParser
import AVFoundation
import AppKit

/// Swift OCR Processor - replaces Python PaddleOCR
/// Uses Apple Vision Framework (VNRecognizeTextRequest) with ANE acceleration
///
/// Output format (compatible with OcrResult Rust struct):
/// {
///   "frame_count": N,
///   "fps": 30.0,
///   "frames": [
///     { "frame": 0, "timestamp": 0.0, "texts": [{ "text": "...", "x": 0, "y": 0, "width": 0, "height": 0, "confidence": 0.0 }] }
///   ]
/// }
@main
struct SwiftOCR: ParsableCommand {
    @Argument(help: "Video file path")
    var inputPath: String

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

    @Option(name: .long, help: "Frames to skip between OCR (default=30)")
    var sampleInterval: Int = 30

    @Option(name: .long, help: "Video FPS (auto-detect if 0)")
    var fps: Double = 0

    @Option(name: .long, help: "UUID for logging")
    var uuid: String = ""

    @Option(name: .long, help: "Recognition level: fast or accurate (default=accurate)")
    var recognitionLevel: String = "accurate"

    mutating func run() throws {
        let startTime = Date()
        print("[SwiftOCR] Starting: \(inputPath)")
        print("[SwiftOCR] Sample interval: \(sampleInterval)")

        let url = URL(fileURLWithPath: inputPath)
        let asset = AVAsset(url: url)

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

        let duration = asset.duration.seconds
        let detectedFps = fps > 0 ? fps : Double(videoTrack.nominalFrameRate)
        let totalFrames = Int(duration * detectedFps)
        print("[SwiftOCR] Duration: \(String(format: "%.1f", duration))s, FPS: \(String(format: "%.1f", detectedFps)), Frames: \(totalFrames)")

        let frameStep = sampleInterval

        // Use shared frame cache if available (set by FrameManager)
        let tempDir: URL
        let framesDir: URL
        if let cacheDir = ProcessInfo.processInfo.environment["MOMENTRY_FRAME_DIR"] {
            framesDir = URL(fileURLWithPath: cacheDir)
            tempDir = framesDir  // No cleanup needed (managed by FrameManager)
            print("[SwiftOCR] Using shared frame cache: \(cacheDir)")
        } else {
            tempDir = FileManager.default.temporaryDirectory.appendingPathComponent("swift_ocr_\(UUID().uuidString)")
            framesDir = tempDir.appendingPathComponent("frames")
            try FileManager.default.createDirectory(at: framesDir, withIntermediateDirectories: true)

            let framePattern = framesDir.appendingPathComponent("frame_%05d.jpg").path
            print("[SwiftOCR] Extracting frames with ffmpeg (interval=\(frameStep))...")

            let extractProc = Process()
            extractProc.executableURL = URL(fileURLWithPath: "/opt/homebrew/bin/ffmpeg")
            extractProc.arguments = ["-y", "-v", "quiet", "-i", inputPath,
                                      "-vf", "select=not(mod(n\\,\(frameStep))),scale=320:-2",
                                      "-vsync", "vfr", "-q:v", "15", framePattern]
            let startExtract = Date()
            try extractProc.run()
            extractProc.waitUntilExit()
            let extractTime = Date().timeIntervalSince(startExtract)
            print("[SwiftOCR] Frame extraction complete: \(String(format: "%.1f", extractTime))s")
        }

        // Sort extracted frame files
        let fileManager = FileManager.default
        let allFiles = (try? fileManager.contentsOfDirectory(atPath: framesDir.path)) ?? []
        let frameFiles = allFiles
            .filter { $0.hasPrefix("frame_") && $0.hasSuffix(".jpg") }
            .sorted()

        let level: VNRequestTextRecognitionLevel = (recognitionLevel == "fast") ? .fast : .accurate
        var ocrFrames: [[String: Any]] = []
        var lastProgress = 0
        let totalFrames_to_process = frameFiles.count

        for (i, frameName) in frameFiles.enumerated() {
            let imgPath = framesDir.appendingPathComponent(frameName).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
            }

            // Extract frame number from filename
            let frameNumber = Int(frameName.replacingOccurrences(of: "frame_", with: "").replacingOccurrences(of: ".jpg", with: "")) ?? (i * frameStep)
            let timestamp = Double(frameNumber) / detectedFps

            let handler = VNImageRequestHandler(cgImage: cgImage, options: [:])
            let request = VNRecognizeTextRequest()
            request.recognitionLevel = level
            request.usesLanguageCorrection = true
            request.preferBackgroundProcessing = true

            guard (try? handler.perform([request])) != nil,
                  let results = request.results else { continue }

            var texts: [[String: Any]] = []
            let cgW = cgImage.width
            let cgH = cgImage.height
            for obs in results {
                guard let candidate = obs.topCandidates(1).first else { continue }
                let conf = candidate.confidence
                guard conf > 0.3 else { continue }
                let bb = obs.boundingBox
                let item: [String: Any] = [
                    "text": candidate.string,
                    "x": Int(bb.origin.x * CGFloat(cgW)),
                    "y": Int((1.0 - bb.origin.y - bb.size.height) * CGFloat(cgH)),
                    "width": Int(bb.size.width * CGFloat(cgW)),
                    "height": Int(bb.size.height * CGFloat(cgH)),
                    "confidence": conf
                ]
                texts.append(item)
            }

            if !texts.isEmpty {
                ocrFrames.append([
                    "frame": frameNumber,
                    "timestamp": timestamp,
                    "texts": texts
                ])
            }

            let pct = Int(Float(i) / Float(totalFrames_to_process) * 100)
            if pct >= lastProgress + 5 {
                print("[SwiftOCR] Progress: \(pct)% (\(i)/\(totalFrames_to_process) samples, \(ocrFrames.count) with text)")
                lastProgress = pct
            }
        }

        // Write output
        let outputDict: [String: Any] = [
            "frame_count": ocrFrames.count,
            "fps": detectedFps,
            "frames": ocrFrames
        ]

        if let jsonData = try? JSONSerialization.data(withJSONObject: outputDict, options: [.prettyPrinted]) {
            try jsonData.write(to: URL(fileURLWithPath: outputPath))
        }

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

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

    func recognizeText(pixelBuffer: CVPixelBuffer, level: VNRequestTextRecognitionLevel) -> [[String: Any]] {
        var texts: [[String: Any]] = []
        let request = VNRecognizeTextRequest()
        request.recognitionLevel = level
        request.usesLanguageCorrection = true
        request.preferBackgroundProcessing = true

        let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer, options: [:])
        guard (try? handler.perform([request])) != nil,
              let results = request.results else { return texts }

        let cgW = CGFloat(CVPixelBufferGetWidth(pixelBuffer))
        let cgH = CGFloat(CVPixelBufferGetHeight(pixelBuffer))

        for obs in results {
            guard let candidate = obs.topCandidates(1).first,
                  candidate.confidence > 0.2 else { continue }
            let bb = obs.boundingBox
            texts.append([
                "text": candidate.string,
                "x": Int(bb.origin.x * cgW),
                "y": Int((1.0 - bb.origin.y - bb.size.height) * cgH),
                "width": Int(bb.size.width * cgW),
                "height": Int(bb.size.height * cgH),
                "confidence": candidate.confidence
            ])
        }
        return texts
    }
}
