import AppKit
import ApplicationServices
import Foundation

let canvasRole = "AXWebArea"

struct CanvasDescriptor {
    let role: String
    let rect: CGRect?
    let hidden: HiddenRead
}

enum HiddenRead {
    case value(Bool)
    case absent
    case invalid
}

enum CanvasReadFailure: Error {
    case nonBenign(AXError)
    case traversalLimit
}

func isBenignAbsence(_ error: AXError) -> Bool {
    error == .noValue || error == .attributeUnsupported
}

func validRect(_ rect: CGRect?) -> Bool {
    guard let rect else { return false }
    return [rect.origin.x, rect.origin.y, rect.width, rect.height].allSatisfy(\.isFinite)
        && rect.width > 0 && rect.height > 0
}

func convertAppKitVisibleFramesToAX(
    primaryFrame: CGRect?,
    visibleFrames: [CGRect]
) -> [CGRect]? {
    guard let primaryFrame, validRect(primaryFrame), !visibleFrames.isEmpty,
          visibleFrames.allSatisfy({ validRect($0) }) else { return nil }
    let primaryTop = primaryFrame.maxY
    return visibleFrames.map { visible in
        CGRect(
            x: visible.minX,
            y: primaryTop - visible.maxY,
            width: visible.width,
            height: visible.height
        )
    }
}

func uniqueCanvas(
    _ result: Result<[CanvasDescriptor], CanvasReadFailure>,
    visibleFrames: [CGRect]
) -> CanvasDescriptor? {
    guard let descriptors = try? result.get(), !visibleFrames.isEmpty else { return nil }
    let matches = descriptors.filter { $0.role == canvasRole }
    guard matches.count == 1 else { return nil }
    switch matches[0].hidden {
    case .absent, .value(false): break
    case .invalid, .value(true): return nil
    }
    guard
          let rect = matches[0].rect,
          validRect(rect),
          visibleFrames.contains(where: { validRect($0) && $0.contains(rect) }) else { return nil }
    return matches[0]
}

func selfTest() -> Int32 {
    let display = CGRect(x: 0, y: 0, width: 1920, height: 1040)
    let primaryFrame = CGRect(x: 0, y: 0, width: 1920, height: 1080)
    let appKitVisibleFrames = [
        CGRect(x: 0, y: 0, width: 1920, height: 1055),
        CGRect(x: 0, y: 1080, width: 1600, height: 1200),
        CGRect(x: 0, y: -900, width: 1600, height: 900),
        CGRect(x: -1280, y: 0, width: 1280, height: 1024),
        CGRect(x: 1920, y: 0, width: 1280, height: 1024),
    ]
    guard let convertedFrames = convertAppKitVisibleFramesToAX(
        primaryFrame: primaryFrame,
        visibleFrames: appKitVisibleFrames
    ) else { return 1 }
    let good = CanvasDescriptor(
        role: canvasRole,
        rect: CGRect(x: 100, y: 100, width: 800, height: 600),
        hidden: .value(false)
    )
    let invalidGeometry = [
        CGRect(x: 0, y: 0, width: 0, height: 1),
        CGRect(x: 0, y: 0, width: 1, height: 0),
        CGRect(x: CGFloat.nan, y: 0, width: 1, height: 1),
        CGRect(x: 0, y: CGFloat.infinity, width: 1, height: 1),
    ]
    let outOfBounds = CanvasDescriptor(
        role: good.role,
        rect: CGRect(x: 1800, y: 100, width: 200, height: 200),
        hidden: .value(false)
    )
    do {
        guard case .absent = try hiddenRead(error: .noValue, value: nil),
              case .absent = try hiddenRead(error: .attributeUnsupported, value: nil),
              case .invalid = try hiddenRead(error: .success, value: nil),
              case .invalid = try hiddenRead(error: .success, value: "invalid" as CFString),
              case .value(true) = try hiddenRead(error: .success, value: kCFBooleanTrue),
              case .value(false) = try hiddenRead(error: .success, value: kCFBooleanFalse) else {
            return 1
        }
        do {
            _ = try hiddenRead(error: .cannotComplete, value: nil)
            return 1
        } catch CanvasReadFailure.nonBenign(.cannotComplete) {
        }
    } catch {
        return 1
    }
    guard convertedFrames[0] == CGRect(x: 0, y: 25, width: 1920, height: 1055),
          convertedFrames[1] == CGRect(x: 0, y: -1200, width: 1600, height: 1200),
          convertedFrames[2] == CGRect(x: 0, y: 1080, width: 1600, height: 900),
          convertedFrames[3] == CGRect(x: -1280, y: 56, width: 1280, height: 1024),
          convertedFrames[4] == CGRect(x: 1920, y: 56, width: 1280, height: 1024),
          convertAppKitVisibleFramesToAX(primaryFrame: nil, visibleFrames: appKitVisibleFrames) == nil,
          convertAppKitVisibleFramesToAX(primaryFrame: primaryFrame, visibleFrames: []) == nil,
          uniqueCanvas(.success([good]), visibleFrames: [display]) != nil,
          uniqueCanvas(.success([]), visibleFrames: [display]) == nil,
          uniqueCanvas(.success([good, good]), visibleFrames: [display]) == nil,
          uniqueCanvas(.failure(.nonBenign(.cannotComplete)), visibleFrames: [display]) == nil,
          invalidGeometry.allSatisfy({ rect in
              uniqueCanvas(
                  .success([CanvasDescriptor(role: good.role, rect: rect, hidden: .value(false))]),
                  visibleFrames: [display]
              ) == nil
          }),
          uniqueCanvas(
              .success([CanvasDescriptor(role: good.role, rect: good.rect, hidden: .value(true))]),
              visibleFrames: [display]
          ) == nil,
          uniqueCanvas(
              .success([CanvasDescriptor(role: good.role, rect: good.rect, hidden: .absent)]),
              visibleFrames: [display]
          ) != nil,
          uniqueCanvas(
              .success([CanvasDescriptor(role: good.role, rect: good.rect, hidden: .invalid)]),
              visibleFrames: [display]
          ) == nil,
          uniqueCanvas(
              .success([CanvasDescriptor(role: good.role, rect: nil, hidden: .absent)]),
              visibleFrames: [display]
          ) == nil,
          uniqueCanvas(
              .success([CanvasDescriptor(role: outOfBounds.role, rect: outOfBounds.rect, hidden: .absent)]),
              visibleFrames: [display]
          ) == nil,
          uniqueCanvas(.success([outOfBounds]), visibleFrames: [display]) == nil,
          uniqueCanvas(.success([good]), visibleFrames: []) == nil else { return 1 }
    print("self_test=PASS primary_display_ax_basis=true multi_display_layouts=true unique_web_area=true zero_rejected=true multiple_rejected=true read_error_rejected=true geometry_rejected=true visibility_rejected=true display_bounds_rejected=true")
    return 0
}

