import Foundation
import IOKit
import IOKit.hid
import CoreGraphics
import ApplicationServices

func log(_ message: String) {
    print("\(ISO8601DateFormatter().string(from: Date())) \(message)")
    fflush(stdout)
}

struct ReportSource: Hashable {
    let device: Int
    let reportID: UInt32
}

final class State {
    var buffers: [IOHIDDevice: UnsafeMutablePointer<UInt8>] = [:]
    var lastBits: [ReportSource: UInt8] = [:]
    var lastFired: [CGKeyCode: TimeInterval] = [:]
    var opened = false
    var lastStatus = ""
}

let state = State()
let statePtr = Unmanaged.passUnretained(state).toOpaque()

func sendCommandKey(_ keyCode: CGKeyCode) {
    guard CGPreflightPostEventAccess(),
          let source = CGEventSource(stateID: .hidSystemState),
          let down = CGEvent(keyboardEventSource: source, virtualKey: keyCode, keyDown: true),
          let up = CGEvent(keyboardEventSource: source, virtualKey: keyCode, keyDown: false)
    else { return }
    down.flags = .maskCommand
    up.flags = .maskCommand
    down.post(tap: .cgSessionEventTap)
    up.post(tap: .cgSessionEventTap)
}

// The mouse reports tilt on either bit pair depending on its mode after (re)connecting:
// 0x20/0x40, or 0x08/0x10 (seen after a replug).
let leftTiltMask: UInt8 = 0x20 | 0x08
let rightTiltMask: UInt8 = 0x40 | 0x10
let tiltMask = leftTiltMask | rightTiltMask
// Ignore a second press of the same tilt within this window (switch bounce).
let debounceInterval: TimeInterval = 0.12

// Preserve the supplied Razer report layout and Command+[ / Command+] mapping.
// The mouse repeats the tilt bit in every report while held, so fire only on the press edge.
let reportCallback: IOHIDReportCallback = {
    context, result, sender, _, reportID, report, reportLength in
    guard result == kIOReturnSuccess, reportLength >= 1, let context = context else { return }
    let state = Unmanaged<State>.fromOpaque(context).takeUnretainedValue()
    let source = ReportSource(device: Int(bitPattern: sender), reportID: reportID)
    let value = report[0] & tiltMask
    let pressed = value & ~(state.lastBits[source] ?? 0)
    state.lastBits[source] = value
    let now = ProcessInfo.processInfo.systemUptime
    for (mask, keyCode) in [(leftTiltMask, CGKeyCode(33)), (rightTiltMask, CGKeyCode(30))] where pressed & mask != 0 {
        if let last = state.lastFired[keyCode], now - last < debounceInterval { continue }
        state.lastFired[keyCode] = now
        sendCommandKey(keyCode)
    }
}

let deviceCallback: IOHIDDeviceCallback = { context, result, _, device in
    guard result == kIOReturnSuccess, let context = context else { return }
    let state = Unmanaged<State>.fromOpaque(context).takeUnretainedValue()
    guard state.buffers[device] == nil else { return }
    let reportedSize = IOHIDDeviceGetProperty(device, kIOHIDMaxInputReportSizeKey as CFString) as? NSNumber
    let bufferSize = max(64, reportedSize?.intValue ?? 64)
    let buffer = UnsafeMutablePointer<UInt8>.allocate(capacity: bufferSize)
    buffer.initialize(repeating: 0, count: bufferSize)
    state.buffers[device] = buffer
    IOHIDDeviceRegisterInputReportCallback(device, buffer, bufferSize, reportCallback, context)
    let product = IOHIDDeviceGetProperty(device, kIOHIDProductKey as CFString)
    log("Watching: \(product ?? "Razer Pro Click V2" as CFString)")
}

let removalCallback: IOHIDDeviceCallback = { context, _, _, device in
    guard let context = context else { return }
    let state = Unmanaged<State>.fromOpaque(context).takeUnretainedValue()
    let deviceKey = Int(bitPattern: Unmanaged.passUnretained(device).toOpaque())
    state.lastBits = state.lastBits.filter { $0.key.device != deviceKey }
    if let buffer = state.buffers.removeValue(forKey: device) {
        IOHIDDeviceRegisterInputReportCallback(device, buffer, 0, nil, nil)
        buffer.deallocate()
        log("Device disconnected; waiting for reconnection.")
    }
}

