# copyright 2003-2023 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
# contact https://www.logilab.fr/ -- mailto:contact@logilab.fr
#
# This file is part of CubicWeb.
#
# CubicWeb is free software: you can redistribute it and/or modify it under the
# terms of the GNU Lesser General Public License as published by the Free
# Software Foundation, either version 2.1 of the License, or (at your option)
# any later version.
#
# CubicWeb is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License along
# with CubicWeb.  If not, see <https://www.gnu.org/licenses/>.
"""Usage: %s [OPTIONS] <instance id> <queries file>

Stress test a CubicWeb repository

OPTIONS:
  -h / --help
     Display this help message and exit.

  -u / --user <user>
     Connect as <user> instead of being prompted to give it.
  -p / --password <password>
     Automatically give <password> for authentication instead of being prompted
     to give it.

  -n / --nb-times <num>
     Repeat queries <num> times.
  -t / --nb-threads <num>
     Execute queries in <num> parallel threads.
  -P / --profile <prof_file>
     dumps profile results (hotshot) in <prof_file>
  -o / --report-output <filename>
     Write profiler report into <filename> rather than on stdout

Copyright (c) 2003-2023 LOGILAB S.A. (Paris, FRANCE), license is LGPL v2.
https://www.logilab.fr/ -- mailto:contact@logilab.fr
"""
import getopt
import os
import sys
import threading
import traceback
from getpass import getpass
from os.path import basename
from time import process_time

from logilab.common.fileutils import lines
from logilab.common.ureports import Table, TextWriter

from cubicweb import repoapi
from cubicweb.server.repository import Repository

TB_LOCK = threading.Lock()


class QueryExecutor:
    def __init__(self, cnx, times, queries, reporter=None):
        self._cnx = cnx
        self._times = times
        self._queries = queries
        self._reporter = reporter

    def run(self):
        with self._cnx as cnx:
            times = self._times
            while times:
                for index, query in enumerate(self._queries):
                    start = process_time()
                    try:
                        cnx.execute(query)
                    except Exception:
                        TB_LOCK.acquire()
                        traceback.print_exc()
                        TB_LOCK.release()
                        return
                    if self._reporter is not None:
                        self._reporter.add_proftime(process_time() - start, index)
                times -= 1


def usage(status=0):
    """print usage string and exit"""
    print(__doc__ % basename(sys.argv[0]))
    sys.exit(status)


class ProfileReporter:
    """a profile reporter gathers all profile informations from several
    threads and can write a report that summarizes all profile informations
    """

    profiler_lock = threading.Lock()

    def __init__(self, queries):
        self._queries = tuple(queries)
        self._profile_results = [(0.0, 0)] * len(self._queries)
        # self._table_report = Table(3, rheaders = True)
        len_max = max([len(query) for query in self._queries]) + 5
        self._query_fmt = "%%%ds" % len_max

    def add_proftime(self, elapsed_time, query_index):
        """add a new time measure for query"""
        ProfileReporter.profiler_lock.acquire()
        cumul_time, times = self._profile_results[query_index]
        cumul_time += elapsed_time
        times += 1.0
        self._profile_results[query_index] = (cumul_time, times)
        ProfileReporter.profiler_lock.release()

    def dump_report(self, output=sys.stdout):
        """dump report in 'output'"""
        table_elems = ["RQL Query", "Times", "Avg Time"]
        total_time = 0.0
        for query, (cumul_time, times) in zip(self._queries, self._profile_results):
            avg_time = cumul_time / float(times)
            table_elems += [str(query), f"{times:f}", f"{avg_time:f}"]
            total_time += cumul_time
        table_elems.append("Total time :")
        table_elems.append(str(total_time))
        table_elems.append(" ")
        table_layout = Table(3, rheaders=True, children=table_elems)
        TextWriter().format(table_layout, output)
        # output.write('\n'.join(tmp_output))


def run(args):
    """run the command line tool"""
    try:
        opts, args = getopt.getopt(
            args,
            "hn:t:u:p:P:o:",
            [
                "help",
                "user=",
                "password=",
                "nb-times=",
                "nb-threads=",
                "profile",
                "report-output=",
            ],
        )
    except Exception as ex:
        print(ex)
        usage(1)
    repeat = 100
    threads = 1
    user = os.environ.get("USER", os.environ.get("LOGNAME"))
    password = None
    report_output = sys.stdout
    prof_file = None
    for opt, val in opts:
        if opt in ("-h", "--help"):
            usage()
        if opt in ("-u", "--user"):
            user = val
        elif opt in ("-p", "--password"):
            password = val
        elif opt in ("-n", "--nb-times"):
            repeat = int(val)
        elif opt in ("-t", "--nb-threads"):
            threads = int(val)
        elif opt in ("-P", "--profile"):
            prof_file = val
        elif opt in ("-o", "--report-output"):
            report_output = open(val, "w")
    if len(args) != 2:
        usage(1)
    queries = [query for query in lines(args[1]) if not query.startswith("#")]
    if user is None:
        user = input("login: ")
    if password is None:
        password = getpass("password: ")
    from cubicweb.cwconfig import instance_configuration

    config = instance_configuration(args[0])
    # get local access to the repository
    print("Creating repo", prof_file)
    repo = Repository(config, prof_file)
    repo.bootstrap()
    cnx = repoapi.connect(repo, user, password=password)
    reporter = ProfileReporter(queries)
    if threads > 1:
        executors = []
        while threads:
            qe = QueryExecutor(cnx, repeat, queries, reporter=reporter)
            executors.append(qe)
            thread = threading.Thread(target=qe.run)
            qe.thread = thread
            thread.start()
            threads -= 1
        for qe in executors:
            qe.thread.join()
    else:
        QueryExecutor(cnx, repeat, queries, reporter=reporter).run()
    reporter.dump_report(report_output)


if __name__ == "__main__":
    run(sys.argv[1:])
