#!/usr/bin/python3
import argparse
import hashlib
import os
import screeninfo
import shutil
import subprocess
import sys
import tempfile
import time

try:
    from PIL import Image, ImageDraw, ImageFont, ImageFilter, ImageOps
except ImportError:
    print('Please install Pillow.', file=sys.stderr)
    exit(1)

def download(url, path):
    import urllib.request
    import ssl

    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    u = urllib.request.urlopen(url, context=ctx)
    meta = u.info()
    if meta.get('Content-Length') is not None:
        file_size = int(meta.get('Content-Length'))
    else:
        file_size = None
    print("Downloading %s, size = %s..." % (path, file_size))

    file_size_dl = 0
    block_sz = 8192
    with open(path, 'wb') as f:
        while True:
            buffer = u.read(block_sz)
            if not buffer:
                break

            file_size_dl += len(buffer)
            f.write(buffer)
            if file_size is not None:
                status = r"%10d  [%3.2f%%]" % (
                    file_size_dl, file_size_dl * 100. / file_size)
            else:
                status = r"%10d" % file_size_dl
            status = status + chr(8)*(len(status)+1)
            print(status)

def cover_or_fit(img, size_dst, quality, cover=False):
    ratio_src = img.size[0] / img.size[1]
    ratio_dst = size_dst[0] / size_dst[1]
    if (ratio_dst > ratio_src) == cover:
        img = img.resize((size_dst[0], int(size_dst[0] // ratio_src)), quality)
    else:
        img = img.resize((int(size_dst[1] * ratio_src), size_dst[1]), quality)
    return img

def cover(img, size_dst, quality=Image.ANTIALIAS):
    return cover_or_fit(img, size_dst, quality, cover=True)

def fit(img, size_dst, quality=Image.ANTIALIAS):
    return cover_or_fit(img, size_dst, quality, cover=False)


class WallpaperChanger(object):
    def __init__(self):
        monitors = screeninfo.get_monitors()
        self.monitors = dict((i + 1, m) for i, m in enumerate(monitors))

    def validate_monitor_number(self, number):
        if number not in self.monitors:
            raise RuntimeError('Monitor number out of range')

    def get_path(self, number):
        return os.path.expanduser('~/.wallpaper%d' % number)

    def get_paths(self):
        return [self.get_path(m) for m in self.monitors.keys]

    def get_temp_path(self, text='tmp'):
        hash = hashlib.md5(str(text).encode('utf-8')).hexdigest()
        return os.path.join(tempfile.gettempdir(), 'wallpaper-' + hash + '.tmp')

    def set(self, args):
        gaps = [args.gap[i] for i in (3, 0, 1, 2)] #u r d l --> l u r d
        self.validate_monitor_number(args.monitor)
        mon_size = (
            self.monitors[args.monitor].width,
            self.monitors[args.monitor].height)
        work_area = (
            mon_size[0] - 2 * args.border_size - gaps[0] - gaps[2],
            mon_size[1] - 2 * args.border_size - gaps[1] - gaps[3])

        target_img = Image.new('RGB', mon_size, args.background)

        if args.path is not None:
            if args.path.startswith('http'):
                source_path = self.get_temp_path(args.path)
                if not os.path.exists(source_path):
                    download(args.path, source_path)
            else:
                source_path = args.path
            if not os.path.exists(source_path):
                raise RuntimeError(source_path + ' does not exist')

            source_img = Image.open(source_path)
            if args.crop:
                crop_area = (
                    int(source_img.size[0] * args.crop[0]),
                    int(source_img.size[1] * args.crop[1]),
                    int(source_img.size[0] * args.crop[2]),
                    int(source_img.size[1] * args.crop[3]))
                source_img = source_img.crop(crop_area)

            if args.fit:
                source_img = fit(source_img, work_area)
            elif args.cover:
                source_img = cover(source_img, work_area)

            resize_area = (
                int(source_img.size[0] * args.scale),
                int(source_img.size[1] * args.scale))

            source_img = source_img.resize(resize_area, Image.ANTIALIAS)

            freedom_area = (
                work_area[0] - resize_area[0],
                work_area[1] - resize_area[1])

            target_position = (
                int(gaps[0] + freedom_area[0] * args.translate[0]),
                int(gaps[1] + freedom_area[1] * args.translate[1]))

            if args.border_size:
                source_img = ImageOps.expand(
                    source_img, border=args.border_size, fill=args.border_color)

            target_img.paste(source_img, target_position, source_img)

        if args.preview:
            target_path = self.get_temp_path()
        else:
            target_path = self.get_path(args.monitor)
        target_img.save(target_path, 'PNG')

        if args.preview:
            paths = self.get_paths()
            paths[args.monitor] = target_path
            self.render(paths)
            time.sleep(1)
        self.refresh()

    def swap(self, args):
        if args.first == args.second:
            return

        self.validate_monitor_number(args.first)
        self.validate_monitor_number(args.second)

        temp_path = self.get_temp_path()
        shutil.copy(self.get_path(args.second), temp_path)
        os.rename(
            self.get_path(args.first),
            self.get_path(args.second))
        shutil.copy(temp_path, self.get_path(args.first))
        os.unlink(temp_path)
        self.refresh()

    def identify(self, args):
        if 'cygwin' in sys.platform:
            font_path = 'C:\\windows\\fonts\\Arial.ttf'
        else:
            font_path = 'arial.ttf'
        font = ImageFont.truetype(font_path, 200, encoding='unic')
        paths = []
        for m, info in self.monitors.items():
            path_src = self.get_path(m)
            path_dst = self.get_temp_path(m)
            paths.append(path_dst)

            img = Image.open(path_src)
            text = str(m)

            text_img = Image.new('RGBA', (400, 400), (0, 0, 0, 0))
            draw = ImageDraw.Draw(text_img)
            text_size = draw.textsize(text, font=font)
            pos = [(text_img.size[i] - text_size[i]) // 2 for i in (0,1)]
            draw.text(pos, text, font=font, fill='#000000')

            blurred_text_img = text_img.filter(ImageFilter.BLUR)
            ImageDraw.Draw(blurred_text_img).text(
                pos, text, font=font, fill='#ffffff')

            pos = tuple((img.size[i] - text_img.size[i]) // 2 for i in (0,1))
            img.paste(blurred_text_img, pos, blurred_text_img)
            img.save(path_dst, 'JPEG', quality=100)

        self.render(paths)
        time.sleep(1)
        self.refresh()
        for path in paths:
            os.unlink(path)

    def render(self, paths):
        raise NotImplementedError()

    def refresh(self):
        self.render([self.get_path(i) for i in self.monitors])

class WindowsWallpaperChanger(WallpaperChanger):
    @staticmethod
    def detect():
        return 'cygwin' in sys.platform

    def render(self, paths):
        full_path = 'C:\\windows\\temp\\wallpaper.bmp'
        total_width = max(m.x + m.width for m in self.monitors.values())
        total_height = max(m.y + m.height for m in self.monitors.values())
        full_img = Image.new('RGBA', (total_width, total_height), (0, 0, 0, 0))
        for i, monitor in enumerate(self.monitors.values()):
            sub_img = Image.open(paths[i])
            sub_img = sub_img.resize((monitor.width, monitor.height))
            full_img.paste(sub_img, (monitor.x, monitor.y))
        full_img.save(full_path)
        ps_cmd = [
            '''Add-Type @"
using System;
using System.Runtime.InteropServices;
using Microsoft.Win32;
namespace Wallpaper {
public class Setter {
    public const int SetDesktopWallpaper = 20;
    public const int UpdateIniFile = 0x01;
    public const int SendWinIniChange = 0x02;

    [DllImport("user32.dll", SetLastError = true, CharSet = CharSet.Auto)]
    private static extern int SystemParametersInfo(
        int uAction, int uParam, string lpvParam, int fuWinIni);

    public static void SetWallpaper(string path) {
        SystemParametersInfo(
            SetDesktopWallpaper, 0, path, UpdateIniFile | SendWinIniChange);
        RegistryKey key = Registry.CurrentUser.OpenSubKey(
            "Control Panel\\\Desktop", true);
        key.SetValue("WallpaperStyle", "1");
        key.SetValue("TileWallpaper", "1");
        key.Close();
    }
}
}''' + "\n\"@",
'[Wallpaper.Setter]::SetWallpaper(\'' + full_path + '\')']
        self.run_powershell(ps_cmd)

    def run_powershell(self, cmd):
        ps = subprocess.Popen(
            ['powershell', '-command', ';'.join(cmd)],
            stdout=subprocess.PIPE)
        return ps.stdout.read().decode('utf-8')

class X11WallpaperChanger(WallpaperChanger):
    @staticmethod
    def detect():
        return 'DISPLAY' in os.environ

    def render(self, paths):
        subprocess.Popen(['feh', '--bg-fill'] + list(paths)).wait()

def add_set_arg_parser(parent_parser, fmt, wc):
    p = parent_parser.add_parser(
        'set', help='set a wallpaper', formatter_class=fmt)
    p.add_argument(
        '-p', '--preview', action='store_true',
        help='preview the changes without applying them')
    p.add_argument(
        '-s', '--scale', type=float, metavar='RATIO', default=1,
        help='scale image by given factor')
    p.add_argument(
        '-f', '--fit', action='store_true',
        help='fit the image within monitor work area')
    p.add_argument(
        '-c', '--cover', action='store_true',
        help='make the image cover whole monitor work area')
    p.add_argument(
        '-b', '--background', metavar='COLOR', default='#000000',
        help='select background color')
    p.add_argument(
        '-t', '--translate', type=float, nargs=2, metavar=('X', 'Y'),
        default=(0.5, 0.5),
        help='place the image at given position')
    p.add_argument(
        '-bs', '--border-size', type=int, metavar='SIZE', default=0,
        help='select border width')
    p.add_argument(
        '-bc', '--border-color', metavar='COLOR', default='black',
        help='select border size')
    p.add_argument(
        '--crop', type=float, nargs=4, metavar=('X1', 'Y1', 'X2', 'Y2'),
        default=(0, 0, 1, 1),
        help='select crop area')
    p.add_argument(
        '-g', '--gap', '--gaps', type=int, nargs=4,
        metavar=('U', 'R', 'D', 'L'),
        default=(0, 0, 0, 0),
        help='keep "border" around work area')
    p.add_argument(
        'monitor', metavar='MONITOR', type=int,
        help='set the monitor number')
    p.add_argument(
        'path', metavar='PATH', nargs='?',
        help='set the input image path')
    p.set_defaults(func=wc.set)

def add_identify_arg_parser(parent_parser, fmt, wc):
    p = parent_parser.add_parser(
        'identify', help='briefly display monitor numbers on the desktop')
    p.set_defaults(func=wc.identify)

def add_swap_arg_parser(parent_parser, fmt, wc):
    p = parent_parser.add_parser(
        'swap', help='swap wallpapers between monitors')
    p.add_argument(
        'first', metavar='FIRST', type=int, nargs='?', default=1,
        help='first monitor number')
    p.add_argument(
        'second', metavar='SECOND', type=int, nargs='?', default=2,
        help='second monitor number')
    p.set_defaults(func=wc.swap)

class CustomHelpFormatter(argparse.HelpFormatter):
    def __init__(self, prog):
        super().__init__(prog, max_help_position=40, width=80)

    def _format_action_invocation(self, action):
        if not action.option_strings:
            return super()._format_action_invocation(action)
        default = self._get_default_metavar_for_optional(action)
        values = self._format_args(action, default)
        opts = action.option_strings
        short = [arg for arg in opts if not arg.startswith('--')]
        long = [arg for arg in opts if arg.startswith('--')]
        return '%*s%s%s %s' % (
            3, ', '.join(short), ', ' if short else '  ', ', '.join(long),
            values)

def get_arg_parser(wc):
    fmt = lambda prog: CustomHelpFormatter(prog)
    parser = argparse.ArgumentParser(
        description='Wallpaper utility', formatter_class=fmt)
    subparsers = parser.add_subparsers(help='choose the subcommand')
    add_set_arg_parser(subparsers, fmt, wc)
    add_swap_arg_parser(subparsers, fmt, wc)
    add_identify_arg_parser(subparsers, fmt, wc)
    return parser

def main():
    changers = [WindowsWallpaperChanger, X11WallpaperChanger]
    chosen = None
    for changer in changers:
        if changer.detect():
            chosen = changer()
            break
    if chosen is None:
        print('This graphical environment is not supported.', file=sys.stderr)
        return
    parser = get_arg_parser(chosen)
    args = parser.parse_args()
    if 'func' not in args:
        chosen.refresh()
    else:
        args.func(args)

if __name__ == '__main__':
    main()
