import Foundation
import AppKit
import CoreGraphics
import CoreMedia
import ScreenCaptureKit
import AVFoundation

class ScreenCapture: NSObject, SCStreamOutput {
    func stream(_ stream: SCStream, didOutputSampleBuffer sampleBuffer: CMSampleBuffer, of type: SCStreamOutputType) {
        guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
            return
        }
        
        DispatchQueue.main.async {
            let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
            let rep = NSCIImageRep(ciImage: ciImage)
            let nsImage = NSImage(size: rep.size)
            nsImage.addRepresentation(rep)
            
            self.lastCapturedImage = nsImage
        }
    }
    
    private var activeStream: SCStream?
    private var lastCapturedImage: NSImage?
    private var excludeCursor: Bool = false
    
    func startCapture(excludeCursor: Bool = false) async throws {
        // Always stop any existing capture before starting a new one
        if activeStream != nil {
            try await stopCapture()
        }
        
        self.excludeCursor = excludeCursor
        
        guard #available(macOS 12.3, *) else { return }
        
        let content = try await SCShareableContent.current
        
        // Get the window from the content
        guard let window = content.windows.first(where: { window in
            let bundleID = window.owningApplication?.bundleIdentifier ?? ""
            let isChessApp = bundleID == "com.chess.iphone"
            let hasValidSize = window.frame.width > 100 && window.frame.height > 100
            return isChessApp && window.isOnScreen && hasValidSize
        }) else {
            throw NSError(domain: "ChessPrism", code: 1, userInfo: [NSLocalizedDescriptionKey: "Chess window not found"])
        }
        
        let filter = SCContentFilter(desktopIndependentWindow: window)
        let config = SCStreamConfiguration()
        config.width = Int(window.frame.width)
        config.height = Int(window.frame.height)
        config.minimumFrameInterval = CMTime(value: 1, timescale: 30) // 30 FPS
        config.queueDepth = 5 // Buffer up to 5 frames
        config.showsCursor = !excludeCursor // Set cursor visibility
        
        let stream = SCStream(filter: filter, configuration: config, delegate: nil)
        try stream.addStreamOutput(self, type: .screen, sampleHandlerQueue: DispatchQueue.global(qos: .userInitiated))
        
        // Store the stream before starting it
        activeStream = stream
        try await stream.startCapture()
    }
    
    func stopCapture() async throws {
        if let stream = activeStream {
            try await stream.stopCapture()
            activeStream = nil
            lastCapturedImage = nil
        }
    }
    
    func getCurrentImage() -> NSImage? {
        return lastCapturedImage
    }
}
