#!/usr/bin/env python

from __future__ import print_function

import re
import os
import sys
import time
import json
import signal
import fnmatch
import argparse
import traceback
import subprocess

from subprocess import CalledProcessError
from os.path import exists, abspath, realpath, relpath, dirname, basename, join, split, splitext

import yaml
import jinja2

import requests
from requests.exceptions import ConnectionError

import colorama
colorama.init()

from conda.utils import human_bytes

try:
    import gevent
except ImportError:
    use_server = False
else:
    use_server = True

try:
    import IPython
except ImportError:
    use_notebook = False
else:
    use_notebook = True

py2 = sys.version_info.major == 2

if py2:
    import boto
    from boto.s3.key import Key as S3Key
    from boto.exception import NoAuthHandlerFound

s3_bucket = "bokeh-travis"
s3 = "https://s3.amazonaws.com/%s" % s3_bucket

default_timeout = int(os.environ.get("BOKEH_DEFAULT_TIMEOUT", 10))
default_diff = os.environ.get("BOKEH_DEFAULT_DIFF", None)
default_upload = default_diff is not None

parser = argparse.ArgumentParser(description="Automated testing of Bokeh's examples")
parser.add_argument("patterns", type=str, nargs="*",
                    help="select a subset of examples to test")
parser.add_argument("-p", "--bokeh-port", type=int, default=5006,
                    help="port on which Bokeh server resides")
parser.add_argument("-n", "--ipython-port", type=int, default=6007,
                    help="port on which IPython Notebook server resides")
parser.add_argument("-j", "--phantomjs", type=str, default="phantomjs",
                    help="phantomjs executable")
parser.add_argument("-t", "--timeout", type=int, default=default_timeout,
                    help="how long can an example run (in seconds)")
parser.add_argument("-v", "--verbose", action="store_true", default=False,
                    help="show console messages")
parser.add_argument("-D", "--no-dev", dest="dev", action="store_false", default=True,
                    help="don't use development JavaScript and CSS files")
parser.add_argument("-all-n", "--all-notebooks", action="store_true", default=False,
                    help="test all the notebooks inside examples/plotting/notebook folder.")
parser.add_argument("-o", "--output-cells", type=str, choices=['complain', 'remove', 'ignore'], default='complain',
                    help="what to do with notebooks' output cells")
parser.add_argument("-l", "--log-file", type=argparse.FileType('w'), default='examples.log',
                    help="where to write the complete log")
parser.add_argument("-d", "--diff", type=str, default=default_diff,
                    help="compare generated images against this ref")
parser.add_argument("-u", "--upload", dest="upload", action="store_true", default=default_upload,
                    help="upload generated images as reference images to S3")

args = parser.parse_args()

def write(*values, **kwargs):
    end = kwargs.get('end', '\n')
    print(*values, end=end)
    print(*values, end=end, file=args.log_file)
    args.log_file.flush()

def red(text):
    return "%s%s%s" % (colorama.Fore.RED, text, colorama.Style.RESET_ALL)

def yellow(text):
    return "%s%s%s" % (colorama.Fore.YELLOW, text, colorama.Style.RESET_ALL)

def blue(text):
    return "%s%s%s" % (colorama.Fore.BLUE, text, colorama.Style.RESET_ALL)

def green(text):
    return "%s%s%s" % (colorama.Fore.GREEN, text, colorama.Style.RESET_ALL)

def fail(msg=None):
    msg = " " + msg if msg is not None else ""
    write("%s%s" % (red("[FAIL]"), msg))

def warn(msg=None):
    msg = " " + msg if msg is not None else ""
    write("%s%s" % (yellow("[WARN]"), msg))

def info(msg=None):
    msg = " " + msg if msg is not None else ""
    write("%s%s" % ("[INFO]", msg))

def ok(msg=None):
    msg = " " + msg if msg is not None else ""
    write("%s%s" % (green("[OK]"), msg))

DEFAULT_NO_DEV = os.environ.get('BOKEH_DEFAULT_NO_DEV', 'false')
if DEFAULT_NO_DEV.lower() in ["true", "yes", "on", "1"]:
    args.dev = False
elif DEFAULT_NO_DEV.lower() not in ["false", "no", "off", "0"]:
    write("unexpected value for BOKEH_DEFAULT_NO_DEV environmental variable, %s" % DEFAULT_NO_DEV)
    sys.exit(1)

