#!/usr/bin/env python3

import os
import sys

SCRIPT_PATH = os.path.realpath(__file__)
SCRIPT_DIR = os.path.dirname(SCRIPT_PATH)
sys.path.append(os.path.join(SCRIPT_DIR, "third_party"))

import argparse
import datetime
import foundation
from foundation import (benchrunner,
                        buildbot,
                        config,
                        context,
                        logger,
                        system_info,
                        util)
import json
import os
import shutil
import sys
import tempfile


########################################
#  Parser helpers.
########################################


def configureArgumentParser():
    parser = argparse.ArgumentParser(
        description="Cycles benchmark helper script.")
    parser.add_argument("-b", "--blender",
                        help="Full file path to Blender's binary " +
                             "to use for rendering",
                        default="")
    parser.add_argument("-d", "--scenes-dir",
                        help="Directory with scenes",
                        default="")
    parser.add_argument('-c', '--scenes',
                        nargs='+',
                        help='Scenes to be rendered',
                        default=[])
    parser.add_argument('-t', '--device-type',
                        help="Type of the device to render on",
                        default="")
    parser.add_argument('-n', '--device-name',
                        help="Device name to render on",
                        default="")
    parser.add_argument('-o', '--device-name-override',
                        help="Override compute device name with this",
                        default="")
    parser.add_argument('-s', '--device-single',
                        help="Use single device when multiple matches",
                        action='store_true',
                        default=False)
    parser.add_argument('-v', '--verbose',
                        help="Do verbose logging",
                        action='store_true',
                        default=False)
    parser.add_argument('-i', '--system-info',
                        help="Only dump system info to the screen",
                        action='store_true',
                        default=False)
    return parser

########################################
#  Configuration helpers.
########################################


def injectDefaultConfiguration(config):
    """
    For a specified configuration object, set all possible properties to their
    default value.
    """
    root_dir = util.getBundleRootDirectory()
    section = {
        "scenes": "",
        "scenes_dir": os.path.join(root_dir, "scenes"),
        "device_name": "",
        "device_single": "False",
    }
    config['farm'] = section


def injectArgparseConfiguration(config, args):
    """
    Override settings wit harguments passed from the command line.
    """
    section = config['farm']
    if args.blender:
        section['blender'] = args.blender
    if args.scenes:
        section['scenes'] = ",".join(args.scenes)
    if args.scenes_dir:
        section['scenes_dir'] = args.scenes_dir
    if args.device_type:
        section['device_type'] = args.device_type
    if args.device_name:
        section['device_name'] = args.device_name
    if args.device_name_override:
        section['device_name_override'] = args.device_name_override
    if args.device_single:
        section['device_single'] = args.device_single


def readConfiguration(args):
    """
    Read configuration file and return BenchmarkConfig with all the settings
    we will need to use.
    """
    config = foundation.config.BenchmarkConfig()
    injectDefaultConfiguration(config)
    read_configs = config.readGlobalConfig("farm")
    if read_configs:
        logger.INFO("Configuration was read from:")
        for cfg in read_configs:
            print("    " + util.stripSensitiveInfo(cfg))
    injectArgparseConfiguration(config, args)
    return config


def checkConfiguration(config):
    """
    Check whether configuration is complete and usable.
    """
    required_keys = ('device_type', )
    known_device_types = ('CPU', 'GPU', 'CUDA', 'OPENCL', )
    logger.INFO("Validating configuration...")
    # Check whether section exists.
    if 'farm' not in config.sections():
        logger.INFO("    Missing configuration section for 'farm'.")
        return False
    section = config['farm']
    # Check whether required keys exists.
    for required_key in required_keys:
        if required_key not in section:
            logger.INFO("    Missing configuration key {}" .
                        format(required_key))
            return False
    # Check whether device configuration is correct.
    device_type = section['device_type']
    if device_type not in known_device_types:
        logger.INFO("    Unknown device type {}" . format(device_type))
        return False
    if device_type != 'CPU':
        if 'device_name' not in section or not section['device_name']:
            logger.INFO("    Need to explicitly specify device name.")
            return False
    # Check whether directories are correct.
    if not os.path.exists(section["scenes_dir"]):
        logger.INFO("    Scenes directory does not exist.")
        return False
    if "output_dir" not in section:
        logger.INFO("    Missing configuration for output directory.")
        return False
    return True


