import AppKit
import CoreImage
import CoreMedia
import CoreGraphics
import Foundation
@preconcurrency import ScreenCaptureKit

public enum ScreenshotProviderError: Error, Equatable, Sendable {
    case windowNotFound(UInt32)
    case captureTimedOut(windowID: UInt32, timeoutMilliseconds: Int)
    case invalidCapture(windowID: UInt32, pixelWidth: Int, pixelHeight: Int, pointWidth: Double, pointHeight: Double)
}

public struct WindowScreenshotCapture: Sendable {
    public let image: CGImage
    public let pointSize: NSSize

    public init(image: CGImage, pointSize: NSSize) {
        self.image = image
        self.pointSize = pointSize
    }
}

public enum WindowServerWindowScreenshotCapturer {
    public static func capture(windowID: UInt32) -> WindowScreenshotCapture? {
        guard let image = captureImage(windowID: windowID) else {
            return nil
        }

        return WindowScreenshotCapture(
            image: image,
            pointSize: NSSize(width: image.width, height: image.height)
        )
    }

    public static func captureImage(windowID: UInt32) -> CGImage? {
        var mutableWindowID = CGWindowID(windowID)
        let images = CGSHWCaptureWindowList(
            CGSMainConnectionID(),
            &mutableWindowID,
            1,
            [.ignoreGlobalClipShape, .bestResolution, .fullSize]
        ).takeRetainedValue() as NSArray

        guard let firstObject = images.firstObject else {
            return nil
        }

        let cfObject = firstObject as CFTypeRef
        guard CFGetTypeID(cfObject) == CGImage.typeID else {
            return nil
        }

        return (firstObject as! CGImage)
    }
}

@MainActor
public protocol WindowScreenshotCapturing: Sendable {
    func captureWindow(windowID: UInt32) async throws -> WindowScreenshotCapture
}

public struct ScreenshotCapturePolicy: Equatable, Sendable {
    public static let round01Default = ScreenshotCapturePolicy(
        timeoutMilliseconds: 300,
        maxRetriesPerSession: 1
    )

    public let timeoutMilliseconds: Int
    public let maxRetriesPerSession: Int

    public init(timeoutMilliseconds: Int, maxRetriesPerSession: Int) {
        self.timeoutMilliseconds = max(1, timeoutMilliseconds)
        self.maxRetriesPerSession = max(0, maxRetriesPerSession)
    }

    var timeoutNanoseconds: UInt64 {
        UInt64(timeoutMilliseconds) * 1_000_000
    }
}

public enum ScreenshotFallbackReason: Equatable, Sendable {
    case screenRecordingDenied
    case captureFailed(String)
    case timedOut(windowID: UInt32, timeoutMilliseconds: Int)
    case invalidCapture(windowID: UInt32, pixelWidth: Int, pixelHeight: Int, pointWidth: Double, pointHeight: Double)
    case retryLimitReached(failedAttempts: Int)
    case syntheticWindowID(windowID: UInt32)
}

public enum ScreenshotResolutionSource: Equatable, Sendable {
    case realScreenshot
    case skeletonFallback(ScreenshotFallbackReason)
}

@MainActor
public struct ScreenshotResolution {
    public let image: NSImage
    public let source: ScreenshotResolutionSource

    public init(image: NSImage, source: ScreenshotResolutionSource) {
        self.image = image
        self.source = source
    }
}

public enum ScreenshotDebugEvent: Equatable, Sendable {
    case captureStarted(windowID: UInt32)
    case captureSucceeded(windowID: UInt32)
    case captureFailed(windowID: UInt32, reason: ScreenshotFallbackReason)
    case retrySkipped(windowID: UInt32, failedAttempts: Int)
    case fallbackUsed(windowID: UInt32, reason: ScreenshotFallbackReason)
}

public extension ScreenshotFallbackReason {
    var diagnosticDescription: String {
        switch self {
        case .screenRecordingDenied:
            return "screenRecordingDenied"
        case .captureFailed(let summary):
            return "captureFailed(\(summary))"
        case .timedOut(let windowID, let timeoutMilliseconds):
            return "timedOut(windowID: \(windowID), timeoutMilliseconds: \(timeoutMilliseconds))"
        case .invalidCapture(let windowID, let pixelWidth, let pixelHeight, let pointWidth, let pointHeight):
            return "invalidCapture(windowID: \(windowID), pixelSize: \(pixelWidth)x\(pixelHeight), pointSize: \(pointWidth)x\(pointHeight))"
        case .retryLimitReached(let failedAttempts):
            return "retryLimitReached(failedAttempts: \(failedAttempts))"
        case .syntheticWindowID(let windowID):
            return "syntheticWindowID(windowID: \(windowID))"
        }
    }
}

