import AppKit
import ApplicationServices
import Foundation

enum ProbeFailure: Error {
    case read
    case dataInvalid
    case limit
}

enum ProbeValue<Value> {
    case value(Value)
    case absent
    case invalid
}

enum SourceDiagnosis: String, CaseIterable {
    case sourceZero = "SOURCE_ZERO"
    case sourceMultiple = "SOURCE_MULTIPLE"
    case traversalLimit = "SOURCE_TRAVERSAL_LIMIT"
    case traversalDataInvalid = "SOURCE_TRAVERSAL_DATA_INVALID"
    case traversalReadError = "SOURCE_TRAVERSAL_READ_ERROR"
    case sourceWindowIDAbsentOrInvalid = "SOURCE_WINDOW_ID_ABSENT_OR_INVALID"
    case otherWindowIDAbsentOrInvalid = "OTHER_WINDOW_ID_ABSENT_OR_INVALID"
    case duplicateWindowIdentity = "DUPLICATE_WINDOW_IDENTITY"
    case sourcePositionAbsentOrInvalid = "SOURCE_POSITION_ABSENT_OR_INVALID"
    case sourceSizeAbsentOrInvalid = "SOURCE_SIZE_ABSENT_OR_INVALID"
    case sourceGeometryNonfiniteOrNonpositive = "SOURCE_GEOMETRY_NONFINITE_OR_NONPOSITIVE"
    case sourceMinimizedAbsentOrInvalid = "SOURCE_MINIMIZED_ABSENT_OR_INVALID"
    case sourceWindowReadError = "SOURCE_WINDOW_READ_ERROR"
    case sourceWindowReady = "SOURCE_WINDOW_READY"

    var line: String { "finder_source=\(rawValue)" }
}

func classifySourceWindow(
    matchCounts: Result<[Int], ProbeFailure>,
    windowIDs: [() -> Result<ProbeValue<Int64>, ProbeFailure>],
    position: () -> Result<ProbeValue<CGPoint>, ProbeFailure>,
    size: () -> Result<ProbeValue<CGSize>, ProbeFailure>,
    minimized: () -> Result<ProbeValue<Bool>, ProbeFailure>
) -> SourceDiagnosis {
    let counts: [Int]
    switch matchCounts {
    case .failure(.limit): return .traversalLimit
    case .failure(.dataInvalid): return .traversalDataInvalid
    case .failure(.read): return .traversalReadError
    case .success(let value): counts = value
    }
    guard counts.count == windowIDs.count, counts.allSatisfy({ $0 >= 0 }) else {
        return .traversalDataInvalid
    }
    let total = counts.reduce(0, +)
    guard total > 0 else { return .sourceZero }
    guard total == 1, let sourceIndex = counts.firstIndex(of: 1) else { return .sourceMultiple }

    let sourceID: Int64
    switch windowIDs[sourceIndex]() {
    case .failure: return .sourceWindowReadError
    case .success(.absent), .success(.invalid): return .sourceWindowIDAbsentOrInvalid
    case .success(.value(let value)): sourceID = value
    }
    let point: CGPoint
    switch position() {
    case .failure: return .sourceWindowReadError
    case .success(.absent), .success(.invalid): return .sourcePositionAbsentOrInvalid
    case .success(.value(let value)): point = value
    }
    let dimensions: CGSize
    switch size() {
    case .failure: return .sourceWindowReadError
    case .success(.absent), .success(.invalid): return .sourceSizeAbsentOrInvalid
    case .success(.value(let value)): dimensions = value
    }
    guard [point.x, point.y, dimensions.width, dimensions.height].allSatisfy(\.isFinite),
          dimensions.width > 0, dimensions.height > 0 else {
        return .sourceGeometryNonfiniteOrNonpositive
    }
    switch minimized() {
    case .failure: return .sourceWindowReadError
    case .success(.absent), .success(.invalid): return .sourceMinimizedAbsentOrInvalid
    case .success(.value): break
    }
    for (index, readID) in windowIDs.enumerated() where index != sourceIndex {
        switch readID() {
        case .failure: return .sourceWindowReadError
        case .success(.absent), .success(.invalid): return .otherWindowIDAbsentOrInvalid
        case .success(.value(let value)) where value == sourceID: return .duplicateWindowIdentity
        case .success(.value): break
        }
    }
    return .sourceWindowReady
}

