#import <Cocoa/Cocoa.h>
#import <ApplicationServices/ApplicationServices.h>
#import <Carbon/Carbon.h>
#import <stdatomic.h>
#import <pthread.h>
#import "darwin.h"

// Go exports
extern void chr_onSelectionTrigger(int kind);
extern void chr_onEvent(int kind, double x, double y);

@interface CHRPanel : NSPanel
@end

@implementation CHRPanel
- (BOOL)canBecomeKeyWindow {
    return NO;
}
- (BOOL)canBecomeMainWindow {
    return NO;
}
@end

@interface CHRView : NSView
@property (nonatomic, assign) NSRect hitRect;
@property (nonatomic, assign) BOOL hasEntered;
@property (nonatomic, strong) NSTrackingArea *trackingArea;

- (void)updateHit:(NSRect)hit;
@end

@implementation CHRView

- (BOOL)isFlipped {
    return YES;
}

- (NSView *)hitTest:(NSPoint)point {
    NSPoint local = [self convertPoint:point fromView:self.superview];
    if (NSPointInRect(local, self.hitRect)) {
        return self;
    }
    return nil;
}

- (void)updateHit:(NSRect)hit {
    self.hitRect = hit;
    if (self.trackingArea) {
        [self removeTrackingArea:self.trackingArea];
        self.trackingArea = nil;
    }
    NSTrackingAreaOptions opts = NSTrackingMouseEnteredAndExited |
                                 NSTrackingMouseMoved |
                                 NSTrackingActiveAlways;
    self.trackingArea = [[NSTrackingArea alloc] initWithRect:self.hitRect
                                                     options:opts
                                                       owner:self
                                                    userInfo:nil];
    [self addTrackingArea:self.trackingArea];
}

- (void)mouseUp:(NSEvent *)event {
    NSPoint loc = [self convertPoint:[event locationInWindow] fromView:nil];
    if (NSPointInRect(loc, self.hitRect)) {
        chr_onEvent(0, loc.x, loc.y); // Click
    }
}

- (void)mouseMoved:(NSEvent *)event {
    NSPoint loc = [self convertPoint:[event locationInWindow] fromView:nil];
    if (NSPointInRect(loc, self.hitRect)) {
        chr_onEvent(1, loc.x, loc.y); // Motion
    }
}

- (void)mouseEntered:(NSEvent *)event {
    self.hasEntered = YES;
}

- (void)mouseExited:(NSEvent *)event {
    if (self.hasEntered) {
        self.hasEntered = NO;
        chr_onEvent(2, 0, 0); // Leave
    }
}

- (void)mouseDown:(NSEvent *)event {
    NSPoint loc = [self convertPoint:[event locationInWindow] fromView:nil];
    if (!NSPointInRect(loc, self.hitRect)) {
        chr_onEvent(3, 0, 0); // Dismiss if clicked on shadow
    }
}

@end

static CHRPanel *gPanel = nil;
static CHRView *gView = nil;
static NSRect gHitOnScreen = {{0, 0}, {0, 0}};

static atomic_bool gAppIsRunning = false;

typedef struct {
    NSRect frame;
    NSRect visibleFrame;
    double scale;
} CHRScreenInfo;

#define MAX_SCREENS 16
static CHRScreenInfo gScreens[MAX_SCREENS];
static int gScreenCount = 0;
static double gPrimaryScreenHeight = 1080.0;
static pthread_mutex_t gScreenLock = PTHREAD_MUTEX_INITIALIZER;

static id gMouseUpMonitor = nil;
static id gMouseDownMonitor = nil;
static id gKeyUpMonitor = nil;
static id gGlobalClickMonitor = nil;
static id gLocalClickMonitor = nil;

static NSPoint gMouseDownLoc = {0, 0};
static NSInteger gMouseDownClickCount = 0;

// pids we have already asked to build their accessibility tree
static NSMutableSet<NSNumber *> *gWoken = nil;

static void updateScreenCache(void) {
    NSArray<NSScreen *> *screens = [NSScreen screens];
    pthread_mutex_lock(&gScreenLock);
    gScreenCount = (int)MIN(screens.count, MAX_SCREENS);
    if (screens.count > 0) {
        gPrimaryScreenHeight = screens[0].frame.size.height;
    }
    for (int i = 0; i < gScreenCount; i++) {
        NSScreen *s = screens[i];
        gScreens[i].frame = s.frame;
        gScreens[i].visibleFrame = s.visibleFrame;
        gScreens[i].scale = s.backingScaleFactor;
    }
    pthread_mutex_unlock(&gScreenLock);
}

