import AppKit
import ApplicationServices
import CoreGraphics
import Foundation

struct WindowRecord {
    let ownerName: String
    let name: String
    let layer: Int
    let pid: Int
    let isOnscreen: Bool
    let bounds: CGRect
    let axBounds: CGRect?

    var verificationBounds: CGRect {
        axBounds ?? bounds
    }
}

struct ScreenRecord {
    let frame: CGRect
    let visibleFrame: CGRect
    let cgFrame: CGRect
}

enum QAError: Error, CustomStringConvertible {
    case windowListUnavailable
    case forbiddenAlignerWindows([WindowRecord])
    case unexpectedOverlayCount(count: Int, expected: Int)
    case tooManyAllowedWindows(title: String, count: Int, maximum: Int)
    case unexpectedWindowCount(title: String, count: Int, expected: Int)
    case invalidOverlayLayer(title: String, layer: Int, maximumExclusive: Int)
    case mouseOutsideScreens(location: CGPoint)
    case overlayNotOnMouseScreen(title: String, bounds: CGRect, mouseLocation: CGPoint, expectedBounds: CGRect, screenFrame: CGRect)
    case overlayDoesNotUseScreenFrame(title: String, bounds: CGRect, expectedBounds: CGRect, screenFrame: CGRect, visibleFrame: CGRect)

    var description: String {
        switch self {
        case .windowListUnavailable:
            return "CGWindowListCopyWindowInfo returned no data."
        case .forbiddenAlignerWindows(let windows):
            let details = windows.map {
                "- owner=\($0.ownerName) pid=\($0.pid) layer=\($0.layer) bounds=\(formatRect($0.bounds)) axBounds=\(formatOptionalRect($0.axBounds)) title=\"\($0.name)\""
            }.joined(separator: "\n")
            return "Unexpected Aligner windows found in Round0:\n\(details)"
        case .unexpectedOverlayCount(let count, let expected):
            return "Unexpected Aligner overlay count: found \(count), expected \(expected)."
        case .tooManyAllowedWindows(let title, let count, let maximum):
            return "Too many allowed Aligner windows titled \"\(title)\": found \(count), maximum \(maximum)."
        case .unexpectedWindowCount(let title, let count, let expected):
            return "Unexpected Aligner window count for \"\(title)\": found \(count), expected \(expected)."
        case .invalidOverlayLayer(let title, let layer, let maximumExclusive):
            return "Invalid overlay layer for \"\(title)\": found \(layer), expected below \(maximumExclusive)."
        case .mouseOutsideScreens(let location):
            return "Mouse is outside all NSScreen frames: \(formatPoint(location))."
        case .overlayNotOnMouseScreen(let title, let bounds, let mouseLocation, let expectedBounds, let screenFrame):
            return "Overlay \"\(title)\" is not on the mouse screen. bounds=\(formatRect(bounds)) expectedCGWindowBounds=\(formatRect(expectedBounds)) mouse=\(formatPoint(mouseLocation)) screen.frame=\(formatRect(screenFrame))."
        case .overlayDoesNotUseScreenFrame(let title, let bounds, let expectedBounds, let screenFrame, let visibleFrame):
            return "Overlay \"\(title)\" does not use full screen.frame. bounds=\(formatRect(bounds)) expectedCGWindowBounds=\(formatRect(expectedBounds)) screen.frame=\(formatRect(screenFrame)) visibleFrame=\(formatRect(visibleFrame))."
        }
    }
}

struct QAOptions {
    let expectDebugOverlay: Bool
    let expectDebugSettingsChild: Bool
    let expectPermissions: Bool
    let expectNoPermissions: Bool
    let expectNoDebugOverlay: Bool
    let expectNoDebugSettingsChild: Bool
    let expectRound1PerformancePoC: Bool
    let expectNoRound1PerformancePoC: Bool
    let expectQuickSwitch: Bool
    let expectNoQuickSwitch: Bool
    let expectSingleAlignerOverlay: Bool
    let expectOverlayOnMouseScreen: Bool
    let expectOverlayUsesScreenFrame: Bool

