import json, platform, os, shutil, sys, subprocess, tempfile, threading, time, urllib, urllib2, hashlib
from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
import SocketServer
from optparse import OptionParser
from urlparse import urlparse, parse_qs

USAGE_EXAMPLE = "%prog"

# The local web server uses the git repo as the document root.
DOC_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__),".."))

GIT_CLONE_CHECK = True
DEFAULT_MANIFEST_FILE = 'test_manifest.json'
EQLOG_FILE = 'eq.log'
BROWSERLOG_FILE = 'browser.log'
REFDIR = 'ref'
TMPDIR = 'tmp'
VERBOSE = False
BROWSER_TIMEOUT = 60

SERVER_HOST = "localhost"

class TestOptions(OptionParser):
    def __init__(self, **kwargs):
        OptionParser.__init__(self, **kwargs)
        self.add_option("-m", "--masterMode", action="store_true", dest="masterMode",
                        help="Run the script in master mode.", default=False)
        self.add_option("--noPrompts", action="store_true", dest="noPrompts",
                        help="Uses default answers (intended for CLOUD TESTS only!).", default=False)
        self.add_option("--manifestFile", action="store", type="string", dest="manifestFile",
                        help="A JSON file in the form of test_manifest.json (the default).")
        self.add_option("-b", "--browser", action="store", type="string", dest="browser",
                        help="The path to a single browser (right now, only Firefox is supported).")
        self.add_option("--browserManifestFile", action="store", type="string",
                        dest="browserManifestFile",
                        help="A JSON file in the form of those found in resources/browser_manifests")
        self.add_option("--reftest", action="store_true", dest="reftest",
                        help="Automatically start reftest showing comparison test failures, if there are any.",
                        default=False)
        self.add_option("--port", action="store", dest="port", type="int",
                        help="The port the HTTP server should listen on.", default=8080)
        self.add_option("--unitTest", action="store_true", dest="unitTest",
                        help="Run the unit tests.", default=False)
        self.set_usage(USAGE_EXAMPLE)

    def verifyOptions(self, options):
        if options.reftest and options.unitTest:
            self.error("--reftest and --unitTest must not be specified at the same time.")
        if options.masterMode and options.manifestFile:
            self.error("--masterMode and --manifestFile must not be specified at the same time.")
        if not options.manifestFile:
            options.manifestFile = DEFAULT_MANIFEST_FILE
        if options.browser and options.browserManifestFile:
            print "Warning: ignoring browser argument since manifest file was also supplied"
        if not options.browser and not options.browserManifestFile:
            print "Starting server on port %s." % options.port

        return options
        
def prompt(question):
    '''Return True iff the user answered "yes" to |question|.'''
    inp = raw_input(question +' [yes/no] > ')
    return inp == 'yes'

MIMEs = {
    '.css': 'text/css',
    '.html': 'text/html',
    '.js': 'application/javascript',
    '.json': 'application/json',
    '.svg': 'image/svg+xml',
    '.pdf': 'application/pdf',
    '.xhtml': 'application/xhtml+xml',
    '.gif': 'image/gif',
    '.ico': 'image/x-icon',
    '.png': 'image/png',
    '.log': 'text/plain',
    '.properties': 'text/plain'
}

class State:
    browsers = [ ]
    manifest = { }
    taskResults = { }
    remaining = { }
    results = { }
    done = False
    numErrors = 0
    numEqFailures = 0
    numEqNoSnapshot = 0
    numFBFFailures = 0
    numLoadFailures = 0
    eqLog = None
    lastPost = { }

class UnitTestState:
    browsers = [ ]
    browsersRunning = 0
    lastPost = { }
    numErrors = 0
    numRun = 0

class Result:
    def __init__(self, snapshot, failure, page):
        self.snapshot = snapshot
        self.failure = failure
        self.page = page

class TestServer(SocketServer.TCPServer):
    allow_reuse_address = True

