#!/usr/bin/env python3

import os
import random
import subprocess
import sys

try:
    from PIL import Image, ImageDraw
except ImportError:
    print("wp: missing dependency: pillow (pip install pillow)", file=sys.stderr)
    sys.exit(1)

HOME = os.environ["HOME"]
DIR = f"{HOME}/img/screen"
os.makedirs(DIR, exist_ok=True)


def get_resolution():
    output = "1920x1080"
    if "WAYLAND_DISPLAY" in os.environ:
        desktop = os.environ.get("XDG_CURRENT_DESKTOP", "")
        if desktop == "Hyprland":
            output = (
                subprocess.check_output(
                    "hyprctl monitors -j | jq -r '.[] | select(.focused==true) | .width, .height' | paste -sd x -",
                    shell=True,
                )
                .decode()
                .strip()
            )
        elif desktop == "sway":
            output = (
                subprocess.check_output(
                    "swaymsg -t get_outputs -r | jq -r '.[] | select(.active==true) | .current_mode.width,.current_mode.height' | paste -sd x -",
                    shell=True,
                )
                .decode()
                .strip()
            )
    else:
        output = (
            subprocess.check_output(
                "xrandr | grep '*' | awk '{print $1}'", shell=True
            )
            .decode()
            .strip()
        )
    return map(int, output.split("x"))


def lock():
    W, H = get_resolution()
    UNIT = 16
    grid_w = W // UNIT
    grid_h = H // UNIT

    bg_color = tuple(random.randint(0, 255) for _ in range(3))
    img = Image.new("RGB", (W, H), bg_color)
    draw = ImageDraw.Draw(img)

    MU = 0.8
    SIGMA = 0.7
    SCALE_FACTOR = 1.8

    S = {(x, y) for x in range(grid_w) for y in range(grid_h)}
    N = len(S)

    while S:
        if random.random() < 1 / N:
            break
        gx, gy = S.pop()
        w_units = max(
            1,
            min(
                grid_w - gx,
                int(round(random.lognormvariate(MU, SIGMA) * SCALE_FACTOR)),
            ),
        )
        h_units = max(
            1,
            min(
                grid_h - gy,
                int(round(random.lognormvariate(MU, SIGMA) * SCALE_FACTOR)),
            ),
        )
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

        dx = random.choice([1, -1])
        dy = random.choice([1, -1])

        gx_end = gx + dx * (w_units - 1)
        gy_end = gy + dy * (h_units - 1)

        x0, x1 = sorted([gx, gx_end])
        y0, y1 = sorted([gy, gy_end])

        draw.rectangle(
            [x0 * UNIT, y0 * UNIT, (x1 + 1) * UNIT - 1, (y1 + 1) * UNIT - 1],
            fill=color,
        )

    img.save(f"{DIR}/lock.jpg", quality=95)


def wall():
    W, H = get_resolution()

    colors = [
        "#aa0000",
        "#ff3333",
        "#ff7777",
        "#ffbb55",
        "#ffcc88",
        "#ff88ff",
        "#aaaaff",
        "#77ddbb",
        "#77ff77",
        "#cccccc",
    ]

    colors.reverse()

    img = Image.new("RGB", (W, H), "white")
    draw = ImageDraw.Draw(img)

    num_bars = len(colors)

    for i, color in enumerate(reversed(colors)):
        top = round(i * H / num_bars)
        bottom = round((i + 1) * H / num_bars)
        draw.rectangle([0, top, W, bottom], fill=color)

    img.save(f"{DIR}/wallpaper.jpg", quality=95)


if len(sys.argv) < 2:
    print("Usage: wp {lock|wall}", file=sys.stderr)
    sys.exit(1)

cmd = sys.argv[1]
if cmd == "lock":
    lock()
elif cmd == "wall":
    wall()
else:
    print("Usage: wp {lock|wall}", file=sys.stderr)
    sys.exit(1)