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

VERSION = "3.2.1"

# MIT License

# Copyright (c) 2016 - 2022 University of St Andrews

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import sys
import signal
import os
import subprocess
import re
import time
import json
import copy
from threading import Thread
import difflib
from optparse import OptionParser, SUPPRESS_HELP
import argparse
import configparser
import tempfile
import shutil
import zipfile
import tarfile

# Stacscheck

# Try importing jinja2, but don't complain if it isn't there
try:
    import jinja2
except ImportError:
    jinja = False
else:
    jinja = True

# The HTML output
JINJAPAGE = u"""
<!DOCTYPE html>
<html lang="en">
<style>
    table.diff {
        table-layout: inherit;
        font-family:Courier;
        border: medium;
        # width: 95vw;
        word-wrap: break-word;
    }
    .code {font-family:Courier; border:medium; width: 2000px}
    .diff_header {background-color:#3a3a3a; color: #ffffff}
    td.diff_header {text-align: center;}
    .diff_next {background-color:#ff00dd; display: none}
    .diff_add {background-color:#24906c}
    .diff_chg {background-color:#ffff77; color:black}
    .diff_sub {background-color:#E44C4B}
    :target {
    background: yellow;
    }

    /* Make the whole thing have a dark background*/
    body {background-color: #292929; width: 100%;}
    /* For test descriptions (PASS build-faro) and stacscheck info */
    .test_desc {font-size: 22px; font-family: Arial; color: #ffffff; font-weight: normal;}
    /* For any normal code */
    .infocode {font-size: 14px; font-family: Courier; color: #ffffff; margin-left: 20px}
    /* For compared code */
    .compcode {font-size: 14px; font-family: Courier; color: #ffffff;
        margin-left: 20px}
    h2 {font-size: 18px; font-family: Arial; color: #ffffff; font-weight: bold; margin-left: 20px}
    /* Collapsible testing */
    .collapsible {display: none;}

    table { table-layout: fixed;}
    # .diff {
    #     display: table;
    #     table-layout: inherit;
    #     word-wrap: break-word;
    #     background-color: coral;
    # }

    # .test-content {
    #     # display: table-cell;
    #     width: 100%;
    #     colspan: 3;
    # }

    td {white-space: normal;}

    # testtable {width: 100%;}

    .buttonholder {width: 2%;}

</style>
<head>
    <title>{{Practical}} - {{ SubmissionID }} </title>
    <h1 class="test_desc">{{ Finalmessage }}</h1>
    <br>
</head>
<body>
    <h1 class="test_desc"> STACSCHECK INFO </h1>
    <div class="infocode">{% for line in InfoText %}
        {{line}}
        <br/>
    {% endfor %}
    <br>
    </div>
    <table id="testtable">
    {% for test in Tests %}
        <tr id="test{{ loop.index }}" bgcolor=
        {% if test.pass %} "#2aa77d" {% else %} "#E44C4B" {% endif %} >
        <td class="buttonholder">
        <button class="collapsebutton" onclick="docollapse('content-{{loop.index}}')">
        expand
        </button>
        </td>
        <td class="test_desc">{% if test.pass %} PASS {% else %} FAIL {% endif %}</td>
        <td class="test_desc" id="{{ test.name }}"> {{ test.name }} </td>
        </tr>
        <div class="test-content" >
        <tr id="content-{{loop.index}}" class="collapsible">
        <td colspan="5" class="code">

        {% if test.type != "Prog" %}
            <h2> Compiler info </h2>
            <div display="inline" class="infocode">
            {% for line in  test.stdout.splitlines() %}
                {{line}}<br>
            {% endfor %}
            </div>
            <div class="infocode">
            {% for line in  test.stderr.splitlines() %}
                {{line}}<br>
            {% endfor %}
            </div>

        {% else %}
            {% if test.inputlines != None %}
                <h2> Input </h2>
                <div class="infocode">
                {% for intext in test.inputlines %}
                    {{ intext }}
                    <br>
                {% endfor %}
                </div>
            {% endif %}
            <h2> Output </h2>

            {#
            <h3> Expected </h3>
            {% for exp in test.comparelines %}
                {{ exp }}
            {% endfor %}
            <h3> Submission </h3>
            {% for sub in test.userlines %}
                {{ sub }}
            {% endfor %}
            #}

            <div class="compcode">
            {% if test.htmldiff %}
                {{ test.htmldiff | safe }}
            </div>
            {% else %}
                <div class="infocode">
                {% for line in  test.stdout.splitlines() %}
                    {{line}}<br>
                {% endfor %}
                </div>
                <div class="infocode">
                {% for line in  test.stderr.splitlines() %}
                    {{line}}<br>
                {% endfor %}
                </div>
            {% endif %}

        {% endif %}
        <br> <br>
        </td>
        </tr>
        </div>
    {% endfor %}
    </table>
    {% if AnyBuildTestFailed %}
    <p>Building is failing, skipping tests</p>
    {% endif %}
</body>
<script>
    function docollapse(contentid) {
        var content = document.getElementById(contentid);
        console.log(content + " banana");
        if (content.style.display == "table-row"){
            content.style.display = "none";
        } else {
            content.style.display = content.style.display = "table-row";
        }
    }
</script>
</html>

"""