class TestHandlerBase(BaseHTTPRequestHandler):
    # Disable annoying noise by default
    def log_request(code=0, size=0):
        if VERBOSE:
            BaseHTTPRequestHandler.log_request(code, size)

    def sendFile(self, path, ext):
        self.send_response(200)
        self.send_header("Content-Type", MIMEs[ext])
        self.send_header("Content-Length", os.path.getsize(path))
        self.end_headers()
        with open(path, "rb") as f:
            self.wfile.write(f.read())

    def do_GET(self):
        url = urlparse(self.path)
        # Ignore query string
        path, _ = urllib.unquote_plus(url.path), url.query
        path = os.path.abspath(os.path.realpath(DOC_ROOT + os.sep + path))
        prefix = os.path.commonprefix(( path, DOC_ROOT ))
        _, ext = os.path.splitext(path.lower())

        if url.path == "/favicon.ico":
            self.sendFile(os.path.join(DOC_ROOT, "test", "resources", "favicon.ico"), ext)
            return

        if os.path.isdir(path):
            self.sendIndex(url.path, url.query)
            return

        if not (prefix == DOC_ROOT
                and os.path.isfile(path) 
                and ext in MIMEs):
            print path
            self.send_error(404)
            return

        if 'Range' in self.headers:
            # TODO for fetch-as-you-go
            self.send_error(501)
            return

        self.sendFile(path, ext)

class UnitTestHandler(TestHandlerBase):
    def sendIndex(self, path, query):
        print "send index"
    def do_POST(self):
        numBytes = int(self.headers['Content-Length'])

        self.send_response(200)
        self.send_header('Content-Type', 'text/plain')
        self.end_headers()

        url = urlparse(self.path)
        result = json.loads(self.rfile.read(numBytes))
        browser = result['browser']
        UnitTestState.lastPost[browser] = int(time.time())
        if url.path == "/tellMeToQuit":
            tellAppToQuit(url.path, url.query)
            UnitTestState.browsersRunning -= 1
            UnitTestState.lastPost[browser] = None
            return
        elif url.path == '/info':
            print result['message']
        elif url.path == '/submit_task_results':
            status, description = result['status'], result['description']
            UnitTestState.numRun += 1
            if status == 'TEST-UNEXPECTED-FAIL':
                UnitTestState.numErrors += 1
            message = status + ' | ' + description + ' | in ' + browser
            if 'error' in result:
                message += ' | ' + result['error']
            print message
        else:
            print 'Error: uknown action' + url.path

class PDFTestHandler(TestHandlerBase):

    def sendIndex(self, path, query):
        if not path.endswith("/"):
          # we need trailing slash
          self.send_response(301)
          redirectLocation = path + "/"
          if query:
            redirectLocation += "?" + query
          self.send_header("Location",  redirectLocation)
          self.end_headers()
          return

        self.send_response(200)
        self.send_header("Content-Type", "text/html")
        self.end_headers()
        if query == "frame":
          self.wfile.write("<html><frameset cols=*,200><frame name=pdf>" +
            "<frame src='" + path + "'></frameset></html>")
          return

        location = os.path.abspath(os.path.realpath(DOC_ROOT + os.sep + path))
        self.wfile.write("<html><body><h1>PDFs of " + path + "</h1>\n")
        for filename in os.listdir(location):
          if filename.lower().endswith('.pdf'):
            self.wfile.write("<a href='/web/viewer.html?file=" +
              urllib.quote_plus(path + filename, '/') + "' target=pdf>" +
              filename + "</a><br>\n")
        self.wfile.write("</body></html>")


    def do_POST(self):
        numBytes = int(self.headers['Content-Length'])

        self.send_response(200)
        self.send_header('Content-Type', 'text/plain')
        self.end_headers()

        url = urlparse(self.path)
        if url.path == "/tellMeToQuit":
            tellAppToQuit(url.path, url.query)
            return

        result = json.loads(self.rfile.read(numBytes))
        browser, id, failure, round, page, snapshot = result['browser'], result['id'], result['failure'], result['round'], result['page'], result['snapshot']
        State.lastPost[browser] = int(time.time())
        taskResults = State.taskResults[browser][id]
        taskResults[round].append(Result(snapshot, failure, page))

        def isTaskDone():
            numPages = result["numPages"]
            rounds = State.manifest[id]["rounds"]
            for round in range(0,rounds):
                if len(taskResults[round]) < numPages:
                    return False
            return True

        if isTaskDone():
            # sort the results since they sometimes come in out of order
            for results in taskResults:
                results.sort(key=lambda result: result.page)
            check(State.manifest[id], taskResults, browser,
                  self.server.masterMode)
            # Please oh please GC this ...
            del State.taskResults[browser][id]
            State.remaining[browser] -= 1

            checkIfDone()