func selfTest() -> Int32 {
    typealias IDRead = () -> Result<ProbeValue<Int64>, ProbeFailure>
    let goodPosition = CGPoint(x: 10, y: 20)
    let goodSize = CGSize(width: 500, height: 320)
    func classify(
        counts: Result<[Int], ProbeFailure> = .success([1]),
        ids: [IDRead] = [{ .success(.value(7)) }],
        position: @escaping () -> Result<ProbeValue<CGPoint>, ProbeFailure> = { .success(.value(CGPoint(x: 10, y: 20))) },
        size: @escaping () -> Result<ProbeValue<CGSize>, ProbeFailure> = { .success(.value(CGSize(width: 500, height: 320))) },
        minimized: @escaping () -> Result<ProbeValue<Bool>, ProbeFailure> = { .success(.value(false)) }
    ) -> SourceDiagnosis {
        classifySourceWindow(matchCounts: counts, windowIDs: ids, position: position, size: size, minimized: minimized)
    }
    let cases: [(SourceDiagnosis, SourceDiagnosis)] = [
        (.sourceWindowReady, classify()),
        (.sourceZero, classify(counts: .success([0]))),
        (.sourceMultiple, classify(counts: .success([1, 1]), ids: [{ .success(.value(7)) }, { .success(.value(8)) }])),
        (.sourceMultiple, classify(counts: .success([2]))),
        (.traversalLimit, classify(counts: .failure(.limit))),
        (.traversalDataInvalid, classify(counts: .failure(.dataInvalid))),
        (.traversalReadError, classify(counts: .failure(.read))),
        (.traversalDataInvalid, classify(counts: .success([1, 0]), ids: [{ .success(.value(7)) }])),
        (.sourceWindowIDAbsentOrInvalid, classify(ids: [{ .success(.absent) }])),
        (.sourceWindowIDAbsentOrInvalid, classify(ids: [{ .success(.invalid) }])),
        (.sourceWindowReadError, classify(ids: [{ .failure(.read) }])),
        (.sourcePositionAbsentOrInvalid, classify(position: { .success(.absent) })),
        (.sourcePositionAbsentOrInvalid, classify(position: { .success(.invalid) })),
        (.sourceWindowReadError, classify(position: { .failure(.read) })),
        (.sourceSizeAbsentOrInvalid, classify(size: { .success(.absent) })),
        (.sourceSizeAbsentOrInvalid, classify(size: { .success(.invalid) })),
        (.sourceWindowReadError, classify(size: { .failure(.read) })),
        (.sourceGeometryNonfiniteOrNonpositive, classify(position: { .success(.value(CGPoint(x: CGFloat.nan, y: 0))) })),
        (.sourceGeometryNonfiniteOrNonpositive, classify(size: { .success(.value(CGSize(width: 0, height: 1))) })),
        (.sourceGeometryNonfiniteOrNonpositive, classify(size: { .success(.value(CGSize(width: 1, height: CGFloat.infinity))) })),
        (.sourceMinimizedAbsentOrInvalid, classify(minimized: { .success(.absent) })),
        (.sourceMinimizedAbsentOrInvalid, classify(minimized: { .success(.invalid) })),
        (.sourceWindowReadError, classify(minimized: { .failure(.read) })),
        (.duplicateWindowIdentity, classify(
            counts: .success([1, 0]),
            ids: [{ .success(.value(7)) }, { .success(.value(7)) }]
        )),
        (.sourceWindowReadError, classify(
            counts: .success([1, 0]),
            ids: [{ .success(.value(7)) }, { .failure(.read) }]
        )),
        (.otherWindowIDAbsentOrInvalid, classify(
            counts: .success([1, 0]),
            ids: [{ .success(.value(7)) }, { .success(.absent) }]
        )),
        (.otherWindowIDAbsentOrInvalid, classify(
            counts: .success([1, 0]),
            ids: [{ .success(.value(7)) }, { .success(.invalid) }]
        )),
    ]
    var terminalReads = 0
    let terminal = classifySourceWindow(
        matchCounts: .success([1]),
        windowIDs: [{ terminalReads += 1; return .success(.absent) }],
        position: { terminalReads += 1; return .success(.value(goodPosition)) },
        size: { terminalReads += 1; return .success(.value(goodSize)) },
        minimized: { terminalReads += 1; return .success(.value(false)) }
    )
    var otherIDReads = 0
    let otherIDTerminal = classify(
        counts: .success([1, 0, 0]),
        ids: [
            { otherIDReads += 1; return .success(.value(7)) },
            { otherIDReads += 1; return .success(.absent) },
            { otherIDReads += 1; return .success(.value(9)) },
        ]
    )
    guard cases.allSatisfy({ $0.0 == $0.1 }),
          terminal == .sourceWindowIDAbsentOrInvalid,
          terminalReads == 1,
          otherIDTerminal == .otherWindowIDAbsentOrInvalid,
          otherIDReads == 2,
          SourceDiagnosis.allCases.allSatisfy({
              $0.line == "finder_source=\($0.rawValue)" && $0.line.split(separator: "\n").count == 1
          }) else { return 1 }
    print("self_test=PASS finder_source_diagnostics=true")
    return 0
}

