3 files changed, 61 insertions, 33 deletions

diff --git a/Invert.h b/Invert.h
index c84685e..91a9f69 100644
--- a/Invert.h
+++ b/Invert.h
@@ -1,6 +1,9 @@
 #import <Cocoa/Cocoa.h>
+#import <CoreGraphics/CoreGraphics.h>
 
-@interface Invert : NSObject
+@interface Invert : NSObject {
+    BOOL _inverted;
+}
 
 - (void)toggleInvertColors:(NSEvent *)event;
 
diff --git a/Invert.m b/Invert.m
index a41eb46..b5675eb 100644
--- a/Invert.m
+++ b/Invert.m
@@ -1,10 +1,67 @@
 #import "Invert.h"
 
+#define MAX_DISPLAYS 8
+
 @implementation Invert
 
+- (id)init
+{
+    self = [super init];
+    if (self) {
+        _inverted = NO;
+    }
+    return self;
+}
+
 - (void)toggleInvertColors:(NSEvent *)event
 {
     NSLog(@"Toggling");
+
+    if (!_inverted) {
+        [self invertColors];
+
+        _inverted = YES;
+    }
+    else {
+        [self restoreColors];
+
+        _inverted = NO;
+    }
+}
+
+- (void)invertColors
+{
+    const CGGammaValue inverted_gamma[2] = {1, 0};
+
+    CGDirectDisplayID active_displays[MAX_DISPLAYS];
+    uint32_t display_count;
+
+    CGError error = CGGetActiveDisplayList(
+        MAX_DISPLAYS,
+        &active_displays[0],
+        &display_count
+    );
+    if (error != kCGErrorSuccess) {
+        // return 69;
+    }
+
+    for (int i = 0; i < display_count; i++) {
+        error = CGSetDisplayTransferByTable(
+            active_displays[i],
+            2,
+            inverted_gamma,
+            inverted_gamma,
+            inverted_gamma
+        );
+        if (error != kCGErrorSuccess) {
+            // TODO: error handling
+        }
+    }
+}
+
+- (void)restoreColors
+{
+    CGDisplayRestoreColorSyncSettings();
 }
 
 @end
diff --git a/main.m b/main.m
index 8ffc3b2..a5db0a9 100644
--- a/main.m
+++ b/main.m
@@ -1,11 +1,8 @@
 #import <Carbon/Carbon.h>
-#import <CoreGraphics/CoreGraphics.h>
 
 #import "Invert.h"
 #import "DDHotKeyCenter.h"
 
-#define MAX_DISPLAYS 8
-
 int main(int argc, const char * argv[]) {
     [NSApplication sharedApplication];
 
@@ -23,37 +20,8 @@ int main(int argc, const char * argv[]) {
         NSLog(@"Error registering hotkey");
     }
 
-    const CGGammaValue inverted_gamma[2] = {1, 0};
-
-    CGDirectDisplayID active_displays[MAX_DISPLAYS];
-    uint32_t display_count;
-
-    CGError error = CGGetActiveDisplayList(
-        MAX_DISPLAYS,
-        &active_displays[0],
-        &display_count
-    );
-    if (error != kCGErrorSuccess) {
-        return 69;
-    }
-
-    for (int i = 0; i < display_count; i++) {
-        error = CGSetDisplayTransferByTable(
-            active_displays[i],
-            2,
-            inverted_gamma,
-            inverted_gamma,
-            inverted_gamma
-        );
-        if (error != kCGErrorSuccess) {
-            // TODO: error handling
-        }
-    }
-
     // for (;;) {}
     // sleep(4);
-    //
-    // CGDisplayRestoreColorSyncSettings();
 
     // TODO: trap?
     // [invert release];