def checkIfDone():
    State.done = True
    for key in State.remaining:
        if State.remaining[key] != 0:
            State.done = False
            return

# Applescript hack to quit Chrome on Mac
def tellAppToQuit(path, query):
    if platform.system() != "Darwin":
        return
    d = parse_qs(query)
    path = d['path'][0]
    cmd = """osascript<<END
tell application "%s"
quit
end tell
END""" % path
    os.system(cmd)

class BaseBrowserCommand(object):
    def __init__(self, browserRecord):
        self.name = browserRecord["name"]
        self.path = browserRecord["path"]
        self.tempDir = None
        self.process = None

        if platform.system() == "Darwin" and (self.path.endswith(".app") or self.path.endswith(".app/")):
            self._fixupMacPath()

        if not os.path.exists(self.path):
            raise Exception("Path to browser '%s' does not exist." % self.path)

    def setup(self):
        self.tempDir = tempfile.mkdtemp()
        self.profileDir = os.path.join(self.tempDir, "profile")
        self.browserLog = open(BROWSERLOG_FILE, "w")

    def teardown(self):
        self.process.terminate()

        # If the browser is still running, wait up to ten seconds for it to quit
        if self.process and self.process.poll() is None:
            checks = 0
            while self.process.poll() is None and checks < 20:
                checks += 1
                time.sleep(.5)
            # If it's still not dead, try to kill it
            if self.process.poll() is None:
                print "Process %s is still running. Killing." % self.name
                self.process.kill()
                self.process.wait()
            
        if self.tempDir is not None and os.path.exists(self.tempDir):
            shutil.rmtree(self.tempDir)

        self.browserLog.close()

    def start(self, url):
        raise Exception("Can't start BaseBrowserCommand")

class FirefoxBrowserCommand(BaseBrowserCommand):
    def _fixupMacPath(self):
        self.path = os.path.join(self.path, "Contents", "MacOS", "firefox-bin")

    def setup(self):
        super(FirefoxBrowserCommand, self).setup()
        shutil.copytree(os.path.join(DOC_ROOT, "test", "resources", "firefox"),
                        self.profileDir)

    def start(self, url):
        cmds = [self.path]
        if platform.system() == "Darwin":
            cmds.append("-foreground")
        cmds.extend(["-no-remote", "-profile", self.profileDir, url])
        self.process = subprocess.Popen(cmds, stdout = self.browserLog, stderr = self.browserLog)

class ChromeBrowserCommand(BaseBrowserCommand):
    def _fixupMacPath(self):
        self.path = os.path.join(self.path, "Contents", "MacOS", "Google Chrome")

    def start(self, url):
        cmds = [self.path]
        cmds.extend(["--user-data-dir=%s" % self.profileDir,
                     "--no-first-run", "--disable-sync", url])
        self.process = subprocess.Popen(cmds, stdout = self.browserLog, stderr = self.browserLog)

def makeBrowserCommand(browser):
    path = browser["path"].lower()
    name = browser["name"]
    if name is not None:
        name = name.lower()

    types = {"firefox": FirefoxBrowserCommand,
             "chrome": ChromeBrowserCommand }
    command = None
    for key in types.keys():
        if (name and name.find(key) > -1) or path.find(key) > -1:
            command = types[key](browser)
            command.name = command.name or key
            break

    if command is None:
        raise Exception("Unrecognized browser: %s" % browser)

    return command 

def makeBrowserCommands(browserManifestFile):
    with open(browserManifestFile) as bmf:
        browsers = [makeBrowserCommand(browser) for browser in json.load(bmf)]
    return browsers

def downloadLinkedPDFs(manifestList):
    for item in manifestList:
        f, isLink = item['file'], item.get('link', False)
        if isLink and not os.access(f, os.R_OK):
            linkFile = open(f +'.link')
            link = linkFile.read()
            linkFile.close()

            sys.stdout.write('Downloading '+ link +' to '+ f +' ...')
            sys.stdout.flush()
            response = urllib2.urlopen(link)

            with open(f, 'wb') as out:
                out.write(response.read())

            print 'done'