# Store the base test directory
TESTBASE = None

# Try harder to run badly formed submissions
TRYHARDER = False

# Store the submission directory
SUBMISSIONBASE = None

# Store parent directory containing practical.config
PARENTBASE = None

# Functions to convert to UTF-8 and print the string
unicode = str
def toUTF8Replace(bytes):
    if isinstance(bytes, str):
        return bytes
    return str(bytes, "utf8", errors="replace")

def output_str(arg):
    print(toUTF8Replace(arg), end='')

# Print info and store in INFO_PRINTED (to keep for html, json, ...)
INFO_PRINTED = []

def info_print_nonewline(string):
    global INFO_PRINTED
    output_str(string)
    INFO_PRINTED.append(string)

def info_print(string):
    info_print_nonewline(string + "\n")

# Simple function to print out more verbose information if -v is passed
VERBOSE = True
def verbose_print(arg):
    if VERBOSE:
        output_str(arg + "\n")
        output_str("\n")


# Check function for problems in practical
def warn_print(arg):
    info_print("\n- ERROR - There is a problem in the test directory.")
    info_print("The problem is: " + str(arg) + ".")
#    info_print("Possible reasons:")
#    info_print("  - Your copy of stacscheck is out of date.")
#    info_print("  - The directory of tests is missing some files, or is broken.")
    info_print("This may be caused by a broken test directory, or")
    info_print("the test directory may have been copied incorrectly.")
def print_file_contents(
        file_loc,
        ignore_prefixes=["#!"], #Do not modify the list in the method, otherwise troubles may occur
        ignore_empty_lines=True
):
    try:
        for line in open(file_loc, "rb"):
            line = toUTF8Replace(line)

            if ignore_empty_lines and line.strip() == "":
                continue

            if any(line.startswith(ignore_prefix) for ignore_prefix in ignore_prefixes):
                continue

            output_str(line.rstrip() + "\n")
    except IOError as e:
        warn_print("Error reading '" + str(file_loc) + "': " + e.strerror)

# Takes a single string with newlines, and returns a list of lines
# We strip all whitespace, then add a "\n" on the end
# This is to deal with differences in invisible whitespaces
def strip_string(string):
    return [l.rstrip() + "\n" for l in string.split("\n") if l.rstrip() != '']


# Take a possibly huge output from a program, and reduce it's size so it fits
# reasonably on screen.
# userlines: The output of the user's program
# comparelines: Expected output (or an empty list if there is no expected output)
def reduce_huge_output(userlines, comparelines):
    if VERBOSE:
        return userlines

    # Don't break the original
    userlines = copy.deepcopy(userlines)
    # If the number of lines is massive, bring it slightly more under control.
    if len(userlines) > len(comparelines) + 25:
        userlines = userlines[:len(comparelines) + 25] + [u"... (output truncated)\n"]

    # Now check if any line is stupidly long
    for i in range(len(userlines)):
        if i < len(comparelines):
            comparelen = len(comparelines[i])
        else:
            comparelen = 0

        keeplen = max(comparelen + 25, 256)
        if len(userlines[i]) > keeplen:
            userlines[i] = userlines[i][:keeplen] + u" ... (line truncated)\n"

    return userlines


