#!/usr/bin/python3
# Reads a given directory and set a random
# wallpaper using the images on that dir.

import subprocess
import os
import os.path
import re
import random
import sys
import tempfile

class Display:
    index           = None
    name            = None
    offset          = None
    size            = None

    wall_filename   = None

    def __lt__(self, other):
        if self.offset[1] != other.offset[1]:
            return self.offset[1] < other.offset[1]
        return self.offset[0] < other.offset[0]

def detect_screen_geometry():
    screen      = None
    displays    = []

    p = subprocess.Popen('xrandr', stdout = subprocess.PIPE, text = True)
    for line in p.stdout:
        if line.startswith('Screen'):
            if screen:
                sys.stderr.write('Multiple X screens, not supported by this script\n')
                raise NotImplementedError

            m = re.search(r'current (\d*) x (\d*),', line)
            try:
                screen = [int(m.group(1)), int(m.group(2))]
            except:
                sys.stderr.write('Unsupported screen specification line: %s' % line)
                raise
        elif ' connected ' in line:
            m = re.match(r'(\S*) connected (\S+ )*(\d*)x(\d*)([+-]\d*)([+-]\d*)', line)
            if m:
                d = Display()

                d.name = m.group(1)
                d.offset = (int(m.group(5)), int(m.group(6)))
                d.size   = (int(m.group(3)), int(m.group(4)))

                if d.size[0] <= 0 or d.size[1] <= 0:
                    sys.stderr.write('Invalid non-positive display dimensions\n')
                    raise ValueError
                if d.offset[0] < 0 or d.offset[1] < 0:
                    sys.stderr.write('Negative display geometry offset, this is not supported\n')
                    raise NotImplementedError

                displays.append(d)

    if screen is None:
        sys.stderr.write('Missing screen specification in xrandr output\n')
        raise ValueError

    # sort the displays in the render scan order
    displays = sorted(displays)
    for i, d in enumerate(displays):
        d.index = i

    return screen, displays


if len(sys.argv) < 3:
    sys.stderr.write('Usage: %s <wallpaper_dir> <state dir> <display number 0> [<display number 1> ... ]\n'%(sys.argv[0]))
    sys.exit(-1)

# parse the commandline
wall_dir    = os.path.expandvars(sys.argv[1])
state_dir   = os.path.expandvars(sys.argv[2])

sys.stderr.write('Trying directory %s.\n'%wall_dir)
if not os.path.isdir(wall_dir):
    sys.stderr.write('Path %s is not a directory, aborting.\n'%wall_dir)
    sys.exit(-1)
if not os.path.isdir(state_dir):
    sys.stderr.write('Invalid state dir supplied, won\'t store/use any state\n')
    state_dir = None

screen, displays = detect_screen_geometry()

sys.stderr.write('The X screen resolution is %dx%d.\n' % (screen[0], screen[1]))
sys.stderr.write('The connected displays are:\n')
for i, d in enumerate(displays):
    sys.stderr.write('    %d: %s - %d x %d at (%d, %d)\n' % (d.index, d.name, d.size[0], d.size[1], d.offset[0], d.offset[1]))

# pick a new wallpaper for each specified display
files = os.listdir(wall_dir)
random.seed()

if len(sys.argv) >= 4 and sys.argv[3] == 'all':
    disp_indices = range(len(displays))
else:
    disp_indices = map(int, sys.argv[3:])

for d_idx in disp_indices:
    if d_idx >= len(displays):
        sys.stderr.write('Display index %d is too large, there are only %d connected displays\n'
                         % (d_idx, len(displays)))
        raise ValueError
    d = displays[d_idx]

    wallfile = random.choice(files)
    sys.stderr.write('Picking wall %s for display %d (%s)\n' % (wallfile, d.index, d.name))

    if state_dir:
        fname = os.path.join(state_dir, '%d.png' % d.index)
    else:
        d.tfile = tempfile.NamedTemporaryFile()
        fname   = d.tfile.name

    subprocess.check_call(['convert', '-size', '%dx%d' %(d.size[0], d.size[1]), 'canvas:black',
                           os.path.join(wall_dir, wallfile), '-resize', '%dx%d' % (d.size[0], d.size[1]),
                           '-gravity', 'center', '-composite', fname])
    d.wall_filename = fname

# compose the wallpaper file
tf = tempfile.NamedTemporaryFile()
cmdline = ['convert', '-size', '%dx%d' % (screen[0], screen[1]), 'canvas:black']
for d in displays:
    fname = d.wall_filename

    if fname is None and state_dir:
        tmpname = os.path.join(state_dir, '%d.png' % d.index)
        if os.path.isfile(tmpname):
            fname = tmpname

    if fname is not None:
        cmdline.extend([fname, '-geometry', '+%d+%d' %(d.offset[0], d.offset[1]), '-composite'])
cmdline.append('png:' + tf.name)
subprocess.check_call(cmdline)

# set the background
try:
    subprocess.check_call(['feh', '--bg-center', '--no-xinerama', tf.name])
except (subprocess.CalledProcessError, OSError):
    sys.stderr.write('Error setting the wallpaper. Is feh installed?\n')
    sys.exit(-1)

sys.exit(0)