# -*- coding: utf-8 -*-
"""
A pytest plugin to trace resource leaks.
"""
from __future__ import print_function, division

import collections
import gc
import time
import os
import sys
import threading

import pytest


def pytest_addoption(parser):
    group = parser.getgroup('resource leaks')
    group.addoption(
        '-L', '--leaks',
        action='store',
        dest='leaks',
        help='''\
List of resources to monitor for leaks before and after each test.
Can be 'all' or a comma-separated list of resource names
(possible values: {known_checkers}).
'''.format(known_checkers=', '.join(sorted("'%s'" % s for s in all_checkers)))
    )
    group.addoption(
        '--leaks-timeout',
        action='store',
        type='float',
        dest='leaks_timeout',
        default=0.5,
        help='''\
Wait at most this number of seconds to mark a test leaking
(default: %(default)s).
'''
    )
    group.addoption(
        '--leaks-fail',
        action='store_true',
        dest='leaks_mark_failed',
        default=False,
        help='''Mark leaked tests failed.'''
    )
    group.addoption(
        '--leak-retries',
        action='store',
        type=int,
        dest='leak_retries',
        default=1,
        help='''\
Max number of times to retry a test when it leaks, to ignore
warmup-related issues (default: 1).
'''
    )


def pytest_configure(config):
    leaks = config.getvalue('leaks')
    if leaks:
        if leaks == 'all':
            leaks = sorted(all_checkers)
        else:
            leaks = leaks.split(',')
        unknown = sorted(set(leaks) - set(all_checkers))
        if unknown:
            raise ValueError("unknown resources: %r" % (unknown,))

        checkers = [all_checkers[leak]() for leak in leaks]
        checker = LeakChecker(checkers=checkers,
                              grace_delay=config.getvalue('leaks_timeout'),
                              mark_failed=config.getvalue('leaks_mark_failed'),
                              max_retries=config.getvalue('leak_retries'),
                              )
        config.pluginmanager.register(checker, 'leaks_checker')


all_checkers = {}


def register_checker(name):
    def decorate(cls):
        assert issubclass(cls, ResourceChecker), cls
        assert name not in all_checkers
        all_checkers[name] = cls
        return cls

    return decorate


class ResourceChecker(object):

    def on_start_test(self):
        pass

    def on_stop_test(self):
        pass

    def on_retry(self):
        pass

    def measure(self):
        raise NotImplementedError

    def has_leak(self, before, after):
        raise NotImplementedError

    def format(self, before, after):
        raise NotImplementedError


@register_checker('fds')
class FDChecker(ResourceChecker):

    def measure(self):
        if os.name == 'posix':
            import psutil
            return psutil.Process().num_fds()
        else:
            return 0

    def has_leak(self, before, after):
        return after > before

    def format(self, before, after):
        return "leaked %d file descriptor(s)" % (after - before)


@register_checker('memory')
class RSSMemoryChecker(ResourceChecker):

    def measure(self):
        import psutil
        return psutil.Process().memory_info().rss

    def has_leak(self, before, after):
        return after > before + 1e7

    def format(self, before, after):
        return "leaked %d MB of RSS memory" % ((after - before) / 1e6)


@register_checker('threads')
class ActiveThreadsChecker(ResourceChecker):

    def measure(self):
        return set(threading.enumerate())

    def has_leak(self, before, after):
        return not after <= before

    def format(self, before, after):
        leaked = after - before
        assert leaked
        return ("leaked %d Python threads: %s"
                % (len(leaked), sorted(leaked, key=str)))


class _ChildProcess(collections.namedtuple('_ChildProcess',
                                           ('pid', 'name', 'cmdline'))):

    @classmethod
    def from_process(cls, p):
        return cls(p.pid, p.name(), p.cmdline())