public extension ScreenshotDebugEvent {
    var diagnosticDescription: String {
        switch self {
        case .captureStarted(let windowID):
            return "captureStarted(\(windowID))"
        case .captureSucceeded(let windowID):
            return "captureSucceeded(\(windowID))"
        case .captureFailed(let windowID, let reason):
            return "captureFailed(\(windowID), \(reason.diagnosticDescription))"
        case .retrySkipped(let windowID, let failedAttempts):
            return "retrySkipped(\(windowID), failedAttempts: \(failedAttempts))"
        case .fallbackUsed(let windowID, let reason):
            return "fallbackUsed(\(windowID), \(reason.diagnosticDescription))"
        }
    }
}

@MainActor
public protocol ScreenshotDebugLogging: AnyObject {
    func record(_ event: ScreenshotDebugEvent)
}

@MainActor
public final class NoopScreenshotDebugLogger: ScreenshotDebugLogging {
    public init() {}

    public func record(_ event: ScreenshotDebugEvent) {}
}

@MainActor
public final class ScreenshotCaptureSession {
    private var failedAttemptsByWindowID: [UInt32: Int] = [:]

    public init() {}

    func failedAttempts(for windowID: UInt32) -> Int {
        failedAttemptsByWindowID[windowID, default: 0]
    }

    func recordFailure(for windowID: UInt32) {
        failedAttemptsByWindowID[windowID, default: 0] += 1
    }

    func recordSuccess(for windowID: UInt32) {
        failedAttemptsByWindowID[windowID] = nil
    }
}

@MainActor
private final class ScreenshotCaptureRace {
    private var continuation: CheckedContinuation<WindowScreenshotCapture, any Error>?
    private var result: Result<WindowScreenshotCapture, any Error>?
    private var didFinish = false

    func wait() async throws -> WindowScreenshotCapture {
        if let result {
            return try result.get()
        }

        return try await withCheckedThrowingContinuation { continuation in
            self.continuation = continuation
        }
    }

    func finish(_ result: Result<WindowScreenshotCapture, any Error>) {
        guard !didFinish else {
            return
        }

        didFinish = true
        if let continuation {
            self.continuation = nil
            continuation.resume(with: result)
        } else {
            self.result = result
        }
    }
}

@MainActor
public final class ScreenCaptureKitWindowScreenshotCapturer: WindowScreenshotCapturing {
    public init() {}

    public func captureWindow(windowID: UInt32) async throws -> WindowScreenshotCapture {
        if let privateCapture = WindowServerWindowScreenshotCapturer.capture(windowID: windowID) {
            return privateCapture
        }

        let content = try await SCShareableContent.current

        guard let captureWindow = content.windows.first(where: { $0.windowID == windowID }) else {
            throw ScreenshotProviderError.windowNotFound(windowID)
        }

        let filter = SCContentFilter(desktopIndependentWindow: captureWindow)
        let contentInfo = SCShareableContent.info(for: filter)
        let configuration = configuration(for: contentInfo)
        let image: CGImage
        do {
            image = try await SCScreenshotManager.captureImage(contentFilter: filter, configuration: configuration)
        } catch {
            if let sampleBufferImage = try? await captureSampleBufferImage(
                filter: filter,
                configuration: configuration,
                windowID: windowID,
                pointSize: NSSize(width: contentInfo.contentRect.width, height: contentInfo.contentRect.height)
            ) {
                image = sampleBufferImage
            } else {
                throw error
            }
        }

        return WindowScreenshotCapture(
            image: image,
            pointSize: NSSize(width: contentInfo.contentRect.width, height: contentInfo.contentRect.height)
        )
    }