if CommandLine.arguments == [CommandLine.arguments[0], "--self-test"] { exit(selfTest()) }
guard CommandLine.arguments.count == 3,
      CommandLine.arguments[1] == "diagnose",
      !CommandLine.arguments[2].isEmpty else { exit(64) }
let exactSourceIdentifier = CommandLine.arguments[2]

func attribute(_ element: AXUIElement, _ name: String) -> Result<ProbeValue<CFTypeRef>, ProbeFailure> {
    var value: CFTypeRef?
    let error = AXUIElementCopyAttributeValue(element, name as CFString, &value)
    if error == .noValue || error == .attributeUnsupported { return .success(.absent) }
    guard error == .success else { return .failure(.read) }
    guard let value else { return .success(.invalid) }
    return .success(.value(value))
}

func stringValue(_ element: AXUIElement, _ name: String) -> Result<ProbeValue<String>, ProbeFailure> {
    switch attribute(element, name) {
    case .failure(let failure): return .failure(failure)
    case .success(.absent): return .success(.absent)
    case .success(.invalid): return .success(.invalid)
    case .success(.value(let value)):
        guard CFGetTypeID(value) == CFStringGetTypeID(), let result = value as? String else {
            return .success(.invalid)
        }
        return .success(.value(result))
    }
}

func windowID(_ element: AXUIElement) -> Result<ProbeValue<Int64>, ProbeFailure> {
    switch attribute(element, "AXWindowNumber") {
    case .failure(let failure): return .failure(failure)
    case .success(.absent): return .success(.absent)
    case .success(.invalid): return .success(.invalid)
    case .success(.value(let value)):
        guard CFGetTypeID(value) == CFNumberGetTypeID() else { return .success(.invalid) }
        var result: Int64 = 0
        guard CFNumberGetValue((value as! CFNumber), .sInt64Type, &result) else { return .success(.invalid) }
        return .success(.value(result))
    }
}

func pointValue(_ element: AXUIElement) -> Result<ProbeValue<CGPoint>, ProbeFailure> {
    switch attribute(element, kAXPositionAttribute) {
    case .failure(let failure): return .failure(failure)
    case .success(.absent): return .success(.absent)
    case .success(.invalid): return .success(.invalid)
    case .success(.value(let value)):
        guard CFGetTypeID(value) == AXValueGetTypeID() else { return .success(.invalid) }
        var point = CGPoint.zero
        guard AXValueGetValue(value as! AXValue, .cgPoint, &point) else { return .success(.invalid) }
        return .success(.value(point))
    }
}

func sizeValue(_ element: AXUIElement) -> Result<ProbeValue<CGSize>, ProbeFailure> {
    switch attribute(element, kAXSizeAttribute) {
    case .failure(let failure): return .failure(failure)
    case .success(.absent): return .success(.absent)
    case .success(.invalid): return .success(.invalid)
    case .success(.value(let value)):
        guard CFGetTypeID(value) == AXValueGetTypeID() else { return .success(.invalid) }
        var size = CGSize.zero
        guard AXValueGetValue(value as! AXValue, .cgSize, &size) else { return .success(.invalid) }
        return .success(.value(size))
    }
}