########################################
#  Temporary directories helpers.
########################################

def createTempDirectory(config):
    """
    Usually creates an unique temporary directory and returns full path to it.

    In the development environment uses pre-defined directory, where all data
    might already be existing (for example, blender might be already downloaded
    and unpacked).
    """
    temp_dir = tempfile.mkdtemp(prefix="blender-benchmark-")
    return temp_dir


def deleteTempDirectory(config, temp_dir):
    """
    Get rid of automatically created temp directory/
    """
    shutil.rmtree(temp_dir)


########################################
#  Latets Blender helper script.
########################################

def downloadLatestBlender(directory):
    """
    Download latets Blender from buildbot to given directory.
    """
    platform = sys.platform
    latest_blender_url = buildbot.buildbotGetLatetsVersion(platform, "64bit")
    if not latest_blender_url:
        logger.ERROR("Unable to figure out latest Blender version")
        return
    blender_filename = latest_blender_url.split('/')[-1]
    logger.INFO("Found latest Blender: {}" . format(blender_filename))
    local_blender = os.path.join(directory, blender_filename)
    logger.INFO("Downloading Blender...")
    util.downloadFile(latest_blender_url, local_blender)
    return local_blender


def findBlenderDirInDirectory(directory):
    """
    Find Blender directory in given folder.
    """
    for filename in os.listdir(directory):
        if filename.startswith("blender"):
            full_filename = os.path.join(directory, filename)
            if os.path.isdir(full_filename):
                return full_filename
    return None


def findBlenderInDirectory(directory):
    """
    Find blender executable file in given directory. Will recurse into
    Blender directory in there
    """
    blender_dir = findBlenderDirInDirectory(directory)
    platform = sys.platform
    if platform == 'linux':
        return os.path.join(blender_dir, "blender")
    elif platform == "win32":
        return os.path.join(blender_dir, "blender.exe")
    else:
        raise Exception("Need to support your OS!")


def getLatetsBlenderBinary(config, temp_dir):
    """
    Get full file path to latest Blender executable which will be used for
    actual benchmark.

    It will either use Blender from already specified location or it will
    download latest Blender from buildbot.
    """
    # Firts try to use Blender specified in configuration.
    if 'blender' in config['farm']:
        return config['farm']['blender']
    # Well, download and unpack the latest Blender from buildbot.
    logger.INFO("Will get latest Blender from buildbot.")
    blender_archive = downloadLatestBlender(temp_dir)
    logger.INFO("Unpacking Blender...")
    util.unpackArchive(blender_archive, temp_dir)
    return findBlenderInDirectory(temp_dir)

########################################
#  Results output.
########################################


def latestDirGet(path):
    """
    Get directory with bigger number in the given path
    """
    max_file = None
    for f in os.listdir(path):
        if f.endswith(".DS_Store"):
            continue
        x = int(f)
        if not max_file or x > max_file:
            max_file = x
    return max_file


def ensureOutputDir(config):
    output_dir = config['farm']['output_dir']
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    max_dir = latestDirGet(output_dir)
    if not max_dir:
        max_dir = 0
    new_dir = str(max_dir + 1).zfill(8)
    new_full_dir = os.path.join(output_dir, new_dir)
    os.mkdir(new_full_dir)
    return new_full_dir


def ensureImageOutputDir(results_output_dir):
    images_output_dir = os.path.join(results_output_dir, "images")
    os.mkdir(images_output_dir)
    return images_output_dir


def getResultJSONString(ctx, results):
    # Convert custom classes to dictionaries for easier JSON dump.
    json_results = results
    stats = json_results['stats']
    for scene in ctx.scenes:
        if scene not in stats:
            continue
        if stats[scene]:
            stats[scene] = stats[scene].asDict()
            stats[scene]['result'] = 'OK'
        else:
            stats[scene] = {'result': 'CRASH'}
    return json.dumps(json_results, sort_keys=True, indent=2)