def verifyPDFs(manifestList):
    error = False
    for item in manifestList:
        f = item['file']
        if os.access(f, os.R_OK):
            fileMd5 = hashlib.md5(open(f, 'rb').read()).hexdigest()
            if 'md5' not in item:
                print 'WARNING: Missing md5 for file "' + f + '".',
                print 'Hash for current file is "' + fileMd5 + '"'
                error = True
                continue
            md5 = item['md5']
            if fileMd5 != md5:
                print 'WARNING: MD5 of file "' + f + '" does not match file.',
                print 'Expected "' + md5 + '" computed "' + fileMd5 + '"'
                error = True
                continue
        else:
            print 'WARNING: Unable to open file for reading "' + f + '".'
            error = True
    return not error

def getTestBrowsers(options):
    testBrowsers = []
    if options.browserManifestFile:
        testBrowsers = makeBrowserCommands(options.browserManifestFile)
    elif options.browser:
        testBrowsers = [makeBrowserCommand({"path":options.browser, "name":None})]

    if options.browserManifestFile or options.browser:
        assert len(testBrowsers) > 0
    return testBrowsers

def setUp(options):
    # Only serve files from a pdf.js clone
    assert not GIT_CLONE_CHECK or os.path.isfile('../src/pdf.js') and os.path.isdir('../.git')

    if options.masterMode and os.path.isdir(TMPDIR):
        print 'Temporary snapshot dir tmp/ is still around.'
        print 'tmp/ can be removed if it has nothing you need.'
        if options.noPrompts or prompt('SHOULD THIS SCRIPT REMOVE tmp/?  THINK CAREFULLY'):
            subprocess.call(( 'rm', '-rf', 'tmp' ))

    assert not os.path.isdir(TMPDIR)

    testBrowsers = getTestBrowsers(options)

    with open(options.manifestFile) as mf:
        manifestList = json.load(mf)

    downloadLinkedPDFs(manifestList)

    if not verifyPDFs(manifestList):
      print 'Unable to verify the checksum for the files that are used for testing.'
      print 'Please re-download the files, or adjust the MD5 checksum in the manifest for the files listed above.\n'

    for b in testBrowsers:
        State.taskResults[b.name] = { }
        State.remaining[b.name] = len(manifestList)
        State.lastPost[b.name] = int(time.time())
        for item in manifestList:
            id, rounds = item['id'], int(item['rounds'])
            State.manifest[id] = item
            taskResults = [ ]
            for r in xrange(rounds):
                taskResults.append([ ])
            State.taskResults[b.name][id] = taskResults

    return testBrowsers

def setUpUnitTests(options):
    # Only serve files from a pdf.js clone
    assert not GIT_CLONE_CHECK or os.path.isfile('../src/pdf.js') and os.path.isdir('../.git')

    testBrowsers = getTestBrowsers(options)

    UnitTestState.browsersRunning = len(testBrowsers)
    for b in testBrowsers:
        UnitTestState.lastPost[b.name] = int(time.time())
    return testBrowsers

def startBrowsers(browsers, options, path):
    for b in browsers:
        b.setup()
        print 'Launching', b.name
        host = 'http://%s:%s' % (SERVER_HOST, options.port) 
        qs = '?browser='+ urllib.quote(b.name) +'&manifestFile='+ urllib.quote(options.manifestFile)
        qs += '&path=' + b.path
        b.start(host + path + qs)

def teardownBrowsers(browsers):
    for b in browsers:
        try:
            b.teardown()
        except:
            print "Error cleaning up after browser at ", b.path
            print "Temp dir was ", b.tempDir
            print "Error:", sys.exc_info()[0]

def check(task, results, browser, masterMode):
    failed = False
    for r in xrange(len(results)):
        pageResults = results[r]
        for p in xrange(len(pageResults)):
            pageResult = pageResults[p]
            if pageResult is None:
                continue
            failure = pageResult.failure
            if failure:
                failed = True
                State.numErrors += 1
                print 'TEST-UNEXPECTED-FAIL | test failed', task['id'], '| in', browser, '| page', p + 1, 'round', r, '|', failure

    if failed:
        return

    kind = task['type']
    if 'eq' == kind:
        checkEq(task, results, browser, masterMode)
    elif 'fbf' == kind:
        checkFBF(task, results, browser)
    elif 'load' == kind:
        checkLoad(task, results, browser)
    else:
        assert 0 and 'Unknown test type'