func boolValue(_ element: AXUIElement, _ name: String) -> Result<ProbeValue<Bool>, ProbeFailure> {
    switch attribute(element, name) {
    case .failure(let failure): return .failure(failure)
    case .success(.absent): return .success(.absent)
    case .success(.invalid): return .success(.invalid)
    case .success(.value(let value)):
        guard CFGetTypeID(value) == CFBooleanGetTypeID() else { return .success(.invalid) }
        return .success(.value(CFBooleanGetValue((value as! CFBoolean))))
    }
}

func children(_ element: AXUIElement) -> Result<ProbeValue<[AXUIElement]>, ProbeFailure> {
    switch attribute(element, kAXChildrenAttribute) {
    case .failure(let failure): return .failure(failure)
    case .success(.absent): return .success(.absent)
    case .success(.invalid): return .success(.invalid)
    case .success(.value(let value)):
        guard CFGetTypeID(value) == CFArrayGetTypeID(), let result = value as? [AXUIElement] else {
            return .success(.invalid)
        }
        return .success(.value(result))
    }
}

func countExactSource(_ element: AXUIElement, depth: Int = 0, visited: inout Int) throws -> Int {
    guard depth <= 14, visited < 20_000 else { throw ProbeFailure.limit }
    visited += 1
    let role: ProbeValue<String>
    switch stringValue(element, kAXRoleAttribute) {
    case .failure(let failure): throw failure
    case .success(let value): role = value
    }
    var count = 0
    if case .value(let roleValue) = role, roleValue == (kAXTextFieldRole as String) {
        switch stringValue(element, kAXValueAttribute) {
        case .failure(let failure): throw failure
        case .success(.invalid): throw ProbeFailure.dataInvalid
        case .success(.value(let value)) where value == exactSourceIdentifier: count += 1
        case .success: break
        }
    } else if case .invalid = role {
        throw ProbeFailure.dataInvalid
    }
    switch children(element) {
    case .failure(let failure): throw failure
    case .success(.invalid): throw ProbeFailure.dataInvalid
    case .success(.absent): break
    case .success(.value(let values)):
        for child in values { count += try countExactSource(child, depth: depth + 1, visited: &visited) }
    }
    return count
}

func emit(_ diagnosis: SourceDiagnosis) -> Never {
    print(diagnosis.line)
    exit(diagnosis == .sourceWindowReady ? 0 : 69)
}

guard let finder = NSRunningApplication.runningApplications(withBundleIdentifier: "com.apple.finder").first else {
    emit(.traversalReadError)
}
let app = AXUIElementCreateApplication(finder.processIdentifier)
let windows: [AXUIElement]
switch attribute(app, kAXWindowsAttribute) {
case .failure: emit(.traversalReadError)
case .success(.absent), .success(.invalid): emit(.traversalDataInvalid)
case .success(.value(let value)):
    guard CFGetTypeID(value) == CFArrayGetTypeID(), let typed = value as? [AXUIElement] else {
        emit(.traversalDataInvalid)
    }
    windows = typed
}
let matchCounts: Result<[Int], ProbeFailure>
do {
    var values: [Int] = []
    for window in windows {
        var visited = 0
        values.append(try countExactSource(window, visited: &visited))
    }
    matchCounts = .success(values)
} catch let failure as ProbeFailure {
    matchCounts = .failure(failure)
} catch {
    matchCounts = .failure(.read)
}
let sourceIndex: Int? = {
    guard case .success(let counts) = matchCounts, counts.reduce(0, +) == 1 else { return nil }
    return counts.firstIndex(of: 1)
}()
let diagnosis = classifySourceWindow(
    matchCounts: matchCounts,
    windowIDs: windows.map { window in { windowID(window) } },
    position: { sourceIndex.map { pointValue(windows[$0]) } ?? .success(.invalid) },
    size: { sourceIndex.map { sizeValue(windows[$0]) } ?? .success(.invalid) },
    minimized: { sourceIndex.map { boolValue(windows[$0], kAXMinimizedAttribute) } ?? .success(.invalid) }
)
emit(diagnosis)