import base
import std/[times,os,osproc,strutils,sequtils]

var screenshot_type = ""
const TYPES = @["region", "fullscreen", "window"]
const DATE_FORMAT = "yyyyMMdd-hhmmss"
const FILENAME = "Screenshot-%d.png"
const TEMP_DIR = "/tmp/"
const SCREENSHOT_CMD = "maim -u %s --format png %f"
const SCREENSHOT_CMD_WL = "grim %s %f"
var RUN_CMD = SCREENSHOT_CMD
let DATE_STR = now().format(DATE_FORMAT)
# where %s is an extra flag or process, i.e. xdotool for getting active window
const ACTIVE_WINDOW_CMD = "-i $(xdotool getactivewindow)"
const REGION_FLAG = "-s"
const REGION_FLAG_WL = "-g \"$(slurp)\""
var REGION_CMD = REGION_FLAG
const CLIPBOARD_CMD = "xclip -selection clipboard -t image/png"
const CLIPBOARD_CMD_WL = "wl-copy"
var CLIP_CMD = CLIPBOARD_CMD

proc saveToClipboard(filename: string) =
  let cmd = "cat " & filename & " | " & CLIP_CMD
  let status = execCmd(cmd)
  if status == 0 and fileExists(filename):
    removeFile(filename)
  return

proc saveToFile(filename: string) =
  if fileExists(filename):
    let new_filename = filename.replace("/tmp/", getHomeDir() & "Screenshots/")
    copyFile(filename, new_filename)
    if fileExists(new_filename):
      removeFile(filename)
  return

proc openFile(filename: string) =
  let cmd = "xdg-open " & filename
  discard execCmd(cmd)
  return

proc showScreenshotSaveSel(filename: string) =
  let info = newInfo("Save Screenshot")
  let args = @["clipboard", "save", "open", "---", "exit"]
  let choice = outputData(info,args)
  if choice == "---":
    showScreenshotSaveSel(filename)
  elif choice == "exit":
    return
  elif choice in args:
    case choice:
      of "clipboard":
        saveToClipboard(filename)
      of "save":
        saveToFile(filename)
      of "open":
        openFile(filename)
  return

proc showScreenshotTypeSel() =
  let info = newInfo("Screenshot type")
  let args = concat(TYPES,@["---","exit"])
  let choice = outputData(info,args)
  if choice in TYPES:
    screenshot_type = choice
  elif choice == "---":
    showScreenshotTypeSel()
  elif choice == "exit":
    return
  return

proc takeScreenshot() =
  let filename = TEMP_DIR & FILENAME.replace("%d",DATE_STR)
  var cmd = RUN_CMD.replace("%f",filename)
  case screenshot_type:
    of "window":
      cmd = cmd.replace("%s",ACTIVE_WINDOW_CMD)
    of "region":
      cmd = cmd.replace("%s",REGION_CMD)
    else: #fullscreen
      cmd = cmd.replace("%s","")
      # sleep for a bit otherwise the screen shot grabs dmenu as well
      sleep(1*500)

  let status = execCmd(cmd)
  if status == 0:
    showScreenshotSaveSel(filename)
  return

if isMainModule:
  if wayland:
    RUN_CMD = SCREENSHOT_CMD_WL
    REGION_CMD = REGION_FLAG_WL
    CLIP_CMD = CLIPBOARD_CMD_WL
  for arg in args:
    if arg in TYPES:
      screenshot_type = arg
      break
  if screenshot_type == "":
    showScreenshotTypeSel()
  if screenshot_type != "":
    takeScreenshot()