let manager = IOHIDManagerCreate(kCFAllocatorDefault, IOOptionBits(kIOHIDOptionsTypeNone))
let matching: [String: Any] = [kIOHIDVendorIDKey: 5426, kIOHIDProductIDKey: 209]
IOHIDManagerSetDeviceMatching(manager, matching as CFDictionary)
IOHIDManagerRegisterDeviceMatchingCallback(manager, deviceCallback, statePtr)
IOHIDManagerRegisterDeviceRemovalCallback(manager, removalCallback, statePtr)
IOHIDManagerScheduleWithRunLoop(manager, CFRunLoopGetCurrent(), CFRunLoopMode.defaultMode.rawValue)

log("Razer Tilt Mapper started. LEFT -> Command+[; RIGHT -> Command+]")

// macOS presents its own permission dialogs. The user grants access in Settings.
// Launchd restarts this app while permissions are missing, so prompt only once per build.
let buildStamp = (try? FileManager.default.attributesOfItem(atPath: CommandLine.arguments[0])[.modificationDate] as? Date)
    .map { String($0.timeIntervalSince1970) } ?? ""
if !AXIsProcessTrusted() && UserDefaults.standard.string(forKey: "promptedBuild") != buildStamp {
    let options = [kAXTrustedCheckOptionPrompt.takeUnretainedValue() as String: true]
    _ = AXIsProcessTrustedWithOptions(options as CFDictionary)
    UserDefaults.standard.set(buildStamp, forKey: "promptedBuild")
}
if IOHIDCheckAccess(kIOHIDRequestTypeListenEvent) == kIOHIDAccessTypeUnknown {
    _ = IOHIDRequestAccess(kIOHIDRequestTypeListenEvent)
}

func accessGranted() -> Bool {
    IOHIDCheckAccess(kIOHIDRequestTypeListenEvent) == kIOHIDAccessTypeGranted && CGPreflightPostEventAccess()
}

func updateAccess() {
    let inputAllowed = IOHIDCheckAccess(kIOHIDRequestTypeListenEvent) == kIOHIDAccessTypeGranted
    let outputAllowed = CGPreflightPostEventAccess()
    var status = "Input Monitoring: \(inputAllowed ? "granted" : "required"); Accessibility: \(outputAllowed ? "granted" : "required")"
    if inputAllowed && !state.opened {
        let result = IOHIDManagerOpen(manager, IOOptionBits(kIOHIDOptionsTypeNone))
        state.opened = result == kIOReturnSuccess
        if !state.opened { status += "; HID open failed: \(result)" }
    }
    if state.opened { status += "; HID manager running" }
    if status != state.lastStatus {
        log(status)
        state.lastStatus = status
    }
}

updateAccess()
// A running process may not see a new grant, so exit while access is missing and let
// launchd (KeepAlive) start a fresh process that picks it up.
let accessTimer: Timer? = accessGranted() ? nil : {
    let started = Date()
    return Timer.scheduledTimer(withTimeInterval: 3, repeats: true) { _ in
        if accessGranted() {
            log("Permissions granted; restarting to apply them.")
            exit(0)
        }
        if Date().timeIntervalSince(started) >= 20 {
            log("Permissions still missing; restarting to recheck.")
            exit(0)
        }
    }
}()

// Invert vertical mouse-wheel scrolling. Trackpads and Magic Mouse send gesture phases;
// wheels do not, even when their high-resolution scrolling is marked continuous.
var scrollTap: CFMachPort?
// Mouse buttons whose down event was swallowed, so the matching drag/up is swallowed too.
var suppressedButtons = Set<Int64>()
// CG button numbers (0-based) of the tilt bits 0x08/0x10/0x20/0x40, i.e. buttons 4-7.
let tiltButtonNumbers: ClosedRange<Int64> = 3...6
let middleButtonNumber: Int64 = 2
// A wheel click this close to a tilt is the wheel being pressed while it is pushed sideways.
let tiltClickWindow: TimeInterval = 0.15

func tiltActive() -> Bool {
    if state.lastBits.values.contains(where: { $0 != 0 }) { return true }
    let now = ProcessInfo.processInfo.systemUptime
    return state.lastFired.values.contains { now - $0 < tiltClickWindow }
}

