import os
import sys
from collections import defaultdict
from os import listdir
from os.path import (
    abspath,
    basename,
    dirname,
    isdir,
    join,
    split,
)

__version__ = '0.10.0'

from hammett.colors import (
    YELLOW,
    RED,
    RESET_COLOR,
    GREEN,
)

MISSING = object()

_orig_print = print


# pytest compatibility. I throw kwargs on the ground on purpose
def param(*args, **kwargs):
    return args


def parametrize(argnames, argvalues, indirect=False, ids=None, scope=None):
    def parametrize_wrapper(f):
        try:
            f.hammett_parametrize_stack.append((argnames, argvalues))
        except AttributeError:
            f.hammett_parametrize_stack = []
            f.hammett_parametrize_stack.append((argnames, argvalues))
        return f

    return parametrize_wrapper


class Item:
    pass


class _Mark:
    def __getattr__(self, name: str):
        if name == parametrize.__name__:
            return parametrize

        else:
            def marker(*args, **kwargs):
                args = list(args)

                def decorate(f):
                    try:
                        f.hammett_markers.append(Marker(name, args, kwargs))
                    except AttributeError:
                        f.hammett_markers = []
                        f.hammett_markers.append(Marker(name, args, kwargs))
                    return f

                if len(args) == 1:
                    if callable(args[0]):
                        f = args[0]
                        args[:] = args[1:]
                        return decorate(f)
                return decorate

            return marker


mark = _Mark()


class Marker:
    def __init__(self, name, args, kwargs):
        self.args = args
        self.kwargs = kwargs
        self.name = name

    def __repr__(self):
        return f'hammett.{self.name}'


class Globals:
    def __init__(self):
        self.source_location = None
        self.modules = None
        self.orig_cwd = os.getcwd()
        self.verbose = False
        self.durations = False
        self.terminal_width = None
        self.results = None
        self.settings = {}
        self.fail_fast = False
        self.quiet = False
        self.drop_into_debugger = False
        self.durations_results = []
        self.disable_assert_analyze = False
        self.output = []
        self.result_db = None
        self.should_stop = False
        self.hijacked_stdout = None
        self.hijacked_stderr = None
        self.use_cache = None
        self.pre_test_callback = None
        self.post_test_callback = None

    def reset(self):
        self.__init__()


g = Globals()


def print(arg='', end='\n', flush=False):
    g.output.append((arg, end, flush))
    if g.quiet:
        return
    _orig_print(arg, end=end, flush=flush, file=sys.__stdout__)


def fixture(*args, **kwargs):
    args = list(args)

    def decorator(f):
        from hammett.impl import register_fixture
        register_fixture(f, *args, **kwargs)
        return f

    if len(args) == 1:
        if callable(args[0]):
            f = args[0]
            args[:] = args[1:]
            return decorator(f)
        else:
            return decorator
    else:
        assert not args
        return decorator


# This is a deprecated feature in pytest, but we'll keep it here for now. pytest-django uses it for example.
def yield_fixture(*args, **kwargs):
    return fixture(*args, **kwargs)


def raises(expected_exception, match=None):
    from hammett.impl import RaisesContext
    return RaisesContext(expected_exception, match=match)


def warns(expected_warning=None, match=None):
    from hammett.impl import WarnsContext
    return WarnsContext(expected_warning, match=match)


def fail(message):
    raise RuntimeError(message)


class Option:
    def __getattr__(self, name):
        if hasattr(g, name):
            return getattr(g, name)
        return False


class Config:
    def __init__(self):
        self.option = Option()

    def getvalue(self, key):
        return g.settings['config'].get(key)

    def getini(self, name):
        return None

    # For xdist emulation add workerinput which is a dict that contains a key "workerid"


class PytestDirShim:
    def __init__(self, dir):
        self.dir = dir

    def dirpath(self):
        return self

    @property
    def purebasename(self):
        return basename(self.dir).rpartition('.')[0]

    def join(self, *args):
        return join(*([dirname(self.dir)] + [str(x) for x in args]))