    static func parse(arguments: [String]) -> QAOptions {
        QAOptions(
            expectDebugOverlay: arguments.contains("--expect-debug-overlay"),
            expectDebugSettingsChild: arguments.contains("--expect-debug-settings-child"),
            expectPermissions: arguments.contains("--expect-permissions"),
            expectNoPermissions: arguments.contains("--expect-no-permissions"),
            expectNoDebugOverlay: arguments.contains("--expect-no-debug-overlay"),
            expectNoDebugSettingsChild: arguments.contains("--expect-no-debug-settings-child"),
            expectRound1PerformancePoC: arguments.contains("--expect-round01-performance-poc")
                || arguments.contains("--expect-round1-performance-poc"),
            expectNoRound1PerformancePoC: arguments.contains("--expect-no-round01-performance-poc")
                || arguments.contains("--expect-no-round1-performance-poc")
                || arguments.contains("--expect-no-performance-poc"),
            expectQuickSwitch: arguments.contains("--expect-quick-switch"),
            expectNoQuickSwitch: arguments.contains("--expect-no-quick-switch"),
            expectSingleAlignerOverlay: arguments.contains("--expect-single-aligner-overlay"),
            expectOverlayOnMouseScreen: arguments.contains("--expect-overlay-on-mouse-screen"),
            expectOverlayUsesScreenFrame: arguments.contains("--expect-overlay-uses-screen-frame")
        )
    }
}

func boolValue(_ value: Any?) -> Bool {
    switch value {
    case let number as NSNumber:
        return number.boolValue
    case let bool as Bool:
        return bool
    default:
        return false
    }
}

func intValue(_ value: Any?) -> Int {
    switch value {
    case let number as NSNumber:
        return number.intValue
    case let int as Int:
        return int
    default:
        return 0
    }
}

func doubleValue(_ value: Any?) -> Double {
    switch value {
    case let number as NSNumber:
        return number.doubleValue
    case let double as Double:
        return double
    case let int as Int:
        return Double(int)
    default:
        return 0
    }
}

func stringValue(_ value: Any?) -> String {
    value as? String ?? ""
}

func rectValue(_ value: Any?) -> CGRect {
    guard let dictionary = value as? [String: Any] else {
        return .zero
    }

    return CGRect(
        x: doubleValue(dictionary["X"]),
        y: doubleValue(dictionary["Y"]),
        width: doubleValue(dictionary["Width"]),
        height: doubleValue(dictionary["Height"])
    )
}

func formatRect(_ rect: CGRect) -> String {
    "x=\(Int(rect.origin.x)) y=\(Int(rect.origin.y)) w=\(Int(rect.size.width)) h=\(Int(rect.size.height))"
}

func formatOptionalRect(_ rect: CGRect?) -> String {
    guard let rect else { return "nil" }
    return formatRect(rect)
}

func formatPoint(_ point: CGPoint) -> String {
    "x=\(Int(point.x)) y=\(Int(point.y))"
}

func allWindows() throws -> [WindowRecord] {
    guard let windowInfo = CGWindowListCopyWindowInfo(
        [.optionOnScreenOnly, .excludeDesktopElements],
        kCGNullWindowID
    ) as? [[String: Any]] else {
        throw QAError.windowListUnavailable
    }

    return windowInfo.map { item in
        WindowRecord(
            ownerName: stringValue(item[kCGWindowOwnerName as String]),
            name: stringValue(item[kCGWindowName as String]),
            layer: intValue(item[kCGWindowLayer as String]),
            pid: intValue(item[kCGWindowOwnerPID as String]),
            isOnscreen: boolValue(item[kCGWindowIsOnscreen as String]),
            bounds: rectValue(item[kCGWindowBounds as String]),
            axBounds: axWindowBounds(
                pid: intValue(item[kCGWindowOwnerPID as String]),
                title: stringValue(item[kCGWindowName as String])
            )
        )
    }
}

