#! /usr/bin/env python2
#
# Perforce benchmarks using Locust.io framework
#
# Copyright (C) 2016, Robert Cowham, Perforce
#

from __future__ import print_function

import os
import sys
import string
import logging
import platform
import errno
from locust import Locust, events, task, TaskSet
from locust.exception import StopLocust
import subprocess
import P4
import yaml
import time
import random
import pprint
from createfiles import create_file, generator

from locust.stats import RequestStats
def noop(*arg, **kwargs):
    logger.info("Stats reset prevented by monkey patch!")
RequestStats.reset_all = noop

python3 = sys.version_info[0] >= 3

logger = logging.getLogger("repo_benchmark")

startdir = os.getcwd()

CONFIG_FILE = "config_p4_basic.yml"   # Configuration parameters

ENCODING = None
if hasattr(sys.stdin, 'encoding'):
    ENCODING = sys.stdin.encoding
if ENCODING is None:
    import locale
    locale_name, ENCODING = locale.getdefaultlocale()
if ENCODING is None:
    ENCODING = "ISO8859-1"

def popen(exe, cmd, decode=False, errors=True):
    cmd.insert(0, exe)
    pipe = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (stdout, stderr) = pipe.communicate()
    if errors and pipe.returncode > 0:
        raise Exception((stderr + stdout).decode(ENCODING))
    return stdout if not decode else stdout.decode(ENCODING)

def readConfig(startdir):
    config = {}
    with open(os.path.join(startdir, CONFIG_FILE), "r") as f:
        config = yaml.load(f)
    return config

def fmtsize(num):
    for x in ['bytes', 'KB', 'MB', 'GB', 'TB']:
        if num < 1024.0:
            return "%3.1f %s" % (num, x)
        num /= 1024.0

LINE_LENGTH = 80
BLOCK_SIZE = 65536

class Timer(object):
    def __init__(self, request_type):
        self.start_time = time.time()
        self.request_type = request_type

    def report_failure(self, name, e):
        total_time = int((time.time() - self.start_time) * 1000)
        events.request_failure.fire(request_type=self.request_type, name=name, response_time=total_time, exception=e)

    def report_success(self, name, count):
        total_time = int((time.time() - self.start_time) * 1000)
        events.request_success.fire(request_type=self.request_type, name=name, response_time=total_time, response_length=count)

class SyncOutput(P4.OutputHandler):
    "Log sync progress"

    def __init__(self, timer, event_name, sync_progress_size_interval=100000):
        P4.OutputHandler.__init__(self)
        self.timer = timer
        self.event_name = event_name
        self.sync_progress_size_interval = sync_progress_size_interval
        self.filesSynced = 0
        self.sizeSynced = 0
        self.previousSizeSynced = 0

    def reportFileSync(self, fileSize):
        self.filesSynced += 1
        self.sizeSynced += fileSize
        if not self.sync_progress_size_interval:
            return
        if self.sizeSynced > self.previousSizeSynced + self.sync_progress_size_interval:
            self.previousSizeSynced = self.sizeSynced
            self.timer.report_success(self.event_name, 1)
            logger.info("Synced %d files, size %s" % (self.filesSynced, fmtsize(self.sizeSynced)))

    def outputStat(self, stat):
        if 'fileSize' in stat:
            self.reportFileSync(int(stat['fileSize']))
        return P4.OutputHandler.HANDLED