void chr_init(void) {
    @autoreleasepool {
        [NSApplication sharedApplication];
        [NSApp setActivationPolicy:NSApplicationActivationPolicyAccessory];

        updateScreenCache();

        [[NSNotificationCenter defaultCenter] addObserverForName:NSApplicationDidChangeScreenParametersNotification
                                                          object:nil
                                                           queue:[NSOperationQueue mainQueue]
                                                      usingBlock:^(NSNotification * _Nonnull note) {
            updateScreenCache();
        }];

        gPanel = [[CHRPanel alloc] initWithContentRect:NSZeroRect
                                             styleMask:NSWindowStyleMaskBorderless | NSWindowStyleMaskNonactivatingPanel
                                               backing:NSBackingStoreBuffered
                                                 defer:NO];
        [gPanel setLevel:NSPopUpMenuWindowLevel];
        [gPanel setOpaque:NO];
        [gPanel setBackgroundColor:[NSColor clearColor]];
        [gPanel setHasShadow:NO];
        [gPanel setCollectionBehavior:NSWindowCollectionBehaviorCanJoinAllSpaces |
                                      NSWindowCollectionBehaviorFullScreenAuxiliary |
                                      NSWindowCollectionBehaviorStationary];
        [gPanel setBecomesKeyOnlyIfNeeded:YES];
        [gPanel setHidesOnDeactivate:NO];
        [gPanel setAcceptsMouseMovedEvents:YES];
        [gPanel setIgnoresMouseEvents:NO];

        gView = [[CHRView alloc] initWithFrame:NSZeroRect];
        [gView setWantsLayer:YES];
        [gPanel setContentView:gView];

        gWoken = [[NSMutableSet alloc] init];
    }
}

void chr_run(void) {
    gAppIsRunning = true;
    [NSApp run];
    gAppIsRunning = false;
}

void chr_stop(void) {
    dispatch_async(dispatch_get_main_queue(), ^{
        gAppIsRunning = false;
        [NSApp stop:nil];
        NSEvent *event = [NSEvent otherEventWithType:NSEventTypeApplicationDefined
                                           location:NSMakePoint(0, 0)
                                      modifierFlags:0
                                          timestamp:0
                                       windowNumber:0
                                            context:nil
                                            subtype:0
                                              data1:0
                                              data2:0];
        [NSApp postEvent:event atStart:YES];
    });
}

void chr_installMonitors(void) {
    @autoreleasepool {
        gMouseDownMonitor = [NSEvent addGlobalMonitorForEventsMatchingMask:NSEventMaskLeftMouseDown
                                                                   handler:^(NSEvent *e) {
            gMouseDownLoc = [NSEvent mouseLocation];
            gMouseDownClickCount = [e clickCount];
        }];

        gMouseUpMonitor = [NSEvent addGlobalMonitorForEventsMatchingMask:NSEventMaskLeftMouseUp
                                                                 handler:^(NSEvent *e) {
            NSPoint mouseUpLoc = [NSEvent mouseLocation];
            double dist = hypot(mouseUpLoc.x - gMouseDownLoc.x, mouseUpLoc.y - gMouseDownLoc.y);
            bool isDragOrMulti = (dist > 3.0) || (gMouseDownClickCount >= 2);
            chr_onSelectionTrigger(isDragOrMulti ? 1 : 0);
        }];

        gKeyUpMonitor = [NSEvent addGlobalMonitorForEventsMatchingMask:NSEventMaskKeyUp
                                                               handler:^(NSEvent *e) {
            NSEventModifierFlags flags = [e modifierFlags];
            unsigned short kc = [e keyCode];
            bool isShiftNav = ((flags & NSEventModifierFlagShift) != 0) ||
                              (kc >= 123 && kc <= 126) ||
                              kc == 115 || kc == 119 || kc == 116 || kc == 121;
            bool isCmdA = ((flags & NSEventModifierFlagCommand) != 0) && (kc == 0);
            if (isShiftNav || isCmdA) {
                chr_onSelectionTrigger(2);
            }
        }];

        gGlobalClickMonitor = [NSEvent addGlobalMonitorForEventsMatchingMask:NSEventMaskLeftMouseDown | NSEventMaskKeyDown
                                                                     handler:^(NSEvent *e) {
            if (e.type == NSEventTypeLeftMouseDown) {
                if (gPanel && [gPanel isVisible]) {
                    NSPoint loc = [NSEvent mouseLocation];
                    if (!NSPointInRect(loc, gHitOnScreen)) {
                        chr_onEvent(3, 0, 0); // Dismiss
                    }
                }
            } else if (e.type == NSEventTypeKeyDown) {
                if ([e keyCode] == 53) { // Escape
                    if (gPanel && [gPanel isVisible]) {
                        chr_onEvent(3, 0, 0); // Dismiss
                    }
                }
            }
        }];

        gLocalClickMonitor = [NSEvent addLocalMonitorForEventsMatchingMask:NSEventMaskKeyDown
                                                                   handler:^NSEvent *(NSEvent *e) {
            if ([e keyCode] == 53) { // Escape
                if (gPanel && [gPanel isVisible]) {
                    chr_onEvent(3, 0, 0); // Dismiss
                }
            }
            return e;
        }];
    }
}