    private func configuration(for contentInfo: SCShareableContentInfo) -> SCStreamConfiguration {
        let configuration = SCStreamConfiguration()
        configuration.width = max(1, Int(contentInfo.contentRect.width * CGFloat(contentInfo.pointPixelScale)))
        configuration.height = max(1, Int(contentInfo.contentRect.height * CGFloat(contentInfo.pointPixelScale)))
        configuration.showsCursor = false
        configuration.ignoreShadowsSingleWindow = true
        if #available(macOS 14.2, *) {
            configuration.includeChildWindows = true
        }
        return configuration
    }

    private func captureSampleBufferImage(
        filter: SCContentFilter,
        configuration: SCStreamConfiguration,
        windowID: UInt32,
        pointSize: NSSize
    ) async throws -> CGImage {
        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<CGImage, any Error>) in
            SCScreenshotManager.captureSampleBuffer(
                contentFilter: filter,
                configuration: configuration
            ) { sampleBuffer, error in
                if let error {
                    continuation.resume(throwing: error)
                    return
                }
                guard let sampleBuffer else {
                    continuation.resume(throwing: ScreenshotProviderError.invalidCapture(
                        windowID: windowID,
                        pixelWidth: 0,
                        pixelHeight: 0,
                        pointWidth: Double(pointSize.width),
                        pointHeight: Double(pointSize.height)
                    ))
                    return
                }

                guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
                    continuation.resume(throwing: ScreenshotProviderError.invalidCapture(
                        windowID: windowID,
                        pixelWidth: 0,
                        pixelHeight: 0,
                        pointWidth: Double(pointSize.width),
                        pointHeight: Double(pointSize.height)
                    ))
                    return
                }

                let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
                guard let cgImage = CIContext().createCGImage(ciImage, from: ciImage.extent) else {
                    continuation.resume(throwing: ScreenshotProviderError.invalidCapture(
                        windowID: windowID,
                        pixelWidth: CVPixelBufferGetWidth(pixelBuffer),
                        pixelHeight: CVPixelBufferGetHeight(pixelBuffer),
                        pointWidth: Double(pointSize.width),
                        pointHeight: Double(pointSize.height)
                    ))
                    return
                }

                continuation.resume(returning: cgImage)
            }
        }
    }
}

private typealias CGSConnectionID = UInt32

private struct CGSWindowCaptureOptions: OptionSet {
    let rawValue: UInt32

    static let ignoreGlobalClipShape = CGSWindowCaptureOptions(rawValue: 1 << 11)
    static let bestResolution = CGSWindowCaptureOptions(rawValue: 1 << 8)
    static let fullSize = CGSWindowCaptureOptions(rawValue: 1 << 19)
}

@_silgen_name("CGSMainConnectionID")
private func CGSMainConnectionID() -> CGSConnectionID

@_silgen_name("CGSHWCaptureWindowList")
private func CGSHWCaptureWindowList(
    _ connectionID: CGSConnectionID,
    _ windowList: UnsafeMutablePointer<CGWindowID>,
    _ windowCount: UInt32,
    _ options: CGSWindowCaptureOptions
) -> Unmanaged<CFArray>

@MainActor
public final class ScreenCaptureKitScreenshotProvider: ScreenshotProviderProtocol {
    private let capturer: any WindowScreenshotCapturing
    private let skeletonProvider: any SkeletonThumbnailProviderProtocol
    private let policy: ScreenshotCapturePolicy
    private let debugLogger: any ScreenshotDebugLogging
    private let permissionChecker: any SystemPermissionChecking
    private let defaultSession = ScreenshotCaptureSession()

    public init(
        capturer: any WindowScreenshotCapturing = ScreenCaptureKitWindowScreenshotCapturer(),
        skeletonProvider: any SkeletonThumbnailProviderProtocol = NativeSkeletonThumbnailProvider(),
        policy: ScreenshotCapturePolicy = .round01Default,
        debugLogger: any ScreenshotDebugLogging = NoopScreenshotDebugLogger(),
        permissionChecker: any SystemPermissionChecking = SystemPermissionChecker()
    ) {
        self.capturer = capturer
        self.skeletonProvider = skeletonProvider
        self.policy = policy
        self.debugLogger = debugLogger
        self.permissionChecker = permissionChecker
    }

    public func screenshot(for window: AlignerWindow) async throws -> NSImage {
        let capture = try await captureWindowWithinTimeout(windowID: window.id)
        try validateCapture(capture, windowID: window.id)
        return NSImage(cgImage: capture.image, size: capture.pointSize)
    }

    public func resolvedScreenshot(for window: AlignerWindow) async -> ScreenshotResolution {
        await resolvedScreenshot(for: window, in: defaultSession)
    }

