import AppKit
import Foundation
import AlignerCore

@MainActor
enum Round1ScreenshotProbe {
    static func run(prettyPrinted: Bool) async throws {
        let spaceEnumerator = try SkyLightSpaceEnumerator()
        let displays = try spaceEnumerator.displays()
        let fullscreenSpaceIDs = Set(
            displays
                .flatMap(\.spaces)
                .filter { $0.type == .fullscreen }
                .map(\.id)
        )
        let windowService = CGWindowAXWindowService(
            spaceIDsByWindowIDProvider: spaceEnumerator.spaceIDsByWindowID(windowIDs:),
            fullscreenSpaceIDs: fullscreenSpaceIDs
        )
        let windows = try windowService.allWindows()
        let snapshot = QuickSwitchSnapshotBuilder.snapshot(displays: displays, windows: windows)
        let logger = Round1ScreenshotProbeLogger()
        let provider = ScreenCaptureKitScreenshotProvider(debugLogger: logger)
        let session = ScreenshotCaptureSession()
        let candidates = snapshot.appGroups.flatMap { group in
            group.windows.map(\.window)
        }

        var sampleResults: [[String: Any]] = []
        for kind in Round1ScreenshotProbeSampleKind.allCases {
            guard let window = candidates.first(where: { kind.matches($0) }) else {
                sampleResults.append([
                    "kind": kind.rawValue,
                    "status": "missingSampleWindow"
                ])
                continue
            }

            let startedAt = Date()
            let resolution = await provider.resolvedScreenshot(for: window, in: session)
            let elapsedMilliseconds = Int(Date().timeIntervalSince(startedAt) * 1000)
            sampleResults.append(
                sampleDictionary(
                    kind: kind,
                    window: window,
                    resolution: resolution,
                    elapsedMilliseconds: elapsedMilliseconds,
                    events: logger.events(for: window.id)
                )
            )
        }

        let json = try JSONSerialization.data(
            withJSONObject: dictionary(
                displays: displays,
                windows: windows,
                snapshot: snapshot,
                samples: sampleResults
            ),
            options: prettyPrinted ? [.prettyPrinted, .sortedKeys] : [.sortedKeys]
        )

        FileHandle.standardOutput.write(json)
        FileHandle.standardOutput.write(Data("\n".utf8))
    }

    private static func dictionary(
        displays: [AlignerDisplay],
        windows: [AlignerWindow],
        snapshot: QuickSwitchSnapshot,
        samples: [[String: Any]]
    ) -> [String: Any] {
        let missingCount = samples.filter { $0["status"] as? String == "missingSampleWindow" }.count
        let realCount = samples.filter { $0["source"] as? String == "realScreenshot" }.count
        let fallbackCount = samples.filter { $0["source"] as? String == "skeletonFallback" }.count

        return [
            "generatedAt": ISO8601DateFormatter().string(from: Date()),
            "displayCount": displays.count,
            "spaceCount": displays.flatMap(\.spaces).count,
            "candidateWindowCount": windows.count,
            "snapshotAppCount": snapshot.appGroups.count,
            "snapshotWindowCount": snapshot.appGroups.flatMap(\.windows).count,
            "sampleKinds": Round1ScreenshotProbeSampleKind.allCases.map(\.rawValue),
            "missingSampleCount": missingCount,
            "realScreenshotCount": realCount,
            "fallbackCount": fallbackCount,
            "completeSampleSet": missingCount == 0,
            "samples": samples
        ]
    }

    private static func sampleDictionary(
        kind: Round1ScreenshotProbeSampleKind,
        window: AlignerWindow,
        resolution: ScreenshotResolution,
        elapsedMilliseconds: Int,
        events: [ScreenshotDebugEvent]
    ) -> [String: Any] {
        var result: [String: Any] = [
            "kind": kind.rawValue,
            "status": "sampled",
            "elapsedMilliseconds": elapsedMilliseconds,
            "window": windowDictionary(window),
            "events": events.map(eventString)
        ]

        switch resolution.source {
        case .realScreenshot:
            result["source"] = "realScreenshot"
        case .skeletonFallback(let reason):
            result["source"] = "skeletonFallback"
            result["fallbackReason"] = fallbackReasonString(reason)
        }

        result["imageSize"] = [
            "width": resolution.image.size.width,
            "height": resolution.image.size.height
        ]

        return result
    }