write("%s Using configuration:" % green(">>>"))
write("%s patterns      = %s" % (green("---"), args.patterns))
write("%s bokeh-port    = %s" % (green("---"), args.bokeh_port))
write("%s ipython-port  = %s" % (green("---"), args.ipython_port))
write("%s phantomjs     = %s" % (green("---"), args.phantomjs))
write("%s timeout       = %s" % (green("---"), args.timeout))
write("%s verbose       = %s" % (green("---"), args.verbose))
write("%s dev           = %s" % (green("---"), args.dev))
write("%s all-notebooks = %s" % (green("---"), args.all_notebooks))
write("%s output-cells  = %s" % (green("---"), args.output_cells))
write("%s log-file      = %s" % (green("---"), args.log_file.name))
write("%s diff          = %s" % (green("---"), args.diff))
write("%s upload        = %s" % (green("---"), args.upload))

def get_version_from_git(ref=None):
    cmd = ["git", "describe", "--tags", "--always"]

    if ref is not None:
        cmd += ref

    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        code = proc.wait()
    except OSError:
        write("Failed to run: %s" % " ".join(cmd))
        sys.exit(1)

    if code != 0:
        write("Failed to get version for %s" % ref)
        sys.exit(1)

    return proc.stdout.read().decode('utf-8').strip()

__version__ = get_version_from_git()
write("%s Bokeh version: %s" % (green(">>>"), __version__))

if args.diff:
    cmd = ["git", "describe", "--tags", "--always", args.diff]

    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        code = proc.wait()
    except OSError:
        write("Failed to run: %s" % " ".join(cmd))
        sys.exit(1)

    if code != 0:
        write("Failed to process --diff=%s" % args.diff)
        sys.exit(1)

    args.diff = proc.stdout.read().decode('utf-8').strip()
    write("%s Reference version: %s" % (green(">>>"), args.diff))

if args.upload and not py2:
    warn("Can't upload to S3, becase boto is not available under Python 3.")
    warn("Proceeding as --no-upload was specified on the command line.")
    args.upload = False

if args.upload:
    try:
        boto.connect_s3()
    except NoAuthHandlerFound:
        write("Upload (--upload) was requested but could not connect to S3.")
        write("Fix your credentials or permissions and re-run tests.")
        sys.exit(1)

    if __version__.endswith("-dirty"):
        write("Can't upload reference images when working directory is dirty.")
        write("Make sure that __version__ doesn't contain -dirty suffix.")
        sys.exit(1)

def is_selected(example):
    if not args.patterns:
        return True
    elif any(pattern in example for pattern in args.patterns):
        return True
    elif any(fnmatch.fnmatch(example, pattern) for pattern in args.patterns):
        return True
    else:
        return False

base_dir = dirname(__file__)
examples = []

class Flags(object):
    file     = 1<<1
    server   = 1<<2
    notebook = 1<<3
    animated = 1<<4
    skip     = 1<<5

def example_type(flags):
    if flags & Flags.file: return "file"
    elif flags & Flags.server: return "server"
    elif flags & Flags.notebook: return "notebook"

def add_examples(examples_dir, example_type=None, skip=None):
    examples_path = join(base_dir, examples_dir)

    if skip is not None:
        skip = set(skip)

    for file in os.listdir(examples_path):
        flags = 0

        if file.startswith(('_', '.')):
            continue
        elif file.endswith(".py"):
            if example_type is not None:
                flags |= example_type
            elif "server" in file or "animate" in file:
                flags |= Flags.server
            else:
                flags |= Flags.file
        elif file.endswith(".ipynb"):
            flags |= Flags.notebook
        else:
            continue

        if "animate" in file:
            flags |= Flags.animated

            if flags & Flags.file:
                raise ValueError("file examples can't be animated")

        if skip and file in skip:
            flags |= Flags.skip

        examples.append((join(examples_path, file), flags))

def detect_examples():
    with open(join(dirname(__file__), "test.yaml"), "r") as f:
        examples = yaml.load(f.read())

    for example in examples:
        path = example["path"]

        try:
            example_type = getattr(Flags, example["type"])
        except KeyError:
            example_type = None

        if not args.all_notebooks:
            skip = example.get("skip") or example.get("skip_travis")
        else:
            skip = example.get("skip")

        add_examples(path, example_type=example_type, skip=skip)