void chr_uninstallMonitors(void) {
    @autoreleasepool {
        if (gMouseDownMonitor) {
            [NSEvent removeMonitor:gMouseDownMonitor];
            gMouseDownMonitor = nil;
        }
        if (gMouseUpMonitor) {
            [NSEvent removeMonitor:gMouseUpMonitor];
            gMouseUpMonitor = nil;
        }
        if (gKeyUpMonitor) {
            [NSEvent removeMonitor:gKeyUpMonitor];
            gKeyUpMonitor = nil;
        }
        if (gGlobalClickMonitor) {
            [NSEvent removeMonitor:gGlobalClickMonitor];
            gGlobalClickMonitor = nil;
        }
        if (gLocalClickMonitor) {
            [NSEvent removeMonitor:gLocalClickMonitor];
            gLocalClickMonitor = nil;
        }
    }
}

bool chr_checkAccessibility(bool prompt) {
    @autoreleasepool {
        NSDictionary *options = @{(__bridge id)kAXTrustedCheckOptionPrompt: @(prompt)};
        return AXIsProcessTrustedWithOptions((__bridge CFDictionaryRef)options);
    }
}

// wakeAccessibility asks an app to build its accessibility tree. Chromium
// browsers (AXEnhancedUserInterface) and Electron apps (AXManualAccessibility)
// only build it when an assistive app asks; before that, Chromium returns the
// AXApplication itself as the focused element with no selected text.
static void wakeAccessibility(pid_t pid, AXUIElementRef app) {
    if (!AXIsProcessTrusted()) return;
    NSNumber *n = @(pid);
    @synchronized (gWoken) {
        if ([gWoken containsObject:n]) return;
        if (gWoken.count > 64) [gWoken removeAllObjects];
        [gWoken addObject:n];
    }
    AXUIElementSetAttributeValue(app, CFSTR("AXManualAccessibility"), kCFBooleanTrue);
    AXUIElementSetAttributeValue(app, CFSTR("AXEnhancedUserInterface"), kCFBooleanTrue);
}

static NSString *stringValue(CFTypeRef v) {
    if (!v) return nil;
    if (CFGetTypeID(v) == CFStringGetTypeID()) return [(__bridge NSString *)v copy];
    if (CFGetTypeID(v) == CFAttributedStringGetTypeID()) return [[(__bridge NSAttributedString *)v string] copy];
    return nil;
}

// selectedText reads the selection of a focused element: AXSelectedText for
// text views and fields, else the selected text marker range that WebKit and
// Chromium use for page content.
static NSString *selectedText(AXUIElementRef el) {
    CFTypeRef v = NULL;
    NSString *s = nil;
    if (AXUIElementCopyAttributeValue(el, kAXSelectedTextAttribute, &v) == kAXErrorSuccess) {
        s = stringValue(v);
    }
    if (v) CFRelease(v);
    if (s.length > 0) return s;

    CFTypeRef range = NULL;
    if (AXUIElementCopyAttributeValue(el, CFSTR("AXSelectedTextMarkerRange"), &range) != kAXErrorSuccess || !range) {
        return s;
    }
    v = NULL;
    if (AXUIElementCopyParameterizedAttributeValue(el, CFSTR("AXStringForTextMarkerRange"), range, &v) == kAXErrorSuccess) {
        s = stringValue(v);
    }
    if (v) CFRelease(v);
    CFRelease(range);
    return s;
}