// Swallow the native button events behind a tilt: many apps treat buttons 4-7 like a
// middle click (e.g. open link in new tab), and tilting can also press the wheel button.
func shouldSuppress(_ type: CGEventType, _ event: CGEvent) -> Bool {
    let button = event.getIntegerValueField(.mouseEventButtonNumber)
    switch type {
    case .otherMouseDown:
        if tiltButtonNumbers.contains(button) || (button == middleButtonNumber && tiltActive()) {
            suppressedButtons.insert(button)
            return true
        }
        return false
    case .otherMouseDragged:
        return suppressedButtons.contains(button)
    case .otherMouseUp:
        return suppressedButtons.remove(button) != nil
    default:
        return false
    }
}
// Marks the replacement events this app posts so the tap lets them through untouched.
let invertedScrollMarker: Int64 = 0x52544D31
let scrollCallback: CGEventTapCallBack = { proxy, type, event, _ in
    if type == .tapDisabledByTimeout || type == .tapDisabledByUserInput {
        if let tap = scrollTap { CGEvent.tapEnable(tap: tap, enable: true) }
        return Unmanaged.passUnretained(event)
    }
    if shouldSuppress(type, event) { return nil }
    guard type == .scrollWheel,
          event.getIntegerValueField(.eventSourceUserData) != invertedScrollMarker
    else { return Unmanaged.passUnretained(event) }
    let phase = event.getIntegerValueField(.scrollWheelEventScrollPhase)
    let momentum = event.getIntegerValueField(.scrollWheelEventMomentumPhase)
    let isWheel = phase == 0 && momentum == 0
    let lineY = event.getIntegerValueField(.scrollWheelEventDeltaAxis1)
    let pointY = event.getIntegerValueField(.scrollWheelEventPointDeltaAxis1)
    let fixedY = event.getDoubleValueField(.scrollWheelEventFixedPtDeltaAxis1)
    guard isWheel, lineY != 0 || pointY != 0 || fixedY != 0 else { return Unmanaged.passUnretained(event) }
    let lineX = event.getIntegerValueField(.scrollWheelEventDeltaAxis2)
    let pointX = event.getIntegerValueField(.scrollWheelEventPointDeltaAxis2)
    let fixedX = event.getDoubleValueField(.scrollWheelEventFixedPtDeltaAxis2)
    // Apps read the hardware data attached to the original event, so editing its delta
    // fields has no visible effect. Post a fresh line-based event (like a real wheel, so
    // apps keep their per-line scroll distance) carrying the original deltas, Y flipped.
    guard let replacement = CGEvent(scrollWheelEvent2Source: nil, units: .line, wheelCount: 2,
                                    wheel1: Int32(clamping: -lineY), wheel2: Int32(clamping: lineX), wheel3: 0)
    else { return Unmanaged.passUnretained(event) }
    replacement.setIntegerValueField(.scrollWheelEventPointDeltaAxis1, value: -pointY)
    replacement.setIntegerValueField(.scrollWheelEventPointDeltaAxis2, value: pointX)
    replacement.setDoubleValueField(.scrollWheelEventFixedPtDeltaAxis1, value: -fixedY)
    replacement.setDoubleValueField(.scrollWheelEventFixedPtDeltaAxis2, value: fixedX)
    replacement.location = event.location
    replacement.flags = event.flags
    replacement.setIntegerValueField(.eventSourceUserData, value: invertedScrollMarker)
    replacement.tapPostEvent(proxy)
    return nil
}
if accessGranted() {
    scrollTap = CGEvent.tapCreate(tap: .cgSessionEventTap, place: .headInsertEventTap, options: .defaultTap,
                                  eventsOfInterest: [CGEventType.scrollWheel, .otherMouseDown, .otherMouseDragged, .otherMouseUp]
                                      .reduce(CGEventMask(0)) { $0 | CGEventMask(1 << $1.rawValue) },
                                  callback: scrollCallback, userInfo: nil)
    if let tap = scrollTap {
        CFRunLoopAddSource(CFRunLoopGetCurrent(), CFMachPortCreateRunLoopSource(kCFAllocatorDefault, tap, 0), .commonModes)
        log("Mouse wheel inversion running")
    } else {
        log("Mouse wheel inversion failed: could not create event tap")
    }
}
CFRunLoopRun()