class P4Benchmark(object):
    """Generic benchmarking class - must be subclassed"""

    def __init__(self, startdir, config, prog="p4_bench"):
        self.id = id(self)      # Object id - unique enough for now.
        if "workspace_root" in config["general"]:
            self.test_root = config["general"]["workspace_root"]
        else:
            self.test_root = os.path.join(startdir, "testdata")
        self.localfilelist = []
        os.environ["P4CONFIG"] = os.path.join(startdir, "bench_p4config.txt")
        self.p4 = P4.P4()
        self.p4.prog = prog     # Identifier for log analysis
        self.p4.logger = logger
        self.config = config
        p4config = config["perforce"]
        logger.info(pprint.pformat(p4config))
        if isinstance(p4config["port"], list):
            self.p4.port = random.choice(p4config["port"])
        else:
            if python3:
                if isinstance(p4config["port"], (str)):
                    self.p4.port = p4config["port"]
                else:
                    raise Exception("Unknown port config")
            else:
                if isinstance(p4config["port"], (str, unicode)):
                    self.p4.port = p4config["port"]
                else:
                    raise Exception("Unknown port config")
        logger.info("Connecting to server: %s" % self.p4.port)
        self.p4.prog = "%s-%s" % (prog, self.p4.port)
        if "charset" in p4config and p4config["charset"]:
            self.p4.charset = p4config["charset"]
        self.repoPath = p4config["repoPath"]
        self.p4.user = p4config["user"]
        self.sync_progress_size_interval = 1000 * 1000 * 1000   # 1GB
        if "sync_progress_size_interval" in p4config:
            self.sync_progress_size_interval = int(eval(p4config["sync_progress_size_interval"]))
        result = self.p4.connect()
        if p4config["password"]:
            self.p4.password = p4config["password"]
            self.p4.run_login()
        self.workspace_name = "{}.{}.{}".format(self.p4.user, self.id, platform.node())
        self.workspace_root = os.path.join(self.test_root, self.workspace_name)
        if not os.path.isdir(self.test_root):
            try:
                os.makedirs(self.test_root, 0o777)
            except OSError as ex:
                if ex.errno != errno.EEXIST:
                    raise ex
                pass

    def getView(self):
        "Randomly select between dirs, or as directed by config file"
        p4config = self.config["perforce"]
        return ["{}/... //{}/...".format(p4config["repoPath"], self.workspace_name)]

    def createWorkspace(self):
        p4config = self.config["perforce"]
        # Important to set client before attempting to create it in case we are talking to replica
        self.p4.client = self.workspace_name
        ws = self.p4.fetch_client(self.workspace_name)
        existed = (ws._root == self.workspace_root)
        ws._root = self.workspace_root
        if "stream" in p4config and p4config["stream"] == "y":
            logger.warn("Overwiting options")
            ws._stream = p4config["repoPath"]
        else:
            ws._view = self.getView()
        ws._options = ws._options.replace("normdir", "rmdir")
        if "options" in p4config and p4config["options"]:
            logger.warn("Overwiting options")
            ws._options = p4config["options"]
        logger.info("Saving workspace view: %s" % ws._view[0])
        logger.info(pprint.pformat(ws))
        result = self.p4.save_client(ws)
        logger.info(pprint.pformat(result))
        return existed

    def syncWorkspace(self, timer):
        result = []
        p4config = self.config["perforce"]
        if "sync_args" in p4config and p4config["sync_args"]:
            cmd = "p4"
            args = [p4config["sync_args"], "-p", self.p4.port, "-u", self.p4.user, "-c", self.p4.client]
            args.extend(["sync", "//{}/...".format(self.p4.client)])
            logger.info("%s %s" % (cmd, " ".join(args)))
            result = popen(cmd, args)
            logger.info(result)
        else:
            if not os.path.isdir(self.workspace_root):
                os.makedirs(self.workspace_root, 0o777)
            os.chdir(self.workspace_root)
            event_name = "sync_partial_%s" % fmtsize(self.sync_progress_size_interval)
            syncCallback = SyncOutput(timer, event_name, self.sync_progress_size_interval)
            with self.p4.at_exception_level(P4.P4.RAISE_ERRORS):
                result = self.p4.run_sync("//{}/...".format(self.p4.client), handler=syncCallback)
            havelist = self.p4.run_have()
            self.localfilelist = [f["path"] for f in havelist]
        return len(result)
        

    def addFile(self, filename):
        self.p4.run_add(filename)

    def editFile(self, filename):
        self.p4.run_edit(filename)

    def deleteFile(self, filename):
        self.p4.run_delete(filename)

    def commit(self):
        """Submit our change"""
        opened = self.p4.run_opened()
        if len(opened) > 0:
            with self.p4.at_exception_level(P4.P4.RAISE_ERRORS):
                files = [x['depotFile'] for x in opened]
                self.p4.run_sync(files)
                self.p4.run_resolve("-ay")
            chg = self.p4.fetch_change()
            chg._description = "A test change"
            self.p4.save_submit(chg)
            with self.p4.at_exception_level(P4.P4.RAISE_ERRORS):
                self.p4.run_revert("//...")
    
    def basicFileActions(self):
        """Randomly edit/delete/add files in workspace"""
        numActions = random.randint(1, 20)
        for i in range(numActions):
            action = random.choice(["add", "edit", "edit", "edit", "edit", "delete"])
            filename = random.choice(self.localfilelist)
            if action == "add":
                addfilename = os.path.join(os.path.dirname(filename), generator(20, eol=""))
                create_file(random.randint(100, 1000000), addfilename, random.choice([True, False]))
                self.addFile(addfilename)
            elif action == "edit":
                self.editFile(filename)
                with open(filename, "a") as f:
                    f.write("Some extra data\n")
            else:
                self.deleteFile(filename)
        return numActions

class P4BasicBenchmark(P4Benchmark):
    """Performs basic benchmark test - Perforce specific subclass"""

    def __init__(self, startdir, config):
        super(P4BasicBenchmark, self).__init__(startdir, config, "basicActions")

def basicActions(bench, request_type):
    "Build farm basically just does a sync"
    count = 0
    t = Timer(request_type)
    name = "create"
    try:
        bench.createWorkspace()
        count = 1
        t.report_success(name, count)
    except Exception as e:
        logger.exception(e)
        t.report_failure(name, e)
    t = Timer(request_type)
    name = "sync"
    try:
        count = bench.syncWorkspace(t)
        t.report_success(name, count)
    except Exception as e:
        logger.exception(e)
        t.report_failure(name, e)
    t = Timer(request_type)
    name = "actions"
    try:
        count = bench.basicFileActions()
        t.report_success(name, count)
    except Exception as e:
        logger.exception(e)
        t.report_failure(name, e)

class AllTasks(TaskSet):
    """Entry point for locust"""

    min_wait = 1000
    max_wait = 10000
    request_type = "p4"

    def __init__(self, *args, **kwargs):
        super(AllTasks, self).__init__(*args, **kwargs)
        self.config = readConfig(startdir)
        self.min_wait = self.config["general"]["min_wait"]
        self.max_wait = self.config["general"]["max_wait"]

    def on_start(self):
        pass

    @task(10)
    def basicActions(self):
        task_name = "p4buildfarm"
        self.bench = P4BasicBenchmark(startdir, self.config)
        basicActions(self.bench, task_name)
        logger.info("Finished %s" % task_name)
        raise StopLocust("task_name")   # Run once only and die

class P4RepoTestLocust(Locust):
    """Will be imported and then run by locust"""
    task_set = AllTasks