class Request:
    current_fixture_setup = None

    def __init__(self, scope, parent, function=None):
        self.scope = scope
        self.function = function
        self.parent = parent
        self.additional_fixtures_wanted = set()
        self.keywords = {}
        self.fixturenames = set()
        self.finalizers = []
        self.fixture_results = {}
        self.funcargnames = []
        self.catch_log_handler = None
        if function:
            self.fspath = PytestDirShim(function.__code__.co_filename)
            self.name = function.__name__

        self.config = Config()

    def hammett_add_fixture_result(self, result):
        from hammett.impl import fixture_scope
        if fixture_scope.get(self.current_fixture_setup, 'function') != self.scope:
            self.parent.hammett_add_fixture_result(result)
        else:
            self.fixture_results[self.current_fixture_setup] = result

    def hammett_get_existing_result(self, name):
        try:
            return self.fixture_results[name]
        except KeyError:
            if self.parent:
                return self.parent.hammett_get_existing_result(name)
            return MISSING

    def teardown(self):
        for x in reversed(self.finalizers):
            x()

    @property
    def node(self):
        return self

    def addfinalizer(self, x):
        assert Request.current_fixture_setup is not None
        from hammett.impl import fixture_scope
        if fixture_scope.get(Request.current_fixture_setup, 'function') != self.scope:
            self.parent.addfinalizer(x)
        else:
            self.finalizers.append(x)

    def get_closest_marker(self, s):
        # TODO: fall back to the parent markers
        try:
            markers = self.function.hammett_markers
        except AttributeError:
            return None

        for marker in markers:
            if marker.name == s:
                return marker

        return None

    def getfixturevalue(self, s):
        return self.additional_fixtures_wanted.add(s)


def parse_markers(markers):
    if markers is None:
        return None

    def parse_one(s):
        if '[' not in s:
            return s, None

        assert s.endswith(']')
        base_name, _, name_plus_end_bracket = s.partition('[')
        return base_name, name_plus_end_bracket[:-1]

    # So far let's do something stupidly simple
    markers = markers.split(';')
    return dict(parse_one(m.strip()) for m in markers)


def guess_modules_and_source_path():
    if isdir('lib'):
        return listdir('lib'), 'lib'
    elif isdir('src'):
        return listdir('src'), 'src'

    this_dir = os.getcwd().split(os.sep)[-1]

    if this_dir.startswith('py-'):
        this_dir = this_dir[len('py-'):]

    if isdir(this_dir):
        return [this_dir], '.'
    elif isdir(this_dir.replace('-', '_')):
        return [this_dir.replace('-', '_')], '.'
    elif isdir(this_dir.replace(' ', '_')):
        return [this_dir.replace(' ', '_')], '.'
    elif isdir(this_dir.replace('-', '')):
        return [this_dir.replace('-', '')], '.'
    elif isdir(this_dir.replace(' ', '')):
        return [this_dir.replace(' ', '')], '.'
    return [], '.'


def is_test_file(f):
    return not f.startswith('.') and f.endswith('.py') and (f.endswith('__tests.py') or f.startswith('test_') or f.endswith('_test.py'))


def handle_dir(result, conftest_files, d):
    for root, dirs, files in os.walk(d):
        for f in files:
            if is_test_file(f):
                result.append(join(root, f) )
            if f == 'conftest.py':
                conftest_files.append(join(root, f))


def collect_files(filenames):
    result = []
    conftest_files = []
    if filenames is None:
        handle_dir(result, conftest_files, f'tests{os.sep}')
        handle_dir(result, conftest_files, f'test{os.sep}')
        handle_dir(result, conftest_files, f'docs{os.sep}')
        for module_name in g.modules:
            handle_dir(result, conftest_files, join(g.source_location, module_name))
    else:
        for filename in filenames:
            if os.path.exists(filename):
                if os.path.isdir(filename):
                    handle_dir(result, conftest_files, filename)
                else:
                    result.append(filename)
            else:
                # This is a symbol, module or function/class? Try module first.
                symbol_path = join(g.source_location, filename.replace('.', os.sep))
                if os.path.exists(symbol_path + '.py'):
                    result.append(symbol_path + '__tests.py')

    return result, conftest_files


def collect_file_data(path):
    data = {}
    for root, dirs, files in os.walk(path):
        dirs[:] = [x for x in dirs if not x.startswith('.') and x not in ['venv', 'env', '__pycache__']]
        for filename in files:
            if not filename.endswith('.py'):
                continue
            full_path = join(root, filename)
            if full_path.startswith(f'.{os.sep}'):
                full_path = full_path[2:]
            data[full_path] = os.stat(full_path).st_mtime_ns
    return data


DB_VERSION = 2
DB_FILENAME = '.hammett-db'