@register_checker('processes')
class ChildProcessesChecker(ResourceChecker):

    def measure(self):
        import psutil
        # We use pid and creation time as keys to disambiguate between
        # processes (and protect against pid reuse)
        # Other properties such as cmdline may change for a given process
        children = {}
        p = psutil.Process()
        for c in p.children(recursive=True):
            try:
                with c.oneshot():
                    if c.ppid() == p.pid and os.path.samefile(c.exe(), sys.executable):
                        cmdline = c.cmdline()
                        if any(a.startswith('from multiprocessing.semaphore_tracker import main')
                               for a in cmdline):
                            # Skip multiprocessing semaphore tracker
                            continue
                        if any(a.startswith('from multiprocessing.forkserver import main')
                               for a in cmdline):
                            # Skip forkserver process, the forkserver's children
                            # however will be recorded normally
                            continue
                    children[(c.pid, c.create_time())] = _ChildProcess.from_process(c)
            except psutil.NoSuchProcess:
                pass
        return children

    def has_leak(self, before, after):
        return not set(after) <= set(before)

    def format(self, before, after):
        leaked = set(after) - set(before)
        assert leaked
        formatted = []
        for key in sorted(leaked):
            p = after[key]
            formatted.append('  - pid={p.pid}, name={p.name!r}, cmdline={p.cmdline!r}'
                             .format(p=p))
        return ("leaked %d processes:\n%s"
                % (len(leaked), '\n'.join(formatted)))


@register_checker('tracemalloc')
class TracemallocMemoryChecker(ResourceChecker):

    def __init__(self):
        global tracemalloc
        import tracemalloc

    def on_start_test(self):
        tracemalloc.start(1)

    def on_stop_test(self):
        tracemalloc.stop()

    def measure(self):
        import tracemalloc
        current, peak = tracemalloc.get_traced_memory()
        snap = tracemalloc.take_snapshot()
        return current, snap

    def has_leak(self, before, after):
        return after[0] > before[0] + 1e6

    def format(self, before, after):
        bytes_before, snap_before = before
        bytes_after, snap_after = after
        diff = snap_after.compare_to(snap_before, 'traceback')
        ndiff = 5
        min_size_diff = 2e5

        lines = []
        lines += ["leaked %.1f MB of traced Python memory"
                  % ((bytes_after - bytes_before) / 1e6)]
        for stat in diff[:ndiff]:
            size_diff = stat.size_diff or stat.size
            if size_diff < min_size_diff:
                break
            count = stat.count_diff or stat.count
            lines += ["  - leaked %.1f MB in %d calls at:" % (size_diff / 1e6, count)]
            lines += ["    " + line for line in stat.traceback.format()]

        return "\n".join(lines)


