#!/usr/bin/env python
# -*- coding: utf-8 -*-

import io
import os
import subprocess
import tempfile
from contextlib import contextmanager
import shutil

from PIL import Image
import click
from marionette_driver.marionette import Marionette
from marionette_driver import expected, Wait, By


@contextmanager
def tempdir(*args, **kwargs):
    path = tempfile.mkdtemp(*args, **kwargs)
    try:
        yield path
    finally:
        shutil.rmtree(path)


@contextmanager
def process(args):
    process = subprocess.Popen(args)
    try:
        yield process
    finally:
        process.kill()


class Environ(object):
    def __getattr__(self, name):
        try:
            return os.environ[name]
        except KeyError:
            raise AttributeError(name)


class MarionetteWrapper(Marionette):
    def full_page_screenshot(self, filename):
        img = Image.open(io.BytesIO(self.screenshot(element=None, format='binary', full=True)))
        img.save(filename)

    @property
    def wait_until_expected(self):
        return WaitUntil(self)

    @property
    def env(self):
        return Environ()


class WaitUntilExpected:
    def __init__(self, client, name):
        self.client = client
        self.name = name

    def by_name(self, name):
        return self(By.NAME, name)

    def by_css_selector(self, selector):
        return self(By.CSS_SELECTOR, selector)

    def __call__(self, *args, **kwargs):
        return Wait(self.client).until(getattr(expected, self.name)(*args, **kwargs))


class WaitUntil:
    def __init__(self, client):
        self.client = client

    def __getattr__(self, name):
        return WaitUntilExpected(self.client, name)


@click.command()
@click.argument('script_py')
def screenshot(script_py):
    global client

    with tempdir() as profile_path:
        with process(['firefox', '--marionette', '--no-remote', '--new-instance', '--profile', profile_path, '--headless']):
            client = MarionetteWrapper('localhost', '2828')
            client.start_session()
            with open(script_py) as fd:
                exec(fd.read(), {})
            client.cleanup()


if __name__ == '__main__':
    screenshot()