def write_result_db(results):
    """
    The results database is a simple pickled dict with some keys:

    db_version: the version. So we can change the format and throw away an old db if needed
    test_results: dict from filename -> (dict from test_name -> dict(stdout=str, stderr=str, status=str))
    file_data: dict filename -> nanosecond modification date
    """
    if not g.use_cache:
        return

    results['db_version'] = DB_VERSION
    from pickle import dump
    with open(DB_FILENAME, 'wb') as f:
        dump(results, f)


def new_result_db():
    return dict(
        db_version=DB_VERSION,
        test_results=defaultdict(dict),
        file_data=None,
    )


def read_result_db():
    if not g.use_cache:
        return new_result_db()
    from pickle import load
    try:
        with open(DB_FILENAME, 'rb') as f:
            results = load(f)
            if results['db_version'] != DB_VERSION:
                raise FileNotFoundError()
    except FileNotFoundError:
        return new_result_db()
    return results


def drop_cache_for_filename(result_db, filename):
    try:
        del result_db['test_results'][filename]
    except KeyError:
        pass


def update_result_db(result_db, new_file_data):
    if not g.use_cache:
        return

    if result_db['file_data'] is None:
        result_db['file_data'] = new_file_data
        assert not result_db['test_results']
        return result_db

    # Clear out test results when the test file or the tested module has changed
    old_file_data = result_db['file_data']
    clear_all_non_module_tests = False
    for filename, modification_time in old_file_data.items():
        if filename in new_file_data and modification_time != new_file_data[filename]:
            if filename.endswith('__tests.py') or filename.startswith('test_'):
                # The test has been changed
                drop_cache_for_filename(result_db, filename)
            else:
                # The module has been changed so translate the filename to the possible test files
                drop_cache_for_filename(result_db, filename[:-(len('.py'))] + '__tests.py')
                _, filename_only = split(filename)
                drop_cache_for_filename(result_db, f"tests{os.sep}{filename[:-(len('.py'))]}" + '__tests.py')
                clear_all_non_module_tests = True
                # TODO: this doesn't clear the db for module__function__tests.py

    if clear_all_non_module_tests:
        from pathlib import Path
        non_module_filenames = [x for x in result_db['test_results'].keys() if Path(x).stem.startswith('test_')]
        for x in non_module_filenames:
            del result_db['test_results'][x]

    result_db['file_data'] = new_file_data


def finish():
    for x in g.result_db['test_results'].values():
        for y in x.values():
            g.results[y.status] += 1


def main(verbose=False, fail_fast=False, quiet=False, filenames=None, drop_into_debugger=False, match=None, durations=False, markers=None, disable_assert_analyze=False, module_unload=False, cwd=None, use_cache=False, pre_test_callback=None, post_test_callback=None, capture_output=True, insert_cwd=True):
    params = main_setup(verbose=verbose, fail_fast=fail_fast, quiet=quiet, filenames=filenames, drop_into_debugger=drop_into_debugger, durations=durations, markers=markers, disable_assert_analyze=disable_assert_analyze, cwd=cwd, use_cache=use_cache, pre_test_callback=pre_test_callback, post_test_callback=post_test_callback, capture_output=capture_output, insert_cwd=insert_cwd)
    return main_run_tests(match=match, module_unload=module_unload, **params)


def main_setup(verbose=False, fail_fast=False, quiet=False, filenames=None, drop_into_debugger=False, durations=False, markers=None, disable_assert_analyze=False, cwd=None, use_cache=False, pre_test_callback=None, post_test_callback=None, capture_output=True, insert_cwd=True):
    import sys
    if sys.version_info[:2] < (3, 6):
        print('hammett requires python 3.6 or later')
        exit(999)

    sys.modules['pytest'] = sys.modules['hammett']

    markers = parse_markers(markers)

    clean_up_sys_path = False

    if cwd is None:
        cwd = os.getcwd()
    cwd = abspath(cwd)
    g.orig_cwd = cwd
    os.chdir(cwd)

    if cwd not in sys.path and insert_cwd:
        sys.path.insert(0, cwd)
        clean_up_sys_path = True

    g.reset()

    g.results = dict(success=0, failed=0, skipped=0, abort=0)
    g.verbose = verbose
    g.fail_fast = fail_fast
    g.quiet = quiet
    g.drop_into_debugger = drop_into_debugger
    g.durations = durations
    g.disable_assert_analyze = disable_assert_analyze
    g.use_cache = use_cache
    g.pre_test_callback = pre_test_callback
    g.post_test_callback = post_test_callback
    g.capture_output = capture_output

    from hammett.impl import read_settings
    read_settings()

    g.source_location = g.settings.get('source_location', '.')
    g.modules = g.settings.get('modules')
    if isinstance(g.modules, str):
        g.modules = [g.modules]
    if g.modules is None:
        m, s = guess_modules_and_source_path()
        g.source_location = g.source_location or s
        g.modules = m

    g.result_db = read_result_db()
    update_result_db(g.result_db, collect_file_data(g.source_location))
    write_result_db(g.result_db)

    filenames, conftest_files = collect_files(filenames)

    return dict(
        filenames=filenames,
        conftest_files=conftest_files,
        markers=markers,
        clean_up_sys_path=clean_up_sys_path,
    )