static NSString *chr_copyFallback(void) {
    if (!gAppIsRunning || ![NSApp isRunning]) return nil;

    __block NSString *result = nil;
    void (^task)(void) = ^{
        @autoreleasepool {
            NSPasteboard *pb = [NSPasteboard generalPasteboard];
            NSInteger oldCount = [pb changeCount];
            NSString *prevString = [pb stringForType:NSPasteboardTypeString];

            CGEventSourceRef src = CGEventSourceCreate(kCGEventSourceStateHIDSystemState);
            if (src) {
                CGEventRef cDown = CGEventCreateKeyboardEvent(src, (CGKeyCode)8, true);
                CGEventRef cUp = CGEventCreateKeyboardEvent(src, (CGKeyCode)8, false);
                if (cDown && cUp) {
                    CGEventSetFlags(cDown, kCGEventFlagMaskCommand);
                    CGEventSetFlags(cUp, kCGEventFlagMaskCommand);
                    CGEventPost(kCGHIDEventTap, cDown);
                    CGEventPost(kCGHIDEventTap, cUp);
                }
                if (cDown) CFRelease(cDown);
                if (cUp) CFRelease(cUp);
                CFRelease(src);
            }

            usleep(80000); // 80 ms

            if ([pb changeCount] != oldCount) {
                result = [[pb stringForType:NSPasteboardTypeString] copy];
                [pb clearContents];
                if (prevString && prevString.length > 0) {
                    [pb setString:prevString forType:NSPasteboardTypeString];
                }
            }
        }
    };

    if ([NSThread isMainThread]) {
        task();
    } else {
        dispatch_sync(dispatch_get_main_queue(), task);
    }
    return result;
}

char *chr_getSelectedText(bool allowCopyFallback) {
    @autoreleasepool {
        AXUIElementRef focused = NULL;
        AXUIElementRef app = NULL;
        pid_t pid = 0;
        NSRunningApplication *front = [[NSWorkspace sharedWorkspace] frontmostApplication];
        if (front) {
            pid = [front processIdentifier];
            app = AXUIElementCreateApplication(pid);
            if (app) {
                AXUIElementCopyAttributeValue(app, kAXFocusedUIElementAttribute, (CFTypeRef *)&focused);
            }
        }
        if (!focused) {
            AXUIElementRef sys = AXUIElementCreateSystemWide();
            if (sys) {
                AXUIElementCopyAttributeValue(sys, kAXFocusedUIElementAttribute, (CFTypeRef *)&focused);
                CFRelease(sys);
            }
        }

        NSString *s = nil;
        if (focused) {
            s = selectedText(focused);
            CFRelease(focused);
        }
        if (s.length == 0 && app) {
            wakeAccessibility(pid, app); // the next read will find the tree
        }
        if (app) {
            CFRelease(app);
        }
        if (s.length == 0 && allowCopyFallback) {
            s = chr_copyFallback();
        }
        return s ? strdup([s UTF8String]) : NULL;
    }
}

CHRAnchor chr_getAnchor(void) {
    @autoreleasepool {
        NSPoint mouseLoc = [NSEvent mouseLocation];
        CHRAnchor a = {0};

        pthread_mutex_lock(&gScreenLock);
        double primaryH = gPrimaryScreenHeight > 0 ? gPrimaryScreenHeight : 1080.0;
        a.x = mouseLoc.x;
        a.y = primaryH - mouseLoc.y;

        int targetIdx = -1;
        for (int i = 0; i < gScreenCount; i++) {
            if (NSPointInRect(mouseLoc, gScreens[i].frame)) {
                targetIdx = i;
                break;
            }
        }
        if (targetIdx < 0 && gScreenCount > 0) {
            targetIdx = 0;
        }

        if (targetIdx >= 0) {
            NSRect vf = gScreens[targetIdx].visibleFrame;
            a.screenX = vf.origin.x;
            a.screenY = primaryH - (vf.origin.y + vf.size.height);
            a.screenW = vf.size.width;
            a.screenH = vf.size.height;
            a.scale = gScreens[targetIdx].scale;
        } else {
            a.screenX = 0; a.screenY = 0;
            a.screenW = 1920; a.screenH = 1080;
            a.scale = 2.0;
        }
        pthread_mutex_unlock(&gScreenLock);
        return a;
    }
}

