#!/usr/bin/env python3

"""
grimpicker: a simple color picker for wlroots

Dependencies:
`slurp`: utility to select a region
`grim`: utility to make screenshots

Recommendations:
`wl-copy`: clipboard utility
`notify-send`: desktop notifications sender
`wayfreeze`: utility for freezing screen
"""

__pkgname__ = "grimpicker"
__version__ = "1.0.0"

import argparse
import signal
import subprocess
import sys
import time


class Color:
    escape = b"\x1b"
    reset_fg = escape + b"[39m"
    reset_bg = escape + b"[49m"
    reset_all = escape + b"[0m"

    escape_str = "\\e"
    reset_fg_str = escape_str + "[39m"
    reset_bg_str = escape_str + "[49m"
    reset_all_str = escape_str + "[0m"

    def __init__(self, r: int, g: int, b: int):
        (self.r, self.g, self.b) = (r, g, b)

    @classmethod
    def decode_ppm_pixel(cls, ppm: bytes):
        ppm_lines = ppm.splitlines()
        scales = ppm_lines[1].split(b" ")
        scale = int(scales[0])

        if not scale == int(scales[1]):
            raise ValueError("Unknown output scaling used")

        if not (
            len(ppm_lines) == 4
            and ppm_lines[0] == b"P6"
            and ppm_lines[2] == b"255"
            and len(ppm_lines[3]) == scale * scale * 3
        ):
            raise ValueError("only 1x1 pixel ppm P6 format without comments is supported, no HDR")

        #if we are dealing with multiple pixels, average them.
        r = 0
        g = 0
        b = 0
        for s in range(0, scale * scale):
            r += ppm_lines[3][s * 3 + 0]
            g += ppm_lines[3][s * 3 + 1]
            b += ppm_lines[3][s * 3 + 2]

        n = scale * scale
        return cls(r // n, g // n, b // n)

    def to_hex(self) -> str:
        return "#{:0>2X}{:0>2X}{:0>2X}".format(self.r, self.g, self.b)

    def to_escape_fg(self) -> bytes:
        return b"%b[38;2;%d;%d;%dm" % (self.escape, self.r, self.g, self.b)

    def to_escape_bg(self) -> bytes:
        return b"%b[48;2;%d;%d;%dm" % (self.escape, self.r, self.g, self.b)

    def to_escape_fg_str(self) -> str:
        return "{}[38;2;{};{};{}m".format(self.escape_str, self.r, self.g, self.b)

    def to_escape_bg_str(self) -> str:
        return "{}[48;2;{};{};{}m".format(self.escape_str, self.r, self.g, self.b)


def run(args) -> None:
    freeze: subprocess.Popen[bytes] | None = None
    if args.freeze:
        # Turn SIGHUP and SIGTERM into SystemExit so the finally below unfreezes the screen.
        for sig in (signal.SIGHUP, signal.SIGTERM):
            signal.signal(sig, lambda signum, _: sys.exit(128 + signum))
        try:
            freeze = subprocess.Popen(("wayfreeze", "--hide-cursor"), stdin=subprocess.DEVNULL,
                                      stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except FileNotFoundError:
            sys.exit("{}: wayfreeze not found".format(__pkgname__))

    try:
        if freeze:
            time.sleep(0.1)
        slurp = subprocess.check_output(("slurp", "-p"))
        grim = subprocess.check_output(("grim", "-g", "-", "-t", "ppm", "-"), input=slurp)
    finally:
        if freeze:
            freeze.terminate()

    color = Color.decode_ppm_pixel(grim)

    if not (args.print or args.draw or args.escape or args.copy or args.notify):
        args.print = True
        args.draw = True

    if args.print:
        print(color.to_hex())
    if args.draw:
        sys.stdout.buffer.write(color.to_escape_bg() + b" " * 7 + color.reset_bg + b"\n")
    if args.escape:
        sys.stdout.buffer.write(
            b"Truecolor terminal shell escape sequences:\n" +

            b"%bTo change foreground:%b " % (color.to_escape_fg(), color.reset_fg) +
            b'echo -e "%b", to reset: ' % color.to_escape_fg_str().encode() +
            b'echo -e "%b"\n' % color.reset_fg_str.encode() +

            b"%bTo change background:%b " % (color.to_escape_bg(), color.reset_bg) +
            b'echo -e "%b", to reset: ' % color.to_escape_bg_str().encode() +
            b'echo -e "%b"\n' % color.reset_bg_str.encode() +

            b'To reset all attributes: echo -e "%b"\n' % color.reset_all_str.encode()
        )
    if args.copy:
        subprocess.run(("wl-copy", color.to_hex()), check=True)
    if args.notify:
        subprocess.run(("notify-send", "--app-name", "grimpicker", color.to_hex()), check=True)


def parse_args() -> argparse.Namespace:
    usage = "{} [OPTIONS]".format(__pkgname__)
    version = "{} {}".format(__pkgname__, __version__)
    epilog = "See `man 1 grimpicker` for further details"
    parser = argparse.ArgumentParser(usage=usage, add_help=False, epilog=epilog)
    parser.add_argument("-p", "--print", dest="print", action="store_true", help="Print to stdout")
    parser.add_argument("-d", "--draw", dest="draw", action="store_true", help="Draw a colored block")
    parser.add_argument("-e", "--escape", dest="escape", action="store_true", help="Print shell escape sequences")
    parser.add_argument("-c", "--copy", dest="copy", action="store_true", help="Copy to clipboard")
    parser.add_argument("-f", "--freeze", dest="freeze", action="store_true", help="Freeze screen before picking")
    parser.add_argument("-n", "--notify", dest="notify", action="store_true", help="Send a notification")
    parser.add_argument("-h", "--help", action="help", help="Show help message and quit")
    parser.add_argument("-v", "--version", action="version", version=version, help="Show version number and quit")
    return parser.parse_args()


if __name__ == "__main__":
    run(parse_args())