def main_run_tests(filenames, conftest_files, markers, clean_up_sys_path, match=None, module_unload=False, tests=None):
    if not filenames:
        print('No module to test found.')
        print('''You might need to add `modules` and `source_location` item under [hammett] in setup.cfg  like:

[hammett]
modules=
    foo
    bar
source_location=.
''')
        return 3

    g.results.update(dict(success=0, failed=0, skipped=0, abort=0))

    plugins_loaded = False
    conftests_loaded = False
    from os.path import split, sep

    session_request = Request(scope='session', parent=None)

    from hammett.impl import should_stop

    for test_filename in sorted(filenames):
        dirname, filename = split(test_filename)

        if dirname.startswith(f'.{os.sep}'):
            dirname = dirname[2:]

        module_name = f'{dirname.replace(sep, ".")}.{filename.replace(".py", "")}'
        # if module_name in sys.modules:
        #     del sys.modules[module_name]

        cache_filename = join(dirname, filename)
        if cache_filename in g.result_db['test_results']:
            from hammett.impl import print_status
            for test_name, result in g.result_db['test_results'][cache_filename].items():
                print_status(test_name, result)
            continue

        # We do this here because if all test results are up to date, we want to avoid loading slow plugins!
        if not plugins_loaded:
            from hammett.impl import load_plugins
            load_plugins()
            plugins_loaded = True

        if should_stop():
            break

        if not conftests_loaded:
            from hammett.impl import load_conftests
            load_conftests(conftest_files)
            conftests_loaded = True

        run_tests_for_filename(test_filename, session_request, markers, match, module_name, tests)

        # if module_unload:
        #     del sys.modules[module_name]

        if should_stop():
            break

    session_request.teardown()

    if not g.verbose:
        print()

    finish()

    color = GREEN
    if g.results['skipped']:
        color = YELLOW
    if g.results['failed']:
        color = RED

    if clean_up_sys_path:
        del sys.path[0]

    if g.durations:
        limit = 10
        print()
        print(f'--- {limit} slowest tests ---')
        g.durations_results.sort(key=lambda x: x[1], reverse=True)
        for name, time, _ in g.durations_results[:10]:
            print(f'{time}   {name}')

        print()
        print(f'--- {limit} slowest startup times for tests ---')
        print('Note: startup times might be shared cost so might show up unfairly')
        g.durations_results.sort(key=lambda x: x[2], reverse=True)
        for name, _, startup_time in g.durations_results[:10]:
            print(f'{startup_time}   {name}')

        print()

    print(f'{color}{g.results["success"]} succeeded, {g.results["failed"]} failed, {g.results["skipped"]} skipped{RESET_COLOR}')
    os.chdir(g.orig_cwd)

    if g.results['abort']:
        write_result_db(g.result_db)
        return 2

    write_result_db(g.result_db)
    return 1 if g.results['failed'] else 0


def load_module(module_name, test_filename):
    import importlib.util
    spec = importlib.util.spec_from_file_location(module_name, test_filename)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    try:
        spec.loader.exec_module(module)
    except Exception:
        print(f'Failed to load module {module_name}:')
        import traceback
        print(traceback.format_exc())
        g.results['abort'] += 1
    return module