def saveResults(ctx, results, output_dir):
    json_string = getResultJSONString(ctx, results)
    results_file = os.path.join(output_dir, "results.json")
    with open(results_file, "w") as f:
        f.write(json_string)

########################################
#  Main logic.
########################################


def main():
    parser = configureArgumentParser()
    args = parser.parse_args()
    logger.VERBOSE = args.verbose
    logger.init()
    logger.HEADER("Cycles Benchmark Suite v{}, farm edition" .
                  format(foundation.VERSION))
    # Some platform independent directories to helper scripts.
    script_directory = os.path.dirname(os.path.realpath(__file__))
    configure_script = os.path.join(script_directory, "configure.py")
    # Read configuration file, so we know what we will be doing.
    config = readConfiguration(args)
    if not checkConfiguration(config):
        logger.ERROR("Configuration is not complete or valid, aborting.")
        return False
    logger.INFO("Configuration looks reasonable, continuing.")
    # Create temporary directory, all runtime files will be stored there.
    temp_dir = createTempDirectory(config)
    results_output_dir = ensureOutputDir(config)
    images_output_dir = ensureImageOutputDir(results_output_dir)
    try:
        blender_binary = getLatetsBlenderBinary(config, temp_dir)
        logger.INFO("Will benchmark the following Blender: {}".
                    format(util.stripSensitiveInfo(blender_binary)))
        # Bechmark context initialization.
        farm_config = config['farm']
        ctx = context.Context()
        ctx.blender = blender_binary
        ctx.configure_script = configure_script
        ctx.scenes_dir = farm_config['scenes_dir']
        ctx.device_type = farm_config['device_type']
        ctx.device_name = farm_config['device_name']
        if farm_config['device_single'].lower() == "true":
            ctx.device_single = True
        else:
            ctx.device_single = False
        ctx.image_output_dir = images_output_dir
        # Don't do any scenes when only querying for system information.
        if args.system_info:
            ctx.scenes = []
        else:
            if farm_config['scenes']:
                ctx.scenes = farm_config['scenes'].split(",")
            else:
                ctx.scenes = ctx.listAllScenes(ctx.scenes_dir)
        # Print additional informaiton.
        if farm_config["device_name_override"]:
            print("Using device name override: {}" .
                          format(farm_config["device_name_override"]))
        # Print prelmiinary information.
        blender_dvice_info = benchrunner.benchmarkGetDeviceInfo(ctx)
        if not blender_dvice_info['device_type']:
            logger.ERROR("Requested device can not be enabled in Blender.")
            logger.INFO("Requested device details:")
            logger.INFO("    Device type: {}" . format(ctx.device_type))
            logger.INFO("    Device name: {}" . format(ctx.device_name))
            all_stats = None
        else:
            logger.INFO("Configured device details:")
            logger.INFO("    Device type: {}" .
                        format(blender_dvice_info["device_type"]))
            logger.INFO("    Compute devices:")
            for compute_device in blender_dvice_info["compute_devices"]:
                logger.INFO("        {}" . format(compute_device))
            # Run benchmark.
            all_stats = benchrunner.benchmarkAll(ctx)
        # Gather all information together.
        timestamp = datetime.datetime.now(datetime.timezone.utc).isoformat()
        results = {
            "timestamp": timestamp,
            "blender_version": system_info.getBlenderVersion(ctx),
            "system_info": system_info.gatherSystemInfo(ctx),
            "device_info": blender_dvice_info,
            "stats": all_stats if all_stats else {}
        }
        if farm_config["device_name_override"]:
            results["device_name_override"] = farm_config["device_name_override"]
        saveResults(ctx, results, results_output_dir)
        return all_stats is not None
    finally:
        if not args.system_info:
            deleteTempDirectory(config, temp_dir)
    return True


if __name__ == "__main__":
    if not main():
        sys.exit(1)