# Store the results of all tests
testStore = []

# Store if any build step failed
anyBuildTestFailed = False

CONFIG = {'course' : '', 'practical' : '', 'timeout': '60' }

#### Beginning of functions
def try_parse_config_file(basedir):
    config = configparser.ConfigParser()

    # Search up until it finds it (limit 10) and store it under a global
    parentdir = basedir
    for x in range(11):
        # Try to get config file path
        filename = os.path.join(parentdir, "practical.config")
        # If it exists, continue
        if os.path.isfile(filename):
            break
        # Otherwise, get parent directory and try again
        split_path = os.path.split(parentdir)
        parentdir = split_path[0]
        if parentdir == "/":
            warn_print("reached root directory without finding practical.config")
            sys.exit(1)

    # Check that something was found, or quit
    if not os.path.isfile(filename):
        warn_print("no practical.config file found in directory or 10 parent directories")
        sys.exit(1)

    # filename = os.path.join(basedir, "practical.config")
    # if not os.path.isfile(filename):
    #     print("'" + basedir + "' is not a valid test directory")
    #     print("(the test directory does not contain a file called 'practical.config')")
    #     sys.exit(1)
    try:
        config.read(filename)
    except IOError as e:
        print("I/O error reading practical.config: " + e.strerror)
        sys.exit(1)
    except:
        warn_print("practical.config is malformed")
        sys.exit(1)

    if not 'info' in config.sections():
        warn_print("practical.config does not have an '[info]' section")
        sys.exit(1)

    practical_version = None
    # Check version early, for useful error messages
    if 'version' in config.options('info'):
        practical_version = config.get('info', 'version')
    elif 'version' in config.sections() and 'required' in config.options('version'):
        practical_version = config.get('version', 'required')

    if practical_version is not None:
        if [int(x) for x in practical_version.split(".")] > [int(x) for x in VERSION.split(".")]:
            print("- ERROR -")
            print("This practical requires stacscheck version " + practical_version)
            print("But this is stacscheck " + VERSION)
            sys.exit(1)


    for section in config.sections():
        if not section in ['info', 'version']:
            warn_print("practical.config has an unexpected section: ", section)
            sys.exit(1)
    
    if not 'info' in config.sections():
        warn_print("practical.config does not have an [info] section")
        sys.exit(1)

    valid_options = ['course', 'practical', 'srcdir', 'timeout', 'version']

    for option in config.options('info'):
        if not (option in valid_options):
            warn_print("Don't understand option '" + option + "'")
            sys.exit(1)
        CONFIG[option] = config.get('info', option)

    if 'version' in config.sections():
        for option in config.options('version'):
            if not(option in ['required']):
                warn_print("Don't understand [version] in practical.config: '" + version +"'")
                sys.exit(1)
            if 'version' in CONFIG:
                warn_print("'version' defined twice in practical.config")
                sys.exit(1)
            CONFIG['version'] = config.get('version', option)


    for section in config.sections():
        if not(section in ['info', 'version']):
            warn_print("practical.config contains invalid section: ", section)
            sys.exit(1)

    return parentdir

# Walk directory tree under 'rootdir', finding all directories with 'name'.
def find_all_directories_with_name(name, rootdir):
    dirlist = []
    for root, subfolders, _ in os.walk(rootdir):
        for dirname in subfolders:
            if dirname == name and ("/__MACOSX" not in os.path.join(root, dirname)):
                dirlist.append(os.path.join(root, dirname))
    return sorted(dirlist)