if CommandLine.arguments == [CommandLine.arguments[0], "--self-test"] {
    exit(selfTest())
}

guard CommandLine.arguments.count == 2,
      let candidatePID = pid_t(CommandLine.arguments[1]),
      candidatePID > 0 else {
    fputs("usage: CanvasTargetInspect candidatePID\n", stderr)
    exit(64)
}

func attribute(_ element: AXUIElement, _ name: String) throws -> CFTypeRef? {
    var value: CFTypeRef?
    let error = AXUIElementCopyAttributeValue(element, name as CFString, &value)
    if error == .success { return value }
    if isBenignAbsence(error) { return nil }
    throw CanvasReadFailure.nonBenign(error)
}

func stringValue(_ element: AXUIElement, _ name: String) throws -> String? {
    guard let value = try attribute(element, name), CFGetTypeID(value) == CFStringGetTypeID() else { return nil }
    return value as? String
}

func hiddenRead(error: AXError, value: CFTypeRef?) throws -> HiddenRead {
    if isBenignAbsence(error) { return .absent }
    guard error == .success else { throw CanvasReadFailure.nonBenign(error) }
    guard let value, CFGetTypeID(value) == CFBooleanGetTypeID() else { return .invalid }
    return .value(CFBooleanGetValue((value as! CFBoolean)))
}

func hiddenValue(_ element: AXUIElement) throws -> HiddenRead {
    var value: CFTypeRef?
    let error = AXUIElementCopyAttributeValue(element, "AXHidden" as CFString, &value)
    return try hiddenRead(error: error, value: value)
}

func rect(_ element: AXUIElement) throws -> CGRect? {
    guard let position = try attribute(element, kAXPositionAttribute),
          CFGetTypeID(position) == AXValueGetTypeID(),
          let size = try attribute(element, kAXSizeAttribute),
          CFGetTypeID(size) == AXValueGetTypeID() else { return nil }
    var point = CGPoint.zero
    var dimensions = CGSize.zero
    guard AXValueGetValue(position as! AXValue, .cgPoint, &point),
          AXValueGetValue(size as! AXValue, .cgSize, &dimensions) else { return nil }
    return CGRect(origin: point, size: dimensions)
}

func collect(
    _ element: AXUIElement,
    depth: Int = 0,
    visited: inout Int,
    into result: inout [CanvasDescriptor]
) throws {
    guard depth <= 32, visited < 20_000 else { throw CanvasReadFailure.traversalLimit }
    visited += 1
    let role = try stringValue(element, kAXRoleAttribute)
    if role == canvasRole {
        result.append(CanvasDescriptor(
            role: role!,
            rect: try rect(element),
            hidden: try hiddenValue(element)
        ))
    }
    guard let childrenValue = try attribute(element, kAXChildrenAttribute) else { return }
    guard CFGetTypeID(childrenValue) == CFArrayGetTypeID(),
          let children = childrenValue as? [AXUIElement] else {
        throw CanvasReadFailure.nonBenign(.illegalArgument)
    }
    for child in children {
        try collect(child, depth: depth + 1, visited: &visited, into: &result)
    }
}

func axVisibleFrames() -> [CGRect]? {
    let screens = NSScreen.screens
    guard let primaryScreen = screens.first else { return nil }
    return convertAppKitVisibleFramesToAX(
        primaryFrame: primaryScreen.frame,
        visibleFrames: screens.map(\.visibleFrame)
    )
}

let app = AXUIElementCreateApplication(candidatePID)
var descriptors: [CanvasDescriptor] = []
var visited = 0
let traversal: Result<[CanvasDescriptor], CanvasReadFailure>
do {
    try collect(app, visited: &visited, into: &descriptors)
    traversal = .success(descriptors)
} catch let error as CanvasReadFailure {
    traversal = .failure(error)
}

guard let visibleFrames = axVisibleFrames(),
      let target = uniqueCanvas(traversal, visibleFrames: visibleFrames),
      let targetRect = target.rect else {
    fputs("candidate canvas binding failed closed\n", stderr)
    exit(67)
}

print("status=PASS role=AXWebArea")
print("geometry=\(targetRect.origin.x),\(targetRect.origin.y),\(targetRect.width),\(targetRect.height)")