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"
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"

proc saveToClipboard(filename: string) =
  let cmd = "cat " & filename & " | xclip -selection clipboar -t image/png"
  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 = SCREENSHOT_CMD.replace("%f",filename)
  case screenshot_type:
    of "window":
      cmd = cmd.replace("%s",ACTIVE_WINDOW_CMD)
    of "region":
      cmd = cmd.replace("%s",REGION_FLAG)
    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:
  for arg in args:
    if arg in TYPES:
      screenshot_type = arg
      break
  if screenshot_type == "":
    showScreenshotTypeSel()
  if screenshot_type != "":
    takeScreenshot()