# Find the directory the code is stored in, using 'srcdir' from CONFIG.
# If 'srcdir' is not defined, assume current directory, else search for a directory called
# CONFIG['srcdir'] (or check if we are already in it!)
def find_code_directory(start_dir=None):
    # if not 'srcdir' in CONFIG:  # ACELGA
    #     return None
    # If there's a starting directory, start searching from there. Otherwise use cwd
    if not start_dir == None:
        current_dir = start_dir
    else:
        current_dir = os.path.realpath(os.getcwd())

    in_srcdir = (os.path.basename(current_dir) == CONFIG['srcdir'])
    recursive_srcdir = find_all_directories_with_name(CONFIG['srcdir'], current_dir)
    # print("--- Trying to find code directory in " + str(start_dir))
    info_print_nonewline("- Looking for submission in a directory called '" + CONFIG['srcdir'] + "': ")

    if not in_srcdir and len(recursive_srcdir) == 0:
        info_print("Cannot find it....")
        if not TRYHARDER:
            raise EnvironmentError("No submission found. Try running with --archive <filename> if submission is compressed.")
        else:
            info_print("Going to try hard, and guess the current directory. I'm probably wrong.")
            return None

    if in_srcdir and len(recursive_srcdir) == 0:
        info_print("Already in it!")
        return None

    if len(recursive_srcdir) == 1 and not in_srcdir:
        reldir = os.path.relpath(recursive_srcdir[0], current_dir)
        if reldir == CONFIG['srcdir']:
            info_print("found in current directory")
        else:
            info_print("found at '" + os.path.relpath(recursive_srcdir[0], current_dir) + "'")
        return recursive_srcdir[0]

    # There is more than one place the source might be. Take a best guess.
    if len(recursive_srcdir) >= 1 and in_srcdir:
        info_print("\n- Warning: in a directory called '" + CONFIG['srcdir'] + "', with subdirectories also with the same name.")
        info_print("- Guessing your submission is in the current directory.")
        info_print("- If that's wrong, go into the correct directory.")
        return None

    # There is multiple recursive_srcdir, and we are not in directory of the right name
    info_print("- Warning, there are multiple subdirectories called '" + CONFIG['srcdir'] + "'.")
    for d in recursive_srcdir:
        info_print(" - found " + d)
    info_print("Please go into the correct directory.")
    sys.exit(1)

def maybe_print_test_input(test):
    if VERBOSE and not test["pass"]:
        print("--- ran script ---")
        print_file_contents(file_loc=test["path"])
        if "input" in test and test["input"] is not None:
            print("--- with input ---")
            print_file_contents(file_loc=test["input"], ignore_prefixes=[], ignore_empty_lines=False)

def print_test_or_info_output(test):
    maybe_print_test_input(test)

    stdout = strip_string(test.get("stdout", "")) #.rstrip().encode("utf8", errors="replace"))
    stderr = strip_string(test.get("stderr", "")) #.rstrip().encode("utf8", errors="replace"))

    # Always print all student's output if this is an 'info', or --verbose is enabled
    if not(VERBOSE or test.get("alwaysoutput", False)):
        stdout = reduce_huge_output(stdout, [])
        stderr = reduce_huge_output(stderr, [])

    if test.get("stderr", False) and test.get("stdout", False):
        output_str("--- submission output ---\n")
        output_str(u"".join(stdout))
        output_str("--- submission stderr output ---\n")
        output_str(u"".join(stderr))
        output_str("---\n\n")
    elif test.get("stderr", False) or test.get("stdout", False):
        output_str("--- submission output ---\n")
        output_str(u"".join(stdout) + u"".join(stderr))
        output_str("---\n\n")
    else:
        output_str("--- no output from submission---\n\n")

# Record that a test was run, printing as appropriate
def register_returnval_test(test):
    if test["returnval"] == 0:
        print("pass")
        test["pass"] = True
    else:
        print("fail")
        test["pass"] = False

    if not test["pass"] or test.get("alwaysoutput", False):
        print_test_or_info_output(test)

    testStore.append(test)


