import CoreGraphics
import CoreMedia
import CoreVideo
import Foundation
import MastermindPOCCore
import ScreenCaptureKit

final class CaptureCoordinator: NSObject, SCStreamOutput, SCStreamDelegate {
    private let outputQueue = DispatchQueue(label: "app.mastermind.poc.screencapture")
    private let onScreenFrame: (Int) -> Void
    private let onSystemAudioFrame: (PCMFrame) -> Void
    private let onError: (String) -> Void

    private var stream: SCStream?
    private var screenFrameCount = 0

    init(
        onScreenFrame: @escaping (Int) -> Void,
        onSystemAudioFrame: @escaping (PCMFrame) -> Void,
        onError: @escaping (String) -> Void
    ) {
        self.onScreenFrame = onScreenFrame
        self.onSystemAudioFrame = onSystemAudioFrame
        self.onError = onError
        super.init()
    }

    func start(excludingWindowIDs: Set<CGWindowID>) async throws {
        await stop()

        let content = try await SCShareableContent.current

        guard let display = content.displays.first(where: { $0.displayID == CGMainDisplayID() }) ?? content.displays.first else {
            throw POCError("No capturable display found")
        }

        let currentPID = ProcessInfo.processInfo.processIdentifier
        let excludedWindows = content.windows.filter { window in
            excludingWindowIDs.contains(window.windowID) || window.owningApplication?.processID == currentPID
        }

        let filter = SCContentFilter(display: display, excludingWindows: excludedWindows)
        let configuration = SCStreamConfiguration()
        configuration.width = max(1, display.width)
        configuration.height = max(1, display.height)
        configuration.minimumFrameInterval = CMTime(value: 1, timescale: 2)
        configuration.pixelFormat = kCVPixelFormatType_32BGRA
        configuration.queueDepth = 3
        configuration.capturesAudio = true
        configuration.sampleRate = 16_000
        configuration.channelCount = 1
        configuration.excludesCurrentProcessAudio = true

        let stream = SCStream(filter: filter, configuration: configuration, delegate: self)
        try stream.addStreamOutput(self, type: .screen, sampleHandlerQueue: outputQueue)
        try stream.addStreamOutput(self, type: .audio, sampleHandlerQueue: outputQueue)
        try await stream.startCapture()

        screenFrameCount = 0
        self.stream = stream
    }

    func stop() async {
        guard let stream else {
            return
        }

        do {
            try await stream.stopCapture()
        } catch {
            onError("Screen capture stop failed: \(error.localizedDescription)")
        }

        self.stream = nil
    }

    func stream(_ stream: SCStream, didOutputSampleBuffer sampleBuffer: CMSampleBuffer, of type: SCStreamOutputType) {
        guard CMSampleBufferIsValid(sampleBuffer) else {
            return
        }

        switch type {
        case .screen:
            screenFrameCount += 1
            onScreenFrame(screenFrameCount)
        case .audio:
            if let frame = SampleBufferPCMExtractor.extractFrame(from: sampleBuffer, source: .system) {
                onSystemAudioFrame(frame)
            }
        case .microphone:
            return
        @unknown default:
            return
        }
    }

    func stream(_ stream: SCStream, didStopWithError error: Error) {
        onError("Screen capture stopped: \(error.localizedDescription)")
    }
}