func axWindowBounds(pid: Int, title: String) -> CGRect? {
    guard pid > 0 else { return nil }

    let appElement = AXUIElementCreateApplication(pid_t(pid))
    var windowsValue: CFTypeRef?
    guard AXUIElementCopyAttributeValue(
        appElement,
        kAXWindowsAttribute as CFString,
        &windowsValue
    ) == .success,
          let windows = windowsValue as? [AXUIElement]
    else {
        return nil
    }

    for window in windows {
        var titleValue: CFTypeRef?
        let axTitle = AXUIElementCopyAttributeValue(
            window,
            kAXTitleAttribute as CFString,
            &titleValue
        ) == .success ? titleValue as? String : nil
        guard axTitle == title else { continue }

        var positionValue: CFTypeRef?
        var sizeValue: CFTypeRef?
        guard AXUIElementCopyAttributeValue(
            window,
            kAXPositionAttribute as CFString,
            &positionValue
        ) == .success,
              AXUIElementCopyAttributeValue(
                window,
                kAXSizeAttribute as CFString,
                &sizeValue
              ) == .success,
              let positionValue,
              let sizeValue,
              CFGetTypeID(positionValue) == AXValueGetTypeID(),
              CFGetTypeID(sizeValue) == AXValueGetTypeID()
        else {
            return nil
        }

        let positionAXValue = positionValue as! AXValue
        let sizeAXValue = sizeValue as! AXValue
        var position = CGPoint.zero
        var size = CGSize.zero
        guard AXValueGetValue(positionAXValue, .cgPoint, &position),
              AXValueGetValue(sizeAXValue, .cgSize, &size)
        else {
            return nil
        }

        return CGRect(origin: position, size: size)
    }

    return nil
}

func screenRecords() -> [ScreenRecord] {
    let screens = NSScreen.screens
    let globalFrame = screens.map(\.frame).reduce(CGRect.null) { partialResult, frame in
        partialResult.union(frame)
    }

    return screens.map { screen in
        let displayID = screen.deviceDescription[NSDeviceDescriptionKey("NSScreenNumber")] as? CGDirectDisplayID ?? 0
        let cgFrame = displayID == 0
            ? CGRect(
                x: screen.frame.minX,
                y: globalFrame.maxY - screen.frame.maxY,
                width: screen.frame.width,
                height: screen.frame.height
            )
            : CGDisplayBounds(displayID)

        return ScreenRecord(
            frame: screen.frame,
            visibleFrame: screen.visibleFrame,
            cgFrame: cgFrame
        )
    }
}

func mouseScreen(from screens: [ScreenRecord]) -> ScreenRecord? {
    let location = NSEvent.mouseLocation
    return screens.first { screen in
        screen.frame.contains(location)
    }
}

func equivalentRect(_ lhs: CGRect, _ rhs: CGRect, tolerance: CGFloat = 2) -> Bool {
    abs(lhs.origin.x - rhs.origin.x) <= tolerance
        && abs(lhs.origin.y - rhs.origin.y) <= tolerance
        && abs(lhs.size.width - rhs.size.width) <= tolerance
        && abs(lhs.size.height - rhs.size.height) <= tolerance
}

func equivalentWindowBounds(_ bounds: CGRect, to screen: ScreenRecord) -> Bool {
    equivalentRect(bounds, screen.cgFrame)
}

let permissionTitle = "Aligner Needs Permissions"
let debugOverlayTitle = "Aligner Round0 Debug Overlay"
let debugSettingsTitle = "Aligner Round0 Debug Settings"
let round1PerformancePoCTitle = "Aligner Round01 Performance PoC"
let quickSwitchTitle = "Aligner Quick Switch"
let overlayTitles = Set([debugOverlayTitle, round1PerformancePoCTitle, quickSwitchTitle])

