#!/usr/bin/env python3
import os
import sys
import subprocess
import tempfile
from pathlib import Path
from datetime import datetime

def check_cmds(cmds):
    for cmd in cmds:
        if subprocess.call(["which", cmd], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) != 0:
            print(f"Missing: {cmd}", file=sys.stderr)
            sys.exit(1)

def show_help():
    print(f"""Usage: {sys.argv[0]} [color|full|window] [output_dir]\nEnvironment options:\n  WATERMARK=1\n  WATERMARK_TEXT='text'\n  WATERMARK_POS=position\n  WATERMARK_SIZE=px\n  WATERMARK_BG='#rrggbbaa'\n""")

def post(file):
    if os.environ.get("WATERMARK") == "1":
        check_cmds(["magick"])
        watermark_text = os.environ.get("WATERMARK_TEXT", datetime.now().strftime("%Y-%m-%d %H:%M"))
        watermark_pos = os.environ.get("WATERMARK_POS", "southeast")
        watermark_size = os.environ.get("WATERMARK_SIZE", "28")
        watermark_bg = os.environ.get("WATERMARK_BG", "#00000080")
        subprocess.run([
            "magick", file,
            "-gravity", watermark_pos,
            "-pointsize", watermark_size,
            "-fill", "white",
            "-undercolor", watermark_bg,
            "-annotate", "+20+20", watermark_text,
            file
        ], check=True)
    subprocess.run(["xclip", "-selection", "clipboard", "-t", "image/png", "-i", file], check=True)
    subprocess.run(["notify-send", "-i", file, f"Screenshot: {os.path.basename(file)}"], check=True)

def colorpicker(tmp):
    check_cmds(["magick"])
    if subprocess.call(["which", "slop"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) == 0:
        slop = subprocess.run(["slop", "--tolerance=0"], capture_output=True, text=True)
        if slop.returncode != 0:
            sys.exit(1)
        crop = slop.stdout.strip()
        subprocess.run(["import", "-window", "root", "-crop", crop, tmp], check=True)
    else:
        subprocess.run(["import", tmp], check=True)
    try:
        hexcolor = subprocess.check_output([
            "magick", tmp, "-scale", "1x1!", "-format", "#%[hex:p{0,0}]", "info:"
        ], stderr=subprocess.DEVNULL, text=True).strip()
    except subprocess.CalledProcessError:
        hexcolor = "#??????"
    subprocess.run(["xclip", "-selection", "clipboard", "-i"], input=hexcolor.encode(), check=True)
    subprocess.run(["notify-send", "-i", tmp, f"Color: {hexcolor}"], check=True)

def capture_full(file):
    subprocess.run(["import", "-window", "root", file], check=True)
    post(file)

def capture_window(file):
    check_cmds(["xdotool"])
    winid = subprocess.check_output(["xdotool", "getwindowfocus", "-f"], text=True).strip()
    subprocess.run(["import", "-window", winid, file], check=True)
    post(file)

def capture_selection(file):
    subprocess.run(["import", file], check=True)
    post(file)

def main():
    check_cmds(["xclip", "import", "notify-send"])
    outdir = Path(sys.argv[2]) if len(sys.argv) > 2 else Path.home() / "files" / "pics" / "screenies"
    outdir.mkdir(parents=True, exist_ok=True)
    file = str(outdir / (datetime.now().strftime("%Y%m%d_%H%M%S") + ".png"))
    with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as tmp:
        if len(sys.argv) > 1:
            arg = sys.argv[1]
        else:
            arg = ""
        if arg in ("-h", "--help"):
            show_help()
        elif arg.startswith("color"):
            colorpicker(tmp.name)
        elif arg == "full":
            capture_full(file)
        elif arg == "window":
            capture_window(file)
        else:
            capture_selection(file)

if __name__ == "__main__":
    main()