 import OSAKit
 import IOKit
 
+// DisplayMode
 struct MyDisplayMode {
-  let mode:CGDisplayMode
-  let modeFromJSON:CGDisplayModeFromJSON
-  init(mode:CGDisplayMode, modeFromJSON:CGDisplayModeFromJSON) {
-    self.mode = mode
-    self.modeFromJSON = modeFromJSON
-  }
-}
-
+    let mode:CGDisplayMode
+    let modeFromJSON:CGDisplayModeFromJSON
+    init(mode:CGDisplayMode, modeFromJSON:CGDisplayModeFromJSON) {
+        self.mode = mode
+        self.modeFromJSON = modeFromJSON
+    }
+}
+
+// a custom data type that stores String or Double
 // from https://medium.com/swlh/f6ea6a2babf8
-// for dealing with CGDisplayModeFromJSON.RefreshRate, which return String or Double 0 in various settings
-// this can be retrieved by printing CGDisplayMode
-struct CGDisplayModeFromJSON: Decodable {
-  var BitsPerPixel = 32;
-  var BitsPerSample = 8;
-  var DepthFormat = 4;
-  var Height = 2160;
-  var IODisplayModeID = "-2147454976";
-  var IOFlags = 1048579;
-  var Mode = 0;
-  var PixelEncoding = "--------RRRRRRRRGGGGGGGGBBBBBBBB";
-  // making a type augmentation as program return String or 0 in various settings
-  var RefreshRate:StringOrDouble// = "30.00001525878906";
-  var SamplesPerPixel = 3;
-  var UsableForDesktopGUI = 1;
-  var Width = 3840;
-  var kCGDisplayBytesPerRow = 15360;
-  var kCGDisplayHorizontalResolution = 163;
-  var kCGDisplayModeIsInterlaced = 0;
-  var kCGDisplayModeIsSafeForHardware = 1;
-  var kCGDisplayModeIsStretched = 0;
-  var kCGDisplayModeIsTelevisionOutput = 1;
-  var kCGDisplayModeIsUnavailable = 0;
-  var kCGDisplayModeSuitableForUI = 1;
-  var kCGDisplayPixelsHigh = 2160;
-  var kCGDisplayPixelsWide = 3840;
-  var kCGDisplayResolution = 1;
-  var kCGDisplayVerticalResolution = 163;
-  var frequency:Int?
-}
-// use the object as key in a later process of mode filtering/winnowing
-extension CGDisplayModeFromJSON: Hashable & Equatable & Comparable {
-  // this is a hack to read private members in CGDisplayMode
-  static func decode(from modes:[CGDisplayMode], displayID:CGDirectDisplayID) -> [CGDisplayModeFromJSON]? {
-    do {
-      // first we operate the string to make it look like JSON String
-      let modesJSONText = try CGDisplayModeFromJSON.replaces(in: String(reflecting:modes), replacement: self.replacement)
-
-      // then we try to decode it as JSON
-      let decoder = JSONDecoder()
-      var modesFromJSON = try decoder.decode([CGDisplayModeFromJSON].self, from: modesJSONText.data(using: .utf8)!)
-
-      let displayRefreshRate = self.getAlternativeRefreshRate(by: displayID)
-      for i in 0..<modesFromJSON.count {
-        modesFromJSON[i].frequency = modesFromJSON[i].getRefreshRate(alternative:displayRefreshRate)
-      }
-
-      // preliminary integrity check
-      guard modesFromJSON.count == modes.count else { return nil }
-
-      return modesFromJSON
-    } catch {
-      // TODO: catch error and warn user properly
-      print(error)
-      return nil
-    }
-  }
-
-  static private func getAlternativeRefreshRate(by displayID:CGDirectDisplayID) -> Int {
+enum StringOrDouble: Decodable {
+    case double(Double)
+    case string(String)
+    
+    init(from decoder: Decoder) throws {
+        if let string = try? decoder.singleValueContainer().decode(String.self) {
+            self = .string(string)
+            return
+        }
+        if let double = try? decoder.singleValueContainer().decode(Double.self) {
+            self = .double(double)
+            return
+        }
+        throw Error.couldNotFindStringOrDouble
+    }
+    enum Error: Swift.Error {
+        case couldNotFindStringOrDouble
+    }
+}
+
+// get computed refreshrate
+fileprivate func _getAlternativeRefreshRate(by displayID:CGDirectDisplayID) -> Int {
     var link:CVDisplayLink?
     CVDisplayLinkCreateWithCGDisplay(displayID, &link)
     

}

func main() {
  let screens = Screens()

  let arguments = CommandLine.arguments
  let count = arguments.count
  guard count >= 2 else {
      seeHelp()
      return
  }
  switch arguments[1] {
  case "-l", "--list", "list":
      screens.listDisplays()
  case "-m", "--mode", "mode":
      var displayIndex = 0
      if count > 2, let index = Int(arguments[2]) {
          displayIndex = index
      }
      if displayIndex < screens.displayCount {
          print("Supported Modes for Display \(displayIndex):")
          screens.listModes(displayIndex)
      } else {
          print("Display index not found. List all available displays by:\n    screen-resolution-switcher -l")
      }
  case "-s", "--set", "set", "-r", "--set-retina", "retina":
      screens.set(with:DisplayUserSetting( arguments ))
  case "-d", "--toggle-dark-mode":
      DarkMode().toggle()
  case "-sl", "--sleep", "sleep":
      //if #available(macOS 12.0, *) {
          sleepDisplay()
      //}
  default:
      seeHelp()
  }
}

#if os(macOS)
  // run it
  main()
#else
  print("This script currently only runs on macOS")
#endif

// 6    id, width, scale
// 7    id, width, height, scale
struct DisplayUserSetting {
    static let MAX_SCALE = 10
    
    var displayIndex = 0, width = 0
    var height, scale:Int?
    init(_ arr:[String]) {
        var args = arr.compactMap { Int($0) }
        
        if args.count < 1 { return }
        
        if args[0] > Screens.MAX_DISPLAYS { args.insert(0 /* displayIndex */, at:0) }
        
        if args.count < 2 { return }
        
        displayIndex = args[0]
        width = args[1]
        
        if args.count == 2 { return }
        
        if args[2] > DisplayUserSetting.MAX_SCALE {
            height = args[2]
            if args.count > 3 { scale = args[3] }
        }
        else {
            scale = args[2]
            if args.count > 3 { height = args[3] }
        }
    }
    
    // override a lesser-used operator to simplify display mode checks
    static func == (lhs: Self, rhs: CGDisplayModeFromJSON) -> Bool {
        return rhs == lhs
    }
}

class Screens {
    // assume at most 8 display connected
    static let MAX_DISPLAYS = 8
    var maxDisplays = MAX_DISPLAYS
    // actual number of display
    var displayCount:Int = 0
    var dm = [DisplayManager]()
    
    init() {
        // actual number of display
        var displayCount32:UInt32 = 0
        var displayIDs = [CGDirectDisplayID](repeating: 0, count:maxDisplays)
        
        guard CGGetOnlineDisplayList(UInt32(maxDisplays), &displayIDs, &displayCount32) == .success else {
            print("Error on getting online display List.")
            return
        }
        displayCount = Int( displayCount32 )
        dm = displayIDs
            .filter { $0 != 0 }
            .map { DisplayManager($0) }
    }
    
    // print a list of all displays
    // used by -l
    func listDisplays() {
        for (i, m) in dm.enumerated() {
            m.printForOneDisplay("Display \(i):")
        }
    }
    
    func listModes(_ displayIndex:Int) {
        dm[displayIndex].printFormatForAllModes()
    }
    
    func set(with setting:DisplayUserSetting) {
        dm[setting.displayIndex].set(with:setting)
    }
}

// darkMode toggle code with JXA ;-)
// Method from Stackoverflow User: bacongravy
// https://stackoverflow.com/questions/44209057
struct DarkMode {
    static let scriptString = """
  pref = Application(\"System Events\").appearancePreferences
  pref.darkMode = !pref.darkMode()
"""
    let script = OSAScript.init(source: scriptString, language: OSALanguage.init(forName: "JavaScript"))
    
    init() {
        var compileError: NSDictionary?
        
        script.compileAndReturnError(&compileError)
    }
    func toggle() {
        var scriptError: NSDictionary?
        
        if let result = script.executeAndReturnError(&scriptError)?.stringValue { print("Dark Mode:", result) }
    }
}

func sleepDisplay() {
    let wrangler = strdup("IOService:/IOResources/IODisplayWrangler")
    // for legacy system
    // current compile threshold is a ballpark guess, tested on Catalina only
#if compiler(<5.5)
    let flag = kIOMasterPortDefault
#else
    let flag = kIOMainPortDefault
#endif
    // https://github.com/glfw/glfw/issues/1985
    let r = IORegistryEntryFromPath(flag, wrangler)
    IORegistryEntrySetCFProperty(r, ("IORequestIdle" as CFString), kCFBooleanTrue)
    IOObjectRelease(r)
}

func seeHelp() {
    print("""
Usage:
screen-resolution-switcher [-h|--help] [-l|--list|list] [-m|--mode|mode displayIndex]
[-s|--set|set displayIndex width scale] [-r|--set-retina|retina displayIndex width],
