#!/usr/bin/env python

from __future__ import print_function

import os
import sys
import json
import fnmatch
import argparse
import traceback
import subprocess

import colorama
colorama.init()

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", "--notebook-port", type=int, default=8888,
                    help="port on which IPython Notebook server resides")
parser.add_argument("-v", "--verbose", action="store_true", default=False,
                    help="show console messages")
args = parser.parse_args()

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

examples = []

def add_examples(examples_dir, includes=set([]), excludes=set([]), etype='file'):
    ext = '.ipynb' if etype == 'notebook' else '.py'

    for file in includes or os.listdir(examples_dir):
        if not file.startswith('__') and file.endswith(ext) and file not in excludes:
            examples.append((os.path.join(examples_dir, file), etype))

glyphs_dir = os.path.join("glyphs")
glyphs_server = set(["data_select_tool.py", "glyph2.py", "maps.py", "prim.py"])
glyphs_excludes = set(["line_animate.py"])
add_examples(glyphs_dir, excludes=glyphs_excludes|glyphs_server, etype='file')
add_examples(glyphs_dir, includes=glyphs_server, etype='server')

plotting_file_dir = os.path.join("plotting", "file")
glyphs_excludes = set(["image.py", "rgba_image.py"])
add_examples(plotting_file_dir, etype='file', excludes=glyphs_excludes)

plotting_server_dir = os.path.join("plotting", "server")
plotting_server_excludes = set(["animated.py", "line_animate.py"])
add_examples(plotting_server_dir, excludes=plotting_server_excludes, etype='server')

plotting_notebook_dir = os.path.join("plotting", "notebook")
plotting_notebook_excludes = set(["animated.ipynb"])
add_examples(plotting_notebook_dir, excludes=plotting_notebook_excludes, etype='notebook')

def warn(msg=None):
    msg = " " + msg if msg is not None else ""
    print("%s[WARN]%s%s" % (colorama.Fore.YELLOW, colorama.Style.RESET_ALL, msg))

def fail(msg=None):
    msg = " " + msg if msg is not None else ""
    print("%s[FAIL]%s%s" % (colorama.Fore.RED, colorama.Style.RESET_ALL, msg))

def ok(msg=None):
    msg = " " + msg if msg is not None else ""
    print("%s[OK]%s%s" % (colorama.Fore.GREEN, colorama.Style.RESET_ALL, msg))

env = os.environ.copy()
env['BOKEH_BROWSER'] = 'dummy'

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

num_examples = len(selected_examples)
run_examples = 0
ok_examples = 0

try:
    for example, etype in sorted(selected_examples):
        print("%s>>>%s Testing %s ..." % (colorama.Fore.YELLOW, colorama.Style.RESET_ALL, example))

        if etype != 'notebook':
            cwd = os.path.dirname(example)
            module = os.path.splitext(os.path.basename(example))[0]

            proc = subprocess.Popen(["python", "-c", "import %s" % module], cwd=cwd, env=env)
            code = proc.wait()

        if etype == 'notebook' or code == 0:
            no_ext = os.path.splitext(os.path.abspath(example))[0]

            png_file = no_ext + '.png'
            html_file = no_ext + '.html'

            server_url = 'http://localhost:%d/bokeh/doc/%s/show'
            notebook_url = 'http://localhost:%d/notebooks/%s.ipynb'

            if etype == 'server':
                url = server_url % (args.bokeh_port, os.path.basename(no_ext))
            elif etype == 'notebook':
                url = notebook_url % (args.notebook_port, os.path.basename(no_ext))
            elif etype == 'file':
                url = 'file://' + html_file

            cmd = ["phantomjs", "test.js", etype, url, png_file]
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
            code = proc.wait()

            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))

                if resources or errors:
                    for resource in resources:
                        print("%s: %s (%s)" % (
                            resource['url'], resource['status'], resource['statusText']))

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

                        for item in error['trace']:
                            print("    %s: %d" % (item['file'], item['line']))
                    fail()
                else:
                    ok_examples += 1
                    ok()
        else:
            fail()

        run_examples += 1
        print()
except KeyboardInterrupt:
    print("%sINTERRUPTED%s" % (colorama.Fore.YELLOW, colorama.Style.RESET_ALL))

if ok_examples != num_examples:
    if run_examples == num_examples:
        fail("FIX FAILURES AND RUN AGAIN")
    else:
        warn("NOT ALL EXAMPLES WERE RUN")

    sys.exit(1)
else:
    ok("ALL TESTS PASS")
    sys.exit(0)