class LeakChecker(object):
    def __init__(self, checkers, grace_delay, mark_failed, max_retries):
        self.checkers = checkers
        self.grace_delay = grace_delay
        self.mark_failed = mark_failed
        self.max_retries = max_retries

        # {nodeid: {checkers}}
        self.skip_checkers = {}
        # {nodeid: {checker: [(before, after)]}}
        self.counters = {}
        # {nodeid: [(checker, before, after)]}
        self.leaks = {}
        # {nodeid: {outcomes}}
        self.outcomes = collections.defaultdict(set)

        # Reentrancy guard
        self._retrying = False

    def cleanup(self):
        gc.collect()

    def checks_for_item(self, nodeid):
        return [c for c in self.checkers
                if c not in self.skip_checkers.get(nodeid, ())]

    def measure(self, nodeid):
        # Return items in order
        return [(c, c.measure()) for c in self.checks_for_item(nodeid)]

    def measure_before_test(self, nodeid):
        for checker in self.checks_for_item(nodeid):
            checker.on_start_test()
        for checker, before in self.measure(nodeid):
            assert before is not None
            self.counters[nodeid][checker].append((before, None))

    def measure_after_test(self, nodeid):
        outcomes = self.outcomes[nodeid]
        assert outcomes
        if outcomes != {'passed'}:
            # Test failed or skipped
            return

        def run_measurements():
            leaks = []
            for checker, after in self.measure(nodeid):
                assert after is not None
                c = self.counters[nodeid][checker]
                before, _ = c[-1]
                c[-1] = (before, after)
                if checker.has_leak(before, after):
                    leaks.append((checker, before, after))
            return leaks

        t1 = time.time()
        deadline = t1 + self.grace_delay
        leaks = run_measurements()
        if leaks:
            self.cleanup()
            for c, _, _ in leaks:
                c.on_retry()
            leaks = run_measurements()

        while leaks and time.time() < deadline:
            time.sleep(0.1)
            self.cleanup()
            for c, _, _ in leaks:
                c.on_retry()
            leaks = run_measurements()

        if leaks:
            self.leaks[nodeid] = leaks
        else:
            self.leaks.pop(nodeid, None)

        for checker in self.checks_for_item(nodeid):
            checker.on_stop_test()

    def maybe_retry(self, item, nextitem=None):
        def run_test_again():
            # This invokes our setup/teardown hooks again
            # Inspired by https://pypi.python.org/pypi/pytest-rerunfailures
            from _pytest.runner import runtestprotocol
            item._initrequest()  # Re-init fixtures
            reports = runtestprotocol(item, nextitem=nextitem, log=False)

        nodeid = item.nodeid
        leaks = self.leaks.get(nodeid)
        if leaks:
            self._retrying = True
            try:
                for i in range(self.max_retries):
                    run_test_again()
            except Exception as e:
                print("--- Exception when re-running test ---")
                import traceback
                traceback.print_exc()
            else:
                leaks = self.leaks.get(nodeid)
            finally:
                self._retrying = False

        return leaks

    # Note on hook execution order:
    #   pytest_runtest_protocol
    #       pytest_runtest_setup
    #       pytest_report_teststatus
    #       pytest_runtest_call
    #       pytest_report_teststatus
    #       pytest_runtest_teardown
    #       pytest_report_teststatus

    # See also https://github.com/abalkin/pytest-leaks/blob/master/pytest_leaks.py

    @pytest.hookimpl(hookwrapper=True)
    def pytest_runtest_protocol(self, item, nextitem):
        if not self._retrying:
            nodeid = item.nodeid
            assert nodeid not in self.counters
            self.counters[nodeid] = {c: [] for c in self.checkers}

            leaking = item.get_marker('leaking')
            if leaking is not None:
                unknown = sorted(set(leaking.args) - set(all_checkers))
                if unknown:
                    raise ValueError("pytest.mark.leaking: unknown resources %r"
                                     % (unknown,))
                classes = tuple(all_checkers[a] for a in leaking.args)
                self.skip_checkers[nodeid] = {c for c in self.checkers
                                              if isinstance(c, classes)}

        yield

    @pytest.hookimpl(hookwrapper=True)
    def pytest_runtest_setup(self, item):
        self.measure_before_test(item.nodeid)
        yield

    @pytest.hookimpl(hookwrapper=True, trylast=True)
    def pytest_runtest_teardown(self, item):
        yield
        self.measure_after_test(item.nodeid)
        if not self._retrying:
            leaks = self.maybe_retry(item)
            if leaks and self.mark_failed:
                # Trigger fail here to allow stopping with `-x`
                pytest.fail()

    @pytest.hookimpl(hookwrapper=True, trylast=True)
    def pytest_report_teststatus(self, report):
        nodeid = report.nodeid
        outcomes = self.outcomes[nodeid]
        outcomes.add(report.outcome)
        outcome = yield
        if not self._retrying:
            if report.when == 'teardown':
                leaks = self.leaks.get(report.nodeid)
                if leaks:
                    if self.mark_failed:
                        outcome.force_result(('failed', 'L', 'LEAKED'))
                        report.outcome = 'failed'
                        report.longrepr = "\n".join(
                            ["%s %s" % (nodeid, checker.format(before, after))
                             for checker, before, after in leaks])
                    else:
                        outcome.force_result(('leaked', 'L', 'LEAKED'))
                # XXX should we log retried tests

    @pytest.hookimpl
    def pytest_terminal_summary(self, terminalreporter, exitstatus):
        tr = terminalreporter
        leaked = tr.getreports('leaked')
        if leaked:
            # If mark_failed is False, leaks are output as a separate
            # results section
            tr.write_sep("=", 'RESOURCE LEAKS')
            for rep in leaked:
                nodeid = rep.nodeid
                for checker, before, after in self.leaks[nodeid]:
                    tr.line("%s %s"
                            % (rep.nodeid, checker.format(before, after)))