def run_tests_for_filename(test_filename, session_request, markers, match, module_name, tests):
    module = load_module(module_name, test_filename)

    if g.results['abort']:
        return

    module_request = Request(scope='module', parent=session_request)

    module_markers = getattr(module, 'pytestmark', [])
    if not isinstance(module_markers, list):
        module_markers = [module_markers]

    from hammett.impl import execute_test_function, execute_test_class
    from unittest import TestCase
    from hammett.impl import should_stop
    for name, f in list(module.__dict__.items()):
        if g.results['abort']:
            break

        full_name = f'{test_filename}::{name}'.replace('./', '')

        is_test_function = name.startswith('test_') and callable(f)
        is_test_class = isinstance(f, type) and issubclass(f, TestCase) or name.startswith('Test')

        if not is_test_function and not is_test_class:
            continue

        if match is not None:
            if match not in name:
                continue

        if tests is not None:
            if full_name not in tests:
                continue

        for m in module_markers:
            f = m(f)

        if markers is not None:
            keep = False
            for f_marker in getattr(f, 'hammett_markers', []):
                if f_marker.name in markers:
                    arg = markers[f_marker.name]
                    if arg is not None:
                        assert len(f_marker.args) == 1, 'hammett only supports filtering on single arguments to markers right now'
                        if str(f_marker.args[0]) == arg:
                            keep = True
                            break
                    else:
                        keep = True
                        break
            if not keep:
                continue

        if is_test_function:
            execute_test_function(full_name, f, module_request)
        elif is_test_class:
            execute_test_class(full_name, f, module_request)

        if should_stop():
            break

    module_request.teardown()


def hookimpl(*_, **__):
    return lambda f: f


def multi_process_main(*, match, module_unload=False, **kwargs):
    import os

    # key == filename

    # manual fork
    result_by_key = {}

    def read_one_child_exit_status():
        pid, status = os.wait()
        result_by_key[key_from_pid[pid]] = (0xFF00 & status) >> 8  # The high byte contains the exit code

    hammett_kwargs = main_setup(**kwargs)

    key_from_pid = {}
    running_children = 0
    max_children = 160

    for key in hammett_kwargs['filenames']:
        pid = os.fork()
        if not pid:
            # In the child

            # this is needed for non-memory DBs
            Config.workerinput = dict(workerinput=f'_{key}')

            hammett_kwargs['filenames'] = [key]

            _orig_print('start:', key)

            result = main_run_tests(**hammett_kwargs, match=match, module_unload=module_unload)
            _orig_print('end:', key)
            if result != 0:
                # TODO: write failure information to stdout?
                pass
            os._exit(result)
        else:
            key_from_pid[pid] = key
            running_children += 1

        if running_children >= max_children:
            read_one_child_exit_status()
            running_children -= 1

    try:
        while True:
            read_one_child_exit_status()
    except ChildProcessError:
        pass
    return 0


def main_cli(args=None):
    if args is None:
        args = sys.argv[1:]
    from argparse import ArgumentParser
    parser = ArgumentParser(prog='hammett')
    parser.add_argument('-v', '--verbose', dest='verbose', action='store_true', default=False)
    parser.add_argument('-x', dest='fail_fast', action='store_true', default=False)
    parser.add_argument('-q', dest='quiet', action='store_true', default=False)
    parser.add_argument('--multi-experimental', dest='multi_process', action='store_true', default=False)
    parser.add_argument('-k', dest='match', default=None)
    parser.add_argument('-m', dest='markers', default=None)
    parser.add_argument('-s', dest='capture_output', default=None)
    parser.add_argument('--use-cache', dest='use_cache', default=False, help='The cache is an experimental feature to run only relevant changes based on looking at what files have been changed.')
    parser.add_argument('--durations', dest='durations', action='store_true', default=False)
    parser.add_argument('--no-assert-analyze', dest='disable_assert_analyze', action='store_true', default=False)
    parser.add_argument('--pdb', dest='drop_into_debugger', action='store_true', default=False)
    parser.add_argument(dest='filenames', nargs='*')
    args = parser.parse_args(args)

    m = main
    if args.multi_process:
        m = multi_process_main

    return m(
        verbose=args.verbose,
        fail_fast=args.fail_fast,
        quiet=args.quiet,
        filenames=args.filenames or None,
        drop_into_debugger=args.drop_into_debugger,
        match=args.match,
        durations=args.durations,
        markers=args.markers,
        disable_assert_analyze=args.disable_assert_analyze,
        use_cache=args.use_cache,
        capture_output=args.capture_output,
    )


if __name__ == '__main__':
    exit(main_cli())