# Register the result of a test whith a known output
def register_diff_test(test, comparefile):
    with open(comparefile, 'rb') as stream:
        comparelines = toUTF8Replace(stream.read())
    comparelines = strip_string(comparelines)
    userlines = strip_string(test["stdout"] + test["stderr"])

    userlines = reduce_huge_output(userlines, comparelines)

    test["comparelines"] = comparelines
    test["userlines"] = userlines

    # If test passed just print pass
    if comparelines == userlines:
        print("pass")
        test["pass"] = True

    # If test failed print expected & real output (& possibly extras)
    else:
        print("fail")
        test["pass"] = False
        # If verbose, print input to stdout
        maybe_print_test_input(test)
        # Create text for expected & real output
        textdiff = []
        if len(comparelines) > 0:
            textdiff.extend([u"--- expected output ---\n"] + comparelines)
        else:
            textdiff.extend([u"--- no output expected ---\n"])

        if len(userlines) > 0:
            textdiff.extend([u"--- submission output ---\n"] + userlines + [u"---\n"])
        else:
            textdiff.extend([u"--- no output from submission ---\n"])

        # Print to stdout
        test["textdiff"] = textdiff
        textout = u"".join(test["textdiff"])
        assert isinstance(textout, str)
        output_str(textout + "\n")



        # Check things are 'str's

        for l in test["textdiff"]:
            assert isinstance(l, str)

    # Create table for html
    test["htmldiff"] = difflib.HtmlDiff().make_table(userlines, comparelines, u"Submission", u"Reference")
    assert isinstance(test["htmldiff"], str)

    testStore.append(test)


#################### FILE UTILITIES ######################
# Given a filename of a test, get a nicer, shorter name which
# describes the test. First drop extension, then remove PARENTBASE
def nice_name(name):
    noextension = os.path.splitext(name)[0]
    dropdir = noextension[len(PARENTBASE):]
    return dropdir.lstrip("/")

# Check if one directory is contained within another
def in_directory(basedir, childdir):
    #make both absolute    
    basedir = os.path.join(os.path.realpath(basedir),'')
    childdir = os.path.join(os.path.realpath(childdir), '')
    return os.path.commonprefix([basedir, childdir]) == basedir

# Read from a stream, recording result in a record
# Caps the total amount read to ~1MB
def read_stream(outdict, name, stream):
    stream_limit = 1024*1024
    outstr = bytearray()
    while True:
        try:
            chunk = stream.read(stream_limit - len(outstr))
            # End of file reached
            if chunk == b'':
                outdict[name] = bytes(outstr)
                return

            # Add chunk
            outstr.extend(chunk)

            if len(outstr) == stream_limit:
                outstr += b"\n ... Output truncated\n"
                outdict[name] = bytes(outstr)
                # Throw away rest of stream
                while stream.read(1024) != b'':
                    pass
                return
        except IOError:
            outdict[name] = bytes(outstr)


# This is an internal function, use run_program below
def run_program_popen(program, stdin, extra_env):
    env_copy = os.environ.copy()

    if extra_env is not None:
        for key in extra_env:
            env_copy[key] = extra_env[key]

    try:
        proc = subprocess.Popen(program, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE, close_fds=True, shell=False,
                                cwd=SUBMISSIONBASE, env=env_copy, preexec_fn=os.setsid)

        starttime = time.time()

        retdict = dict()

        thread_err = Thread(target=read_stream, args=(retdict, "stderr", proc.stderr))
        thread_err.start()
        thread_out = Thread(target=read_stream, args=(retdict, "stdout", proc.stdout))
        thread_out.start()

        if stdin is not None:
            stdinfd = open(stdin, "rb")
            try:
                for line in stdinfd.readlines():
                    proc.stdin.write(line)
                    time.sleep(0.002)
            except IOError:
                pass
            stdinfd.close()
        # Either we have filled stdin, or we are putting nothing in it
        proc.stdin.close()

        while (proc.poll() is None) and (time.time() - starttime < float(CONFIG['timeout'])):
            time.sleep(0.1)

        timeouthit = False
        if (proc.poll() is None):
            output_str("!! " + CONFIG['timeout'] + " second timeout reached\n")
            timeouthit = True
            os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
            proc.terminate()
            proc.kill()

        thread_err.join()
        thread_out.join()
        proc.wait()
        retdict["stderr"] = toUTF8Replace(retdict["stderr"])
        if timeouthit:
            retdict["stderr"] += "\n... Terminating due to timeout\n"
        retdict["stdout"] = toUTF8Replace(retdict["stdout"])


        retdict["returnval"] = proc.returncode

        return retdict
    except OSError as e:
        warn_print(program[0] + " is broken: " + e.strerror)
        return {"returnval": 1, "stdout": "", "stderr": ""}