    private static func windowDictionary(_ window: AlignerWindow) -> [String: Any] {
        var result = [
            "id": window.id,
            "titleHash": DevelopmentDiagnostics.stableFingerprint(window.title),
            "titleLength": window.title.count,
            "titleIsEmpty": window.title.isEmpty,
            "isMinimized": window.isMinimized,
            "isFullscreen": window.isFullscreen,
            "spaceIDs": window.spaceIDs,
            "app": [
                "bundleIdentifier": window.app.bundleIdentifier,
                "name": window.app.name,
                "category": window.app.category.rawValue,
                "processIdentifier": window.app.processIdentifier as Any
            ]
        ] as [String: Any]
        if DevelopmentDiagnostics.includesSensitiveFields {
            result["title"] = window.title
        }
        return result
    }

    private static func eventString(_ event: ScreenshotDebugEvent) -> String {
        switch event {
        case .captureStarted(let windowID):
            return "captureStarted(\(windowID))"
        case .captureSucceeded(let windowID):
            return "captureSucceeded(\(windowID))"
        case .captureFailed(let windowID, let reason):
            return "captureFailed(\(windowID), \(fallbackReasonString(reason)))"
        case .retrySkipped(let windowID, let failedAttempts):
            return "retrySkipped(\(windowID), failedAttempts: \(failedAttempts))"
        case .fallbackUsed(let windowID, let reason):
            return "fallbackUsed(\(windowID), \(fallbackReasonString(reason)))"
        }
    }

    private static func fallbackReasonString(_ reason: ScreenshotFallbackReason) -> String {
        reason.diagnosticDescription
    }
}

private enum Round1ScreenshotProbeSampleKind: String, CaseIterable {
    case finder
    case browser
    case terminal
    case document
    case systemSettings

    func matches(_ window: AlignerWindow) -> Bool {
        guard !window.isMinimized else { return false }

        let bundleIdentifier = window.app.bundleIdentifier
        let appName = window.app.name.localizedLowercase

        switch self {
        case .finder:
            return bundleIdentifier == "com.apple.finder" || window.app.category == .finder
        case .browser:
            return window.app.category == .browser
        case .terminal:
            return window.app.category == .terminal
        case .document:
            return window.app.category == .documentNotes
                || [
                    "com.apple.TextEdit",
                    "com.apple.Notes",
                    "com.apple.Preview",
                    "com.microsoft.Word",
                    "io.typora",
                    "md.obsidian"
                ].contains(bundleIdentifier)
        case .systemSettings:
            return bundleIdentifier == "com.apple.systempreferences"
                || appName == "system settings"
                || appName == "system preferences"
                || appName == "系统设置"
                || appName == "系统偏好设置"
        }
    }
}

@MainActor
private final class Round1ScreenshotProbeLogger: ScreenshotDebugLogging {
    private var eventsByWindowID: [UInt32: [ScreenshotDebugEvent]] = [:]

    func record(_ event: ScreenshotDebugEvent) {
        eventsByWindowID[event.windowID, default: []].append(event)
    }

    func events(for windowID: UInt32) -> [ScreenshotDebugEvent] {
        eventsByWindowID[windowID, default: []]
    }
}

private extension ScreenshotDebugEvent {
    var windowID: UInt32 {
        switch self {
        case .captureStarted(let windowID),
             .captureSucceeded(let windowID),
             .captureFailed(let windowID, _),
             .retrySkipped(let windowID, _),
             .fallbackUsed(let windowID, _):
            return windowID
        }
    }
}

struct Round1ScreenshotProbeOptions {
    let enabled: Bool
    let prettyPrinted: Bool

    static func parse(arguments: [String]) -> Round1ScreenshotProbeOptions {
        Round1ScreenshotProbeOptions(
            enabled: arguments.contains("--round01-screenshot-probe"),
            prettyPrinted: arguments.contains("--round01-screenshot-probe-pretty")
        )
    }
}