#!/usr/bin/env python

from __future__ import print_function

import os
import sys
import time
import json
import signal
import fnmatch
import argparse
import traceback
import subprocess
from subprocess import CalledProcessError

import yaml

import requests
import requests.exceptions

import colorama
colorama.init()

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

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

from bokeh import __version__

if os.environ.get("TRAVIS_PYTHON_VERSION") is not None:
    default_timeout = 30
else:
    default_timeout = 10

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.")
args = parser.parse_args()

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"]:
    raise ValueError("unexpected value for BOKEH_DEFAULT_NO_DEV environmental variable, %s" % DEFAULT_NO_DEV)

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 = os.path.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 = os.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((os.path.join(examples_path, file), flags))

def detect_examples():
    with open(os.path.join(os.path.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 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 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 ""
    print("%s%s" % (red("[FAIL]"), msg))

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

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

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", os.path.basename(example)]
    cwd = os.path.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)

        try:
            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 = os.path.split(path)
        parts.append(tail)
        path = newpath
        if tail == 'examples':
            break
    parts.reverse()
    return parts

def test_example(example, flags):
    no_ext = os.path.splitext(os.path.abspath(example))[0]
    url_path = os.path.join(*get_path_parts(os.path.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, os.path.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, os.path.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:
        print("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:
                    print(msg)
                elif line is None:
                    print("%s: %s" % (source, msg))
                else:
                    print("%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:
            print(error['msg'])

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

        if resource_errors or errors:
            fail(example)
        else:
            ok(example)
            return True

    return False

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)
    except OSError:
        print("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 requests.exceptions.ConnectionError:
                    return False

            return wait_until(helper)

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

        return proc

def create_ipython_profile():
    create = ["ipython", "profile", "create", "bokeh_test"]
    try:
        subprocess.check_call(create)
    except OSError:
        print("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)
        print("bokeh_test profile found.")
    except CalledProcessError:
        print("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 = os.path.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 = os.path.split(os.path.dirname(os.path.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)
        except OSError:
            print("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

if use_notebook:
    import nbexecuter
    example_nbconverted = os.path.join("examples", "glyphs", "glyph.ipynb")
    print(example_nbconverted) 
    nbexecuter.main(example_nbconverted)

run_examples = 0
fail_examples = []
skip_examples = []

try:
    for example, flags in sorted(selected_examples):
        if (flags & Flags.skip or
            flags & Flags.server and not bokeh_server or
            flags & Flags.notebook and not ipython_notebook):
            print("%s Skipping %s" % (yellow(">>>"), example))
            skip_examples.append(example)
        else:
            print("%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()
            elif flags & Flags.notebook:
                result = test_example(example, flags)
            else:
                raise ValueError("invalid example type: %s" % example_type(flags))

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

        print()
except KeyboardInterrupt:
    print(yellow("INTERRUPTED"))
finally:
    if bokeh_server is not None:
        print("Shutting down bokeh-server ...")
        bokeh_server.kill()

    if ipython_notebook is not None:
        print("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)