    public func resolvedScreenshot(for window: AlignerWindow, in session: ScreenshotCaptureSession) async -> ScreenshotResolution {
        if MinimizedWindowFallbackPolicy.isSyntheticWindow(window) {
            return fallbackResolution(for: window, reason: .syntheticWindowID(windowID: window.id))
        }
        if permissionChecker.status(for: .screenRecording) != .granted {
            return fallbackResolution(for: window, reason: .screenRecordingDenied)
        }

        let failedAttempts = session.failedAttempts(for: window.id)
        if failedAttempts > policy.maxRetriesPerSession {
            let reason = ScreenshotFallbackReason.retryLimitReached(failedAttempts: failedAttempts)
            debugLogger.record(.retrySkipped(windowID: window.id, failedAttempts: failedAttempts))
            return fallbackResolution(for: window, reason: reason)
        }

        debugLogger.record(.captureStarted(windowID: window.id))

        do {
            let image = try await screenshot(for: window)
            session.recordSuccess(for: window.id)
            debugLogger.record(.captureSucceeded(windowID: window.id))
            return ScreenshotResolution(image: image, source: .realScreenshot)
        } catch {
            let reason = fallbackReason(for: error)
            session.recordFailure(for: window.id)
            debugLogger.record(.captureFailed(windowID: window.id, reason: reason))
            return fallbackResolution(for: window, reason: reason)
        }
    }

    private func captureWindowWithinTimeout(windowID: UInt32) async throws -> WindowScreenshotCapture {
        let capturer = self.capturer
        let race = ScreenshotCaptureRace()
        let captureTask = Task { @MainActor in
            do {
                let capture = try await capturer.captureWindow(windowID: windowID)
                race.finish(.success(capture))
            } catch {
                race.finish(.failure(error))
            }
        }
        let timeoutNanoseconds = policy.timeoutNanoseconds
        let timeoutMilliseconds = policy.timeoutMilliseconds
        let timeoutTask = Task {
            do {
                try await Task.sleep(nanoseconds: timeoutNanoseconds)
                await MainActor.run {
                    race.finish(.failure(ScreenshotProviderError.captureTimedOut(
                        windowID: windowID,
                        timeoutMilliseconds: timeoutMilliseconds
                    )))
                }
            } catch {}
        }

        return try await withTaskCancellationHandler {
            defer {
                captureTask.cancel()
                timeoutTask.cancel()
            }
            return try await race.wait()
        } onCancel: {
            captureTask.cancel()
            timeoutTask.cancel()
            Task { @MainActor in
                race.finish(.failure(CancellationError()))
            }
        }
    }

    private func fallbackResolution(for window: AlignerWindow, reason: ScreenshotFallbackReason) -> ScreenshotResolution {
        let title = window.title.isEmpty ? window.app.name : window.title
        let image = skeletonProvider.withOverlayTitle(title, for: window.app)
        debugLogger.record(.fallbackUsed(windowID: window.id, reason: reason))
        return ScreenshotResolution(image: image, source: .skeletonFallback(reason))
    }

    private func fallbackReason(for error: Error) -> ScreenshotFallbackReason {
        if case let ScreenshotProviderError.captureTimedOut(windowID, timeoutMilliseconds) = error {
            return .timedOut(windowID: windowID, timeoutMilliseconds: timeoutMilliseconds)
        }
        if case let ScreenshotProviderError.invalidCapture(windowID, pixelWidth, pixelHeight, pointWidth, pointHeight) = error {
            return .invalidCapture(
                windowID: windowID,
                pixelWidth: pixelWidth,
                pixelHeight: pixelHeight,
                pointWidth: pointWidth,
                pointHeight: pointHeight
            )
        }
        return .captureFailed(Self.sanitizedCaptureFailureSummary(for: error))
    }

    private static func sanitizedCaptureFailureSummary(for error: Error) -> String {
        let nsError = error as NSError
        return [
            "domain=\(nsError.domain)",
            "code=\(nsError.code)",
            "descriptionHash=\(stableFingerprint(nsError.localizedDescription))",
            "descriptionLength=\(nsError.localizedDescription.count)"
        ].joined(separator: " ")
    }

    private static func stableFingerprint(_ value: String) -> String {
        var hash: UInt64 = 0xcbf29ce484222325
        for byte in value.utf8 {
            hash ^= UInt64(byte)
            hash = hash &* 0x100000001b3
        }
        return String(format: "%016llx", hash)
    }

    private func validateCapture(_ capture: WindowScreenshotCapture, windowID: UInt32) throws {
        let pointWidth = Double(capture.pointSize.width)
        let pointHeight = Double(capture.pointSize.height)
        guard capture.image.width > 0, capture.image.height > 0, pointWidth > 0, pointHeight > 0 else {
            throw ScreenshotProviderError.invalidCapture(
                windowID: windowID,
                pixelWidth: capture.image.width,
                pixelHeight: capture.image.height,
                pointWidth: pointWidth,
                pointHeight: pointHeight
            )
        }
    }
}