# Run a program, given as a list [prog, arg1, arg2], with
# an optional file to read as stdin, and optional extra environment variables
def run_program(program, stdin, extra_env):
    with tempfile.TemporaryDirectory() as scratchdir:
        extra_env["SCRATCHDIR"] = scratchdir
        return run_program_popen(program, stdin, extra_env)
    
def run_bash_script(script, stdin, extra_env):
    return run_program(["bash", script], stdin, extra_env)

# Find files in 'directory' matching 'regex'
def files_in_dir_matching_regex(directory, regex):
    try:
        names = [f for f in sorted(os.listdir(directory))
                if re.match(regex, f) and not f.startswith('.')]
        files = [os.path.join(directory, f) for f in names
                if os.path.isfile(os.path.join(directory, f))]
    except OSError as e:
        warn_print("Error reading '" + directory + "': " + e.strerror)
        return []

    return files

# Accept a directory name relative to
def run_tests_recursive(testdir):
    # First check for a build*.sh
    extra_env = {"TESTDIR" : testdir}

    # In this case run all tests
    inside_TESTBASE = in_directory(TESTBASE, testdir)

    # In this case just run build tests
    prefix_of_TESTBASE = in_directory(testdir, TESTBASE)

    if not inside_TESTBASE and not prefix_of_TESTBASE:
        return

    # BUILDSCRIPTS
    buildscripts = files_in_dir_matching_regex(testdir, r'build.*\.sh')

    for buildsh in buildscripts:
        name = nice_name(buildsh)
        output_str("* BUILD TEST - " + name + " : ")
        buildshret = run_bash_script(buildsh, None, extra_env)
        buildshret["name"] = name
        buildshret["path"] = buildsh
        buildshret["type"] = "Build"

        # State there is no fixed input file
        buildshret["inputlines"] = ["No .in input file (search manually)."]

        register_returnval_test(buildshret)
        if buildshret["returnval"] != 0:
            verbose_print("Exiting early due to failed " + buildsh)
            global anyBuildTestFailed
            anyBuildTestFailed = True
            return

    if inside_TESTBASE:
        # TESTSCRIPTS
        testscripts = files_in_dir_matching_regex(testdir, r'test.*\.sh')

        for test in testscripts:
            name = nice_name(test)
            output_str("* TEST - " + name + " : ")
            result = run_bash_script(test, None, extra_env)
            result["name"] = name
            result["path"] = test
            result["type"] = "Test"

            # State there is no fixed input file
            result["inputlines"] = ["No .in input file (search manually)."]

            register_returnval_test(result)

        # INFOSCRIPTS
        infoscripts = files_in_dir_matching_regex(testdir, r'info.*\.sh')

        for info in infoscripts:
            name = nice_name(info)
            output_str("* INFO - " + name + " : ")
            result = run_bash_script(info, None, extra_env)
            result["name"] = name
            result["path"] = info
            result["type"] = "Info"
            result["alwaysoutput"] = True

            # State there is no fixed input file
            result["inputlines"] = ["No .in input file (search manually)."]

            register_returnval_test(result)

        # PROGSCRIPTS
        progscripts = files_in_dir_matching_regex(testdir, r'prog.*\.sh')

        outtextfiles = files_in_dir_matching_regex(testdir, r'.*\.out')

        # If there isn't exactly one .out file per .sh, warn
        if (len(outtextfiles) == 0) != (len(progscripts) == 0):
            if len(outtextfiles) == 0:
                warn_print("Found prog*.sh without *.out files in " + testdir)
            else:
                warn_print("Found *.out files without prog*.sh in " + testdir)
        for progsh in progscripts:
            for out in outtextfiles:
                infile = out[:-4] + ".in"
                if not os.path.isfile(infile):
                    infile = None
                name = nice_name(progsh) + "-" + os.path.basename(out)
                output_str("* COMPARISON TEST - " + name + " : ")
                result = run_bash_script(progsh, infile, extra_env)
                result["name"] = name
                result["input"] = infile
                result["path"] = progsh
                result["type"] = "Prog"

                # Add contents of input file to result (DOESN'T WORK UNLESS THERE'S AN .IN FILE)
                if (result["input"] != None):
                    with open(result["input"], 'rb') as stream:
                        inputlines = toUTF8Replace(stream.read())

                    inputlines = strip_string(inputlines)
                    result["inputlines"] = inputlines
                else:
                    result["inputlines"] = None

                # Add diff for html to result
                register_diff_test(result, out)
        
    # Always search in subdirectories
    # Run tests for subdirectories
    try:
        subdirs = [os.path.join(testdir, d) for d in sorted(os.listdir(testdir))
                if os.path.isdir(os.path.join(testdir, d))]

        for d in subdirs:
            run_tests_recursive(d)
    except OSError as e:
        warn_print(str(e))