def checkEq(task, results, browser, masterMode):
    pfx = os.path.join(REFDIR, sys.platform, browser, task['id'])
    results = results[0]
    taskId = task['id']

    passed = True
    for page in xrange(len(results)):
        snapshot = results[page].snapshot
        ref = None
        eq = True

        path = os.path.join(pfx, str(page + 1))
        if not os.access(path, os.R_OK):
            State.numEqNoSnapshot += 1
            if not masterMode:
                print 'WARNING: no reference snapshot', path
        else:
            f = open(path)
            ref = f.read()
            f.close()

            eq = (ref == snapshot)
            if not eq:
                print 'TEST-UNEXPECTED-FAIL | eq', taskId, '| in', browser, '| rendering of page', page + 1, '!= reference rendering'

                if not State.eqLog:
                    State.eqLog = open(EQLOG_FILE, 'w')
                eqLog = State.eqLog

                # NB: this follows the format of Mozilla reftest
                # output so that we can reuse its reftest-analyzer
                # script
                eqLog.write('REFTEST TEST-UNEXPECTED-FAIL | ' + browser +'-'+ taskId +'-page'+ str(page + 1) + ' | image comparison (==)\n')
                eqLog.write('REFTEST   IMAGE 1 (TEST): ' + snapshot + '\n')
                eqLog.write('REFTEST   IMAGE 2 (REFERENCE): ' + ref + '\n')

                passed = False
                State.numEqFailures += 1

        if masterMode and (ref is None or not eq):
            tmpTaskDir = os.path.join(TMPDIR, sys.platform, browser, task['id'])
            try:
                os.makedirs(tmpTaskDir)
            except OSError, e:
                if e.errno != 17: # file exists
                    print >>sys.stderr, 'Creating', tmpTaskDir, 'failed!'
        
            of = open(os.path.join(tmpTaskDir, str(page + 1)), 'w')
            of.write(snapshot)
            of.close()

    if passed:
        print 'TEST-PASS | eq test', task['id'], '| in', browser

def checkFBF(task, results, browser):
    round0, round1 = results[0], results[1]
    assert len(round0) == len(round1)

    passed = True
    for page in xrange(len(round1)):
        r0Page, r1Page = round0[page], round1[page]
        if r0Page is None:
            break
        if r0Page.snapshot != r1Page.snapshot:
            print 'TEST-UNEXPECTED-FAIL | forward-back-forward test', task['id'], '| in', browser, '| first rendering of page', page + 1, '!= second'
            passed = False
            State.numFBFFailures += 1
    if passed:
        print 'TEST-PASS | forward-back-forward test', task['id'], '| in', browser


def checkLoad(task, results, browser):
    # Load just checks for absence of failure, so if we got here the
    # test has passed
    print 'TEST-PASS | load test', task['id'], '| in', browser


def processResults():
    print ''
    numFatalFailures = (State.numErrors + State.numFBFFailures)
    if 0 == State.numEqFailures and 0 == numFatalFailures:
        print 'All regression tests passed.'
    else:
        print 'OHNOES!  Some tests failed!'
        if 0 < State.numErrors:
            print '  errors:', State.numErrors
        if 0 < State.numEqFailures:
            print '  different ref/snapshot:', State.numEqFailures
        if 0 < State.numFBFFailures:
            print '  different first/second rendering:', State.numFBFFailures


def maybeUpdateRefImages(options, browser):
    if options.masterMode and (0 < State.numEqFailures or 0 < State.numEqNoSnapshot): 
        print "Some eq tests failed or didn't have snapshots."
        print 'Checking to see if master references can be updated...'
        numFatalFailures = (State.numErrors + State.numFBFFailures)
        if 0 < numFatalFailures:
            print '  No.  Some non-eq tests failed.'
        else:
            print '  Yes!  The references in tmp/ can be synced with ref/.'
            if options.reftest:                                                                                                              
                startReftest(browser, options)
            if options.noPrompts or prompt('Would you like to update the master copy in ref/?'):
                sys.stdout.write('  Updating ref/ ... ')

                # XXX unclear what to do on errors here ...
                # NB: do *NOT* pass --delete to rsync.  That breaks this
                # entire scheme.
                subprocess.check_call(( 'rsync', '-arvq', 'tmp/', 'ref/' ))

                print 'done'
            else:
                print '  OK, not updating.'