def make_env():
    env = os.environ.copy()
    env['BOKEH_RESOURCES'] = 'relative'
    env['BOKEH_BROWSER'] = 'none'
    if args.dev:
        env['BOKEH_RESOURCES'] += '-dev'
        env['BOKEH_PRETTY'] = 'yes'
    return env

class Timeout(Exception):
    pass

def run_example(example):
    cmd = ["python", basename(example)]
    cwd = dirname(example)
    env = make_env()

    def alarm_handler(sig, frame):
        raise Timeout

    signal.signal(signal.SIGALRM, alarm_handler)
    signal.alarm(args.timeout)

    try:
        proc = subprocess.Popen(cmd, cwd=cwd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        try:
            def dump(file):
                for line in iter(file.readline, b""):
                    write(line.decode("utf-8"), end="")

            dump(proc.stdout)
            dump(proc.stderr)

            return proc.wait()
        except KeyboardInterrupt:
            proc.kill()
            raise
    except Timeout:
        warn("Timeout")
        proc.kill()
        return 0
    finally:
        signal.alarm(0)

def get_path_parts(path):
    parts = []
    while True:
        newpath, tail = split(path)
        parts.append(tail)
        path = newpath
        if tail == 'examples':
            break
    parts.reverse()
    return parts

def test_example(example, flags):
    no_ext = splitext(abspath(example))[0]
    url_path = join(*get_path_parts(abspath(example)))

    if flags & Flags.file:
        html_file = "%s.html" % no_ext
        url = 'file://' + html_file
    elif flags & Flags.server:
        server_url = 'http://localhost:%d/bokeh/doc/%s/show'
        url = server_url % (args.bokeh_port, basename(no_ext))
    elif flags & Flags.notebook:
        notebook_url = 'http://localhost:%d/notebooks/%s'
        url = notebook_url % (args.ipython_port, url_path)
    else:
        raise ValueError("invalid example type: %s" % example_type(flags))

    png_file = "%s-%s.png" % (no_ext, __version__)

    cmd = [args.phantomjs, join(base_dir, "test.js"),
           example_type(flags), url, png_file, str(args.timeout)]

    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        code = proc.wait()
    except OSError:
        write("Failed to run: %s" % " ".join(cmd))
        sys.exit(1)

    result = json.loads(proc.stdout.read().decode("utf-8"))

    status = result['status']
    errors = result['errors']
    messages = result['messages']
    resources = result['resources']

    if status == 'fail':
        fail("failed to load %s" % url)
    else:
        if args.verbose:
            for message in messages:
                msg = message['msg']
                line = message.get('line')
                source = message.get('source')

                if source is None:
                    write(msg)
                elif line is None:
                    write("%s: %s" % (source, msg))
                else:
                    write("%s:%s: %s" % (source, line, msg))

        resource_errors = False

        for resource in resources:
            url = resource['url']

            if url.endswith(".png"):
                fn, color = warn, yellow
            else:
                fn, color, resource_errors = fail, red, True

            fn("%s: %s (%s)" % (url, color(resource['status']), resource['statusText']))

        for error in errors:
            write(error['msg'])

            for item in error['trace']:
                write("    %s: %d" % (item['file'], item['line']))

        if resource_errors or errors:
            fail(example)
        else:
            if args.diff:
                example_path = relpath(splitext(example)[0], base_dir)
                ref_loc = join(args.diff, example_path + ".png")
                ref_url = join(s3, ref_loc)
                response = requests.get(ref_url)

                if not response.ok:
                    info("referece image %s doesn't exist" % ref_url)
                else:
                    ref_png_file = abspath(join(dirname(__file__), "refs", ref_loc))
                    diff_png_file = splitext(png_file)[0] + "-diff.png"

                    ref_png_path = dirname(ref_png_file)
                    if not exists(ref_png_path):
                        os.makedirs(ref_png_path)

                    with open(ref_png_file, "wb") as f:
                        f.write(response.content)

                    cmd = ["perceptualdiff", "-output", diff_png_file, png_file, ref_png_file]

                    try:
                        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                        code = proc.wait()
                    except OSError:
                        write("Failed to run: %s" % " ".join(cmd))
                        sys.exit(1)

                    if code != 0:
                        warn("generated and reference images differ")
                        warn("generated: " + png_file)
                        warn("reference: " + ref_png_file)
                        warn("diff: " + diff_png_file)

            ok(example)
            return True

    return False

def test_nbexecuter():
    if use_notebook and need_notebook:
        import nbexecuter
        example_nbconverted = join("examples", "glyphs", "glyph.ipynb")
        try:
            nbexecuter.main(example_nbconverted)
        except IOError as e:
            example_nbconverted = example_nbconverted[example_nbconverted.find('glyphs'):]
            nbexecuter.main(example_nbconverted)
        write(example_nbconverted)

def start_bokeh_server():
    cmd = ["python", "-c", "import bokeh.server; bokeh.server.run()"]
    argv = ["--bokeh-port=%s" % args.bokeh_port, "--backend=memory"]

    if args.dev:
        argv.extend(["--splitjs", "--debugjs", "--filter-logs"])

    try:
        proc = subprocess.Popen(cmd + argv, stdout=args.log_file, stderr=args.log_file)
    except OSError:
        write("Failed to run: %s" % " ".join(cmd + argv))
        sys.exit(1)
    else:
        def wait_until(func, timeout=5.0, interval=0.01):
            start = time.time()

            while True:
                if func():
                    return True
                if time.time() - start > timeout:
                    return False
                time.sleep(interval)

        def wait_for_bokeh_server():
            def helper():
                try:
                    return requests.get('http://localhost:%s/bokeh/ping' % args.bokeh_port)
                except ConnectionError:
                    return False

            return wait_until(helper)

        if not wait_for_bokeh_server():
            write("Timeout when running: %s" % " ".join(cmd + argv))
            sys.exit(1)

        return proc

def deal_with_output_cells(example):
    def load_nb(example):
        with open(example, "r") as f:
            return json.load(f)

    def save_nb(example, nb):
        with open(example, "w") as f:
            json.dump(nb, f, ident=2, sort_keys=True)

    def bad_code_cells(nb):
        for worksheet in nb.get("worksheets", []):
            for cell in worksheet.get("cells", []):
                if cell.get("cell_type") == "code":
                    outputs = cell.get("outputs")

                    if isinstance(outputs, list) and len(outputs) > 0:
                        yield cell, outputs

    def complain(fn):
        fn("%s notebook contains output cells" % example)

    if args.output_cells == 'complain':
        nb = load_nb(example)

        for _, _ in bad_code_cells(nb):
            complain(fail)
            return False
    elif args.output_cells == 'remove':
        nb = load_nb(example)
        changes = False

        for cell, _ in bad_code_cells(nb):
            cell["outputs"] = []
            changes = True

        if changes:
            complain(warn)
            save_nb(example, nb)

    return True

def create_ipython_profile():
    create = ["ipython", "profile", "create", "bokeh_test"]
    try:
        subprocess.check_call(create)
    except OSError:
        write("Failed to run: %s" % " ".join(create))
        sys.exit(1)

def locate_ipython_profile():
    locate = ["ipython", "locate", "profile", "bokeh_test"]

    try:
        loc = subprocess.check_output(locate)
        write("bokeh_test profile found.")
    except CalledProcessError:
        write("bokeh_test profile not found. Creating one...")
        create_ipython_profile()
        loc = subprocess.check_output(locate)

    body = """
$([IPython.events]).on('status_started.Kernel', function(){
  IPython.notebook.execute_all_cells();
});
"""

    customjs = join(loc[:-1].decode("utf-8"), "static", "custom", "custom.js")
    open(customjs, "w").write(body)

    return loc

def start_ipython_notebook():
    located_profile = locate_ipython_profile()

    notebook_dir = split(dirname(realpath(__file__)))[0]

    cmd = ["ipython", "notebook"]
    argv = ["--no-browser","--profile=bokeh_test", "--port=%s" % args.ipython_port,
            "--notebook-dir=%s" % notebook_dir]

    if located_profile:
        try:
            proc = subprocess.Popen(cmd + argv, stdout=args.log_file, stderr=args.log_file)
        except OSError:
            write("Failed to run: %s" % " ".join(cmd + argv))
            sys.exit(1)

    return proc

detect_examples()

selected_examples = [ (example, flags) for (example, flags) in examples if is_selected(example) ]

need_server = any(flags & Flags.server for _, flags in selected_examples)
need_notebook = any(flags & Flags.notebook for _, flags in selected_examples)

bokeh_server = start_bokeh_server() if use_server and need_server else None

ipython_notebook = start_ipython_notebook() if use_notebook and need_notebook else None

run_examples = 0
fail_examples = []
skip_examples = []

def skip(example, reason):
    write("%s Skipping %s" % (yellow(">>>"), example))
    write("%s because %s" % (yellow("---"), reason))
    skip_examples.append(example)

try:
    for example, flags in sorted(selected_examples):
        if flags & Flags.skip:
            skip(example, "requested to skip")
        elif flags & Flags.server and not bokeh_server:
            skip(example, "bokeh server is not available")
        elif flags & Flags.notebook and not ipython_notebook:
            skip(example, "ipython notebook is not available")
        else:
            write("%s Testing %s ..." % (green(">>>"), example))
            result = False

            if flags & (Flags.file|Flags.server):
                if run_example(example) == 0:
                    result = test_example(example, flags)
                else:
                    fail(example)
            elif flags & Flags.notebook:
                result = deal_with_output_cells(example) and test_example(example, flags)
            else:
                raise ValueError("invalid example type: %s" % example_type(flags))

            if not result:
                fail_examples.append(example)
            run_examples += 1

        write()

    test_nbexecuter()

    if args.upload:
        write("%s Uploading data to S3 to %s/%s ..." % (green(">>>"), s3_bucket, __version__))

        conn = boto.connect_s3()
        bucket = conn.get_bucket(s3_bucket)

        def upload(key, filename):
            if exists(filename):
                key = S3Key(bucket, key)
                key.set_contents_from_filename(filename, policy="public-read")

        entries = []

        for example, _ in sorted(selected_examples):
            no_ext = splitext(abspath(example))[0]
            example_path = relpath(no_ext, base_dir)
            s3_path = join(__version__, example_path)

            png_file = "%s-%s.png" % (no_ext, __version__)
            png_diff = "%s-%s-diff.png" % (no_ext, __version__)

            s3_png_file = s3_path + ".png"
            s3_png_diff = s3_path + "-diff.png"

            upload(s3_png_file, png_file)
            upload(s3_png_diff, png_diff)

            failed = example in fail_examples
            skipped = example in skip_examples

            s3_png_file = join(s3, s3_png_file)

            if args.diff:
                s3_png_diff = join(s3, s3_png_diff)
                s3_ref_png = join(s3, args.diff, example_path) + ".png"
            else:
                s3_png_diff = None
                s3_ref_png = None

            entries.append((example_path, args.diff, failed, skipped, s3_png_file, s3_png_diff, s3_ref_png))
            write(green("---") + " " + example)

        with open(join(dirname(__file__), "examples.html")) as f:
            template = jinja2.Template(f.read())

        html = template.render(version=__version__, entries=entries)

        report_size = len(html)
        write("%s Uploading report (%s) ..." % (green(">>>"), human_bytes(report_size)))
        report = join(__version__, "report.html")
        key = S3Key(bucket, report)
        key.set_metadata("Content-Type", "text/html")
        key.set_contents_from_string(html, policy="public-read")
        write("%s Access report at: %s" % (green("---"), join(s3, report)))

        log_size = os.stat(args.log_file.name).st_size
        write("%s Uploading log (%s) ..." % (green(">>>"), human_bytes(log_size)))
        log = join(__version__, "examples.log")
        key = S3Key(bucket, log)
        key.set_metadata("Content-Type", "text/plain")
        key.set_contents_from_filename(args.log_file.name, policy="public-read")
        write("%s Access log at: %s" % (green("---"), join(s3, log)))
except KeyboardInterrupt:
    write(yellow("INTERRUPTED"))
finally:
    if bokeh_server is not None:
        write("Shutting down bokeh-server ...")
        bokeh_server.kill()

    if ipython_notebook is not None:
        write("Shutting down ipython-notebook ...")
        ipython_notebook.kill()

num_examples = len(selected_examples)

if len(fail_examples) != 0:
    fail("FIX FAILURES AND RUN AGAIN")
    for failed in fail_examples:
        fail(" %s" % failed)
    sys.exit(1)
elif run_examples != num_examples:
    warn("NOT ALL EXAMPLES WERE RUN")
    for skipped in skip_examples:
        warn(" %s" % skipped)
    sys.exit(0)
else:
    ok("ALL TESTS PASS")
    sys.exit(0)