def extract_to_temp(sub_found, temp_dir):
    sub_found.extractall(path=temp_dir)
    sub_found.close()

    # Note: this works because both tarfile and zipfile libraries have these functions. If
    # something else is added that doesn't have these functions it will need changing.

def extract_file_with_submission(archive_name, temp_dir):
    files_with_srcdir = []
    objects_with_srcdir = []

    file_object = None
    content_list = None

    current_path = os.path.abspath(os.getcwd())
    
    # Check that the file specified with --archive exists in cwd
    if not os.path.exists(os.path.join(current_path, ARCHIVE)):
        print("- File not found: " + archive_name);
        sys.exit(1)
    # Make object and get list of contents
    if '.zip' in archive_name:
        file_object = zipfile.ZipFile(archive_name, mode='r')
        content_list = file_object.namelist()
    elif '.tar.xz' in archive_name:
        file_object = tarfile.open(name=archive_name, mode='r:xz')
        content_list = file_object.getnames()
    elif '.tar.gz' in archive_name:
        file_object = tarfile.open(name=archive_name, mode='r:gz')
        content_list = file_object.getnames()
    elif '.tar.bz2' in archive_name:
        file_object = tarfile.open(name=archive_name, mode='r:bz2')
        content_list = file_object.getnames()
    elif '.tar' in archive_name:
        file_object = tarfile.open(name=archive_name, mode='r')
        content_list = file_object.getnames()
    else:
        print("- File {} is not a recognised type of archive (only zip and tar files are supported)".format(ARCHIVE))
        sys.exit(1)

    # If submission folder is inside (or looks like it), extract. Otherwise move on
    submission_name = CONFIG['srcdir'] + "/"
    foundMatch = False

    for content in content_list:
        if submission_name in content:
            foundMatch = True
            break
    if foundMatch == False:
        file_object.close()
        print("- Couldn't find a directory called " + CONFIG['srcdir'] + " in the archive " + archive_name)
        sys.exit(1)

    # print("- Extracting {} to {}.".format(files_with_srcdir[0], temp_dir))    # useful for testing
    print("- Extracting {} to temp folder".format(archive_name))
    extract_to_temp(file_object, temp_dir)
    file_object.close()