def startReftest(browser, options):
    url = "http://%s:%s" % (SERVER_HOST, options.port)
    url += "/test/resources/reftest-analyzer.xhtml"
    url += "#web=/test/eq.log"
    try:
        browser.setup()
        browser.start(url)
        print "Waiting for browser..."
        browser.process.wait()
    finally:
        teardownBrowsers([browser])
    print "Completed reftest usage."

def runTests(options, browsers):
    t1 = time.time()
    try:
        startBrowsers(browsers, options, '/test/test_slave.html')
        while not State.done:
            for b in State.lastPost:
                if State.remaining[b] > 0 and int(time.time()) - State.lastPost[b] > BROWSER_TIMEOUT:
                    print 'TEST-UNEXPECTED-FAIL | test failed', b, "has not responded in", BROWSER_TIMEOUT, "s"
                    State.numErrors += State.remaining[b]
                    State.remaining[b] = 0
                    checkIfDone()
            time.sleep(1)
        processResults()
    finally:
        teardownBrowsers(browsers)
    t2 = time.time()
    print "Runtime was", int(t2 - t1), "seconds"
    if State.eqLog:
        State.eqLog.close();
    if options.masterMode:
        maybeUpdateRefImages(options, browsers[0])
    elif options.reftest and State.numEqFailures > 0:
        print "\nStarting reftest harness to examine %d eq test failures." % State.numEqFailures
        startReftest(browsers[0], options)

def runUnitTests(options, browsers):
    t1 = time.time()
    try:
        startBrowsers(browsers, options, '/test/unit/unit_test.html')
        while UnitTestState.browsersRunning > 0:
            for b in UnitTestState.lastPost:
                if UnitTestState.lastPost[b] != None and int(time.time()) - UnitTestState.lastPost[b] > BROWSER_TIMEOUT:
                    print 'TEST-UNEXPECTED-FAIL | test failed', b, "has not responded in", BROWSER_TIMEOUT, "s"
                    UnitTestState.lastPost[b] = None
                    UnitTestState.browsersRunning -= 1
                    UnitTestState.numErrors += 1
            time.sleep(1)
        print ''
        print 'Ran', UnitTestState.numRun, 'tests'
        if UnitTestState.numErrors > 0:
            print 'OHNOES!  Some tests failed!'
            print '  ', UnitTestState.numErrors, 'of', UnitTestState.numRun, 'failed'
        else:
            print 'All unit tests passed.'
    finally:
        teardownBrowsers(browsers)
    t2 = time.time()
    print "Unit test Runtime was", int(t2 - t1), "seconds"

def main():
    optionParser = TestOptions()
    options, args = optionParser.parse_args()
    options = optionParser.verifyOptions(options)
    if options == None:
        sys.exit(1)

    if options.unitTest:
        httpd = TestServer((SERVER_HOST, options.port), UnitTestHandler)
        httpd_thread = threading.Thread(target=httpd.serve_forever)
        httpd_thread.setDaemon(True)
        httpd_thread.start()

        browsers = setUpUnitTests(options)
        if len(browsers) > 0:
            runUnitTests(options, browsers)
    else:
        httpd = TestServer((SERVER_HOST, options.port), PDFTestHandler)
        httpd.masterMode = options.masterMode
        httpd_thread = threading.Thread(target=httpd.serve_forever)
        httpd_thread.setDaemon(True)
        httpd_thread.start()

        browsers = setUp(options)
        if len(browsers) > 0:
            runTests(options, browsers)
        else:
            # just run the server
            print "Running HTTP server. Press Ctrl-C to quit."
            try:
                while True:
                    time.sleep(1)
            except (KeyboardInterrupt):
                print "\nExiting."

if __name__ == '__main__':
    main()