func run(options: QAOptions) throws {
    let windows = try allWindows()
    let alignerWindows = windows.filter { $0.ownerName == "Aligner" }
    let alignerOverlayWindows = alignerWindows.filter { overlayTitles.contains($0.name) }
    let allowedTitleMaximums = [
        permissionTitle: 1,
        debugOverlayTitle: 1,
        debugSettingsTitle: 1,
        round1PerformancePoCTitle: 1,
        quickSwitchTitle: 1
    ]
    let allowedTitles = Set(allowedTitleMaximums.keys)
    let forbiddenWindows = alignerWindows.filter { window in
        !allowedTitles.contains(window.name)
    }

    print("WindowLogic QA")
    print("  screens: \(NSScreen.screens.count)")
    print("  onscreen windows: \(windows.count)")
    print("  aligner windows: \(alignerWindows.count)")

    for window in alignerWindows {
        print("  Aligner window: pid=\(window.pid) layer=\(window.layer) bounds=\(formatRect(window.bounds)) axBounds=\(formatOptionalRect(window.axBounds)) title=\"\(window.name)\"")
    }

    guard forbiddenWindows.isEmpty else {
        throw QAError.forbiddenAlignerWindows(forbiddenWindows)
    }

    for (title, maximum) in allowedTitleMaximums {
        let count = alignerWindows.filter { $0.name == title }.count
        if count > maximum {
            throw QAError.tooManyAllowedWindows(title: title, count: count, maximum: maximum)
        }
    }

    try assertWindow(title: debugOverlayTitle, in: alignerWindows, expected: options.expectDebugOverlay ? 1 : nil)
    try assertWindow(title: debugSettingsTitle, in: alignerWindows, expected: options.expectDebugSettingsChild ? 1 : nil)
    try assertWindow(title: permissionTitle, in: alignerWindows, expected: options.expectPermissions ? 1 : nil)
    try assertWindow(title: round1PerformancePoCTitle, in: alignerWindows, expected: options.expectRound1PerformancePoC ? 1 : nil)
    try assertWindow(title: quickSwitchTitle, in: alignerWindows, expected: options.expectQuickSwitch ? 1 : nil)

    if options.expectNoDebugOverlay {
        try assertWindow(title: debugOverlayTitle, in: alignerWindows, expected: 0)
    }
    if options.expectNoDebugSettingsChild {
        try assertWindow(title: debugSettingsTitle, in: alignerWindows, expected: 0)
    }
    if options.expectNoPermissions {
        try assertWindow(title: permissionTitle, in: alignerWindows, expected: 0)
    }
    if options.expectNoRound1PerformancePoC {
        try assertWindow(title: round1PerformancePoCTitle, in: alignerWindows, expected: 0)
    }
    if options.expectNoQuickSwitch {
        try assertWindow(title: quickSwitchTitle, in: alignerWindows, expected: 0)
    }
    if options.expectSingleAlignerOverlay && alignerOverlayWindows.count != 1 {
        throw QAError.unexpectedOverlayCount(count: alignerOverlayWindows.count, expected: 1)
    }

    let mainMenuLayer = Int(CGWindowLevelForKey(.mainMenuWindow))
    for overlay in alignerOverlayWindows {
        if overlay.layer >= mainMenuLayer {
            throw QAError.invalidOverlayLayer(
                title: overlay.name,
                layer: overlay.layer,
                maximumExclusive: mainMenuLayer
            )
        }
    }

    if options.expectOverlayOnMouseScreen || options.expectOverlayUsesScreenFrame {
        let screens = screenRecords()
        let mouseLocation = NSEvent.mouseLocation
        guard let screen = mouseScreen(from: screens) else {
            throw QAError.mouseOutsideScreens(location: mouseLocation)
        }

        for overlay in alignerOverlayWindows {
            let bounds = overlay.verificationBounds
            if options.expectOverlayOnMouseScreen && !equivalentWindowBounds(bounds, to: screen) {
                throw QAError.overlayNotOnMouseScreen(
                    title: overlay.name,
                    bounds: bounds,
                    mouseLocation: mouseLocation,
                    expectedBounds: screen.cgFrame,
                    screenFrame: screen.frame
                )
            }
            if options.expectOverlayUsesScreenFrame && !equivalentWindowBounds(bounds, to: screen) {
                throw QAError.overlayDoesNotUseScreenFrame(
                    title: overlay.name,
                    bounds: bounds,
                    expectedBounds: screen.cgFrame,
                    screenFrame: screen.frame,
                    visibleFrame: screen.visibleFrame
                )
            }
        }
    }

    print("  Overlay QA check: passed")
}

func assertWindow(title: String, in windows: [WindowRecord], expected: Int?) throws {
    guard let expected else { return }

    let count = windows.filter { $0.name == title }.count
    guard count == expected else {
        throw QAError.unexpectedWindowCount(title: title, count: count, expected: expected)
    }
}

do {
    try run(options: QAOptions.parse(arguments: CommandLine.arguments))
} catch {
    fputs("WindowLogic QA failed: \(error)\n", stderr)
    exit(1)
}