void chr_getPointer(double *outX, double *outY) {
    @autoreleasepool {
        NSPoint mouseLoc = [NSEvent mouseLocation];
        pthread_mutex_lock(&gScreenLock);
        double primaryH = gPrimaryScreenHeight > 0 ? gPrimaryScreenHeight : 1080.0;
        pthread_mutex_unlock(&gScreenLock);
        *outX = mouseLoc.x;
        *outY = primaryH - mouseLoc.y;
    }
}

void chr_showPopup(const void *bytes, int stride, int width, int height,
                   double x, double y,
                   double hitX, double hitY, double hitW, double hitH,
                   double contentsScale) {
    @autoreleasepool {
        size_t byteCount = (size_t)stride * height;
        CFDataRef data = CFDataCreate(kCFAllocatorDefault, (const UInt8 *)bytes, byteCount);
        if (!data) return;

        dispatch_async(dispatch_get_main_queue(), ^{
            @autoreleasepool {
                CGColorSpaceRef cs = CGColorSpaceCreateDeviceRGB();
                CGDataProviderRef provider = CGDataProviderCreateWithCFData(data);
                CGImageRef img = CGImageCreate(
                    width, height, 8, 32, stride, cs,
                    kCGImageAlphaPremultipliedLast | kCGBitmapByteOrderDefault,
                    provider, NULL, false, kCGRenderingIntentDefault
                );
                CFRelease(provider);
                CFRelease(data);
                CGColorSpaceRelease(cs);

                if (!img) return;

                pthread_mutex_lock(&gScreenLock);
                double primaryH = gPrimaryScreenHeight > 0 ? gPrimaryScreenHeight : 1080.0;
                pthread_mutex_unlock(&gScreenLock);

                double logicalW = (double)width / contentsScale;
                double logicalH = (double)height / contentsScale;
                double appkitY = primaryH - (y + logicalH);

                NSRect frame = NSMakeRect(x, appkitY, logicalW, logicalH);
                [gPanel setFrame:frame display:YES];
                [gView setFrame:NSMakeRect(0, 0, logicalW, logicalH)];

                gView.layer.contentsScale = contentsScale;
                gView.layer.contents = (__bridge id)img;
                CGImageRelease(img);

                NSRect hit = NSMakeRect(hitX, hitY, hitW, hitH);
                [gView updateHit:hit];

                gHitOnScreen = NSMakeRect(frame.origin.x + hitX,
                                          frame.origin.y + (frame.size.height - (hitY + hitH)),
                                          hitW, hitH);

                NSPoint mouseLoc = [NSEvent mouseLocation];
                gView.hasEntered = NSPointInRect(mouseLoc, gHitOnScreen);

                if (![gPanel isVisible]) {
                    [gPanel setAlphaValue:0.0];
                    [gPanel orderFront:nil];
                    [NSAnimationContext runAnimationGroup:^(NSAnimationContext *context) {
                        [context setDuration:0.090];
                        [[gPanel animator] setAlphaValue:1.0];
                    } completionHandler:nil];
                }
            }
        });
    }
}

void chr_hidePopup(void) {
    dispatch_async(dispatch_get_main_queue(), ^{
        @autoreleasepool {
            if (gPanel) {
                [gPanel orderOut:nil];
            }
            if (gView) {
                gView.layer.contents = nil;
                gView.hasEntered = NO;
            }
            gHitOnScreen = NSMakeRect(0, 0, 0, 0);
        }
    });
}

bool chr_savePopupSnapshot(const char *path) {
    if (!gAppIsRunning || ![NSApp isRunning]) return false;

    __block bool ok = false;
    void (^task)(void) = ^{
        @autoreleasepool {
            if (!gView || !gView.layer) return;
            NSRect bounds = gView.bounds;
            if (bounds.size.width <= 0 || bounds.size.height <= 0) return;
            NSBitmapImageRep *rep = [gView bitmapImageRepForCachingDisplayInRect:bounds];
            [gView cacheDisplayInRect:bounds toBitmapImageRep:rep];
            NSData *png = [rep representationUsingType:NSBitmapImageFileTypePNG properties:@{}];
            if (png) {
                NSString *p = [NSString stringWithUTF8String:path];
                ok = [png writeToFile:p atomically:YES];
            }
        }
    };
    if ([NSThread isMainThread]) {
        task();
    } else {
        dispatch_sync(dispatch_get_main_queue(), task);
    }
    return ok;
}