##################################################################
# Main program
def run():
    global VERBOSE, TESTBASE, SUBMISSIONBASE, PARENTBASE, TRYHARDER, ARCHIVE
    # parser = argparse.ArgumentParser(usage="%prog [options] testdirectory")
    parser = argparse.ArgumentParser(prog="PROG")
    parser.add_argument("--id", dest="subid", default="<unknown>",
                      help="Give identifier for submission")
    parser.add_argument("--html", dest="htmlout",
                      help="Produce HTML overview", metavar="FILE")                 
    # This isn't documented, because the output format isn't (yet) stable.
    # You can use this, but don't complain if the format changes in future.
    parser.add_argument("--json", dest="jsonout",
                      help=SUPPRESS_HELP, metavar="FILE")
    parser.add_argument("--tryharder", dest="tryharder", default=False,
                      action="store_true", help=SUPPRESS_HELP)
    parser.add_argument("-v", "--verbose",
                      action="store_true", dest="verbose", default=False,
                      help="Output more information during testing")
    parser.add_argument("--version", action="version", version=VERSION)
    parser.add_argument("--archive", dest="archive",
                      help="Look for .zip/.tar/.tar.gz/.tar.xz submission with given name")
    parser.add_argument("testdir")

    args = parser.parse_args()

    # if len(args) != 1:
    if args.testdir is None:
        print("stacscheck " + VERSION)
        sys.exit("You must give a directory of tests to run!")

    if args.htmlout is not None and not jinja:
        print("stacscheck " + VERSION)
        sys.exit("Can't output html without the 'jinja2' library. Exiting.\nYou could try 'pip install jinja2'?")

    VERBOSE = args.verbose
    TRYHARDER = args.tryharder
    ARCHIVE = args.archive

    if not os.path.exists(args.testdir):
        print("stacscheck " + VERSION)
        print("There is no directory called '" + args.testdir + "'")
        print("Check you have the test directory name correct!")
        sys.exit(1)

    if not os.path.isdir(args.testdir):
        print("stacscheck " + VERSION)
        print("'" + args.testdir + "' is not a directory")
        print("Check you have the test directory name correct!")
        sys.exit(1)

    # Get the absolute path instead of the relative path
    TESTBASE = os.path.realpath(args.testdir)

    # If config file isn't found, try looking up. store where you find it
    PARENTBASE = try_parse_config_file(TESTBASE)

    if not 'srcdir' in CONFIG:
        print("A srcdir needs to be specified in practical.config.")
        sys.exit(1)

    if CONFIG['course'] != '' or CONFIG['practical'] != '':
        print("Testing " + CONFIG['course'] + " " + CONFIG['practical'])

    try:
        # If a .zip/.tar/other name was specified, skip normal searching
        if args.archive is None:
            SUBMISSIONBASE = find_code_directory()
            run_tests_recursive(PARENTBASE)
        # Try running with compressed file specified with --archive
        else:
            current_path = os.path.abspath(os.getcwd())
            # Extract files and run tests
            with tempfile.TemporaryDirectory() as temp_dir:
                extract_file_with_submission(ARCHIVE, temp_dir)
                SUBMISSIONBASE = find_code_directory(temp_dir)
                run_tests_recursive(PARENTBASE)

        if len(testStore) == 0:
            # Use args.testdir as it is shorter
            print("ERROR: No tests found in '" + args.testdir + "'")
    except EnvironmentError as e:
        print("Error during testing: " + str(e))

    finalmessage = str(len([t for t in testStore if t["pass"]])) + " out of " + str(len(testStore)) + " tests passed"
    print(finalmessage)

    if anyBuildTestFailed:
        print("Building is failing, so some tests were skipped")


    try:
        if args.htmlout is not None:
            env = jinja2.Environment(autoescape=True)
            template = env.from_string(JINJAPAGE)
            with open(args.htmlout, "wb") as html:
                htmloutput = template.render(Practical=CONFIG['course'] + " " + CONFIG['practical'],
                                        SubmissionID=args.subid,
                                        Tests=testStore,
                                        Finalmessage=finalmessage,
                                        AnyBuildTestFailed=anyBuildTestFailed,
                                        InfoText=INFO_PRINTED
                                        )
                # html.write(toUTF8Replace(htmloutput).encode("utf8"))
                html.write(htmloutput.encode("utf8"))
    except EnvironmentError as e:
        print("Error outputting HTML: " + str(e))

    try:
        if args.jsonout is not None:
            with open(args.jsonout, "w") as jsonfile:
                jsonoutput = json.dumps({ "practical" : CONFIG['practical'],
                                            "course" : CONFIG['course'],
                                            "subid" : args.subid,
                                            "tests" : testStore,
                                            "infotext" : INFO_PRINTED},
                                            indent = 2, separators=(',', ': '), ensure_ascii=True, sort_keys=True)
                jsonfile.write(jsonoutput)
                jsonfile.write("\n")
    except EnvironmentError as e:
        print("Error outputting json: " + str(e))

if __name__ == "__main__":
    try:
        run()
    except EnvironmentError as e:
        print("A fatal error has occurred: " + str(e))
