import pytest
from _pytest.pathlib import Path
import typing as T
import json
import sys
from .__version__ import __version__
import tqdm
import math


PYTEST_ARGS = ('-qq', '-p', 'pytest_cleanslate.reduce')
MODULE_LIST_ARG = '--module-list-from'
TEST_LIST_ARG = '--test-list-from'
RESULTS_ARG = '--results-to'


class ReducePlugin:
    def __init__(self, config: pytest.Config) -> None:
        if (modules_file := config.getoption(MODULE_LIST_ARG)):
            with modules_file.open('r') as f:
                self._modules = {Path(line.strip()).resolve() for line in f}
        else:
            self._modules = None

        if (tests_file := config.getoption(TEST_LIST_ARG)):
            with tests_file.open('r') as f:
                self._tests = set(line.strip() for line in f)
        else:
            self._tests = None

        self._results_file = config.getoption(RESULTS_ARG)
        self._collect = []
        self._run = []


    @pytest.hookimpl
    def pytest_ignore_collect(self, collection_path: Path, config: pytest.Config) -> T.Union[None, bool]:
        if self._modules is not None and collection_path.suffix == '.py':
            return collection_path.resolve() not in self._modules


    @pytest.hookimpl
    def pytest_collectreport(self, report: pytest.CollectReport) -> None:
        if self._results_file and report.nodeid.endswith('.py'):
            self._collect.append({
                'id': report.nodeid,
                'outcome': report.outcome,
                'result': [n.nodeid for n in report.result]
            })


    @pytest.hookimpl(tryfirst=True)
    def pytest_collection_modifyitems(self, items: T.List[pytest.Item], config: pytest.Config) -> T.Union[None, bool]:
        if self._tests is not None:
            selected = []
            deselected = []

            for item in items:
                if item.nodeid in self._tests:
                    selected.append(item)
                else:
                    deselected.append(item)

            if deselected:
                config.hook.pytest_deselected(items=deselected)
                items[:] = selected


    @pytest.hookimpl
    def pytest_runtest_logreport(self, report: pytest.TestReport) -> None:
        # to cut down on the log, we only save non-call run reports for failures
        if self._results_file and (report.outcome != 'passed' or report.when == 'call'):
            self._run.append({
                'id': report.nodeid,
                'outcome': report.outcome,
            })


    def write_results(self) -> None:
        if self._results_file:
            with self._results_file.open("w") as f:
                json.dump({
                    'collect': self._collect,
                    'run': self._run
                }, f)


@pytest.hookimpl
def pytest_addoption(parser: pytest.Parser, pluginmanager: pytest.PytestPluginManager) -> None:
    parser.addoption(MODULE_LIST_ARG, type=Path, help="Only collect modules in the given file")
    parser.addoption(TEST_LIST_ARG, type=Path, help="Only run tests whose node IDs are in the given file")
    parser.addoption(RESULTS_ARG, type=Path, help="Write test collection/run results to the given file")


@pytest.hookimpl
def pytest_configure(config: pytest.Config) -> None:
    if config.getoption(MODULE_LIST_ARG) or config.getoption(TEST_LIST_ARG) or \
       config.getoption(RESULTS_ARG):
        config._cleanslate_reduce_plugin = ReducePlugin(config)
        config.pluginmanager.register(config._cleanslate_reduce_plugin)


@pytest.hookimpl
def pytest_unconfigure(config: pytest.Config) -> None:
    if (plugin := getattr(config, "_cleanslate_reduce_plugin", None)):
        plugin.write_results()


class Results:
    """Facilitates access to test results file generated by ReducePlugin."""

    def __init__(self, results_file: Path):
        with results_file.open("r") as f:
            self._results = json.load(f)

        self._outcomes = None

    def get_outcome(self, nodeid: str) -> str:
        if self._outcomes is None:
            self._outcomes = {r['id']: r['outcome'] for r in self._results['collect'] + self._results['run']}

        return self._outcomes[nodeid]

    def get_modules(self) -> T.List[str]:
        return [r['id'] for r in self._results['collect']]

    def get_tests(self) -> T.List[str]:
        return [r['id'] for r in self._results['run']]

    def get_failed(self) -> T.Iterator[str]:
        return iter(o['id'] for o in self._results['collect'] + self._results['run'] if o['outcome'] == 'failed')

    def get_first_failed(self) -> T.Union[None, str]:
        return next(self.get_failed(), None)


def get_module(testid: str) -> str:
    return testid.split('::')[0]


def _is_module(testid: str) -> bool:
    return '::' not in testid


def get_function(testid: str) -> str:
    if '::' in testid:
        return testid.split('::')[1].split('[')[0]


def run_pytest(tests_path: Path, pytest_args=(), *,
               modules: T.List[Path] = None, tests: T.List[str] = None, trace: bool = False) -> dict:
    import tempfile
    import subprocess

    # Specifying dir='.' here works around weird failures running tests on MacOS
    with tempfile.TemporaryDirectory(dir='.') as tmpdir:
        tmpdir = Path(tmpdir)

        results = tmpdir / "results.json"

        if modules:
            modulelist = tmpdir / "modules.txt"
            modulelist.write_text('\n'.join(modules))

        if tests:
            testlist = tmpdir / "tests.txt"
            testlist.write_text('\n'.join(tests))

        command = [
            sys.executable, '-m', 'pytest', *PYTEST_ARGS, *pytest_args,
            RESULTS_ARG, results,
            *((MODULE_LIST_ARG, modulelist) if modules else ()),
            *((TEST_LIST_ARG, testlist) if tests else ()),
            tests_path
        ]

        if trace:
            print(f"Running {command}", flush=True)

        p = subprocess.run(command,
                           check=False, **({} if trace else {'stdout': subprocess.DEVNULL}))
        if p.returncode not in (pytest.ExitCode.OK, pytest.ExitCode.TESTS_FAILED,
                                pytest.ExitCode.INTERRUPTED, pytest.ExitCode.NO_TESTS_COLLECTED):
            p.check_returncode()

        return Results(results)


def _bisect_items(items: T.List[str], failing: str, fails: T.Callable[[T.List[str], str], bool],
                  *, bar: "tqdm") -> T.List[str]:
    assert failing not in items

    while len(items) > 1:
        middle = len(items) // 2

        bar.refresh() # for when using --trace
        bar.set_postfix({"remaining": len(items)})
        bar.update()

        if fails(items[:middle]+[failing]):
            items = items[:middle]
            continue

        if fails(items[middle:]+[failing]):
            items = items[middle:]
            continue

        # TODO could do the rest of delta debugging here
        break

    if len(items) == 1 and fails([failing]):
        items = []

    bar.refresh() # for when using --trace
    bar.set_postfix({"remaining": len(items)})
    bar.update()

    return items


def _reduce_tests(tests_path: Path, tests: T.List[str], failing_test: str, modules: T.List[str],
                  *, trace: bool = False, pytest_args: T.List[str] = ()) -> T.List[str]:
    def fails(test_set: T.List[str]):
        trial = run_pytest(tests_path, (*pytest_args, '--continue-on-collection-errors'),
                           tests=test_set, modules=modules, trace=trace)
        return trial.get_outcome(failing_test) == 'failed'

    module_set = {*modules}
    tests = [t for t in tests if t != failing_test and get_module(t) in module_set]
    if not tests:
        return tests

    steps=math.ceil(math.log(len(tests), 2))
    with tqdm.tqdm(desc="Trying to reduce tests.....", total=steps) as bar:
        return _bisect_items(tests, failing_test, fails, bar=bar)


def _reduce_modules(tests_path: Path, tests: T.List[str], failing_id: str,
                    modules: T.List[str], failing_module: str,
                    *, trace: bool = False, pytest_args: T.List[str] = ()) -> T.List[str]:

    def fails(module_set: T.List[str]):
        trial = run_pytest(tests_path, (*pytest_args, '--continue-on-collection-errors',),
                           tests=tests, modules=module_set, trace=trace)
        return trial.get_outcome(failing_id) == 'failed'

    modules = [m for m in modules if m != failing_module]
    if not modules:
        return modules

    steps = math.ceil(math.log(len(modules), 2))
    with tqdm.tqdm(desc="Trying to reduce modules...", total=steps) as bar:
        return _bisect_items(modules, failing_module, fails, bar=bar)


def reduce(*, tests_path: Path, results: Results = None, pytest_args: T.List[str] = (),
           trace: bool = False, **args) -> dict:
    if not results:
        print("Running tests...", flush=True)
        results = run_pytest(tests_path, (*pytest_args, '-x'), trace=trace)

    failed_id = results.get_first_failed()
    if failed_id is None:
        print("No tests failed!", flush=True)
        return {
            'failed': failed_id,
            'error': 'No tests failed',
        }

    failed_is_module = _is_module(failed_id)
    if failed_is_module:
        if trace: print()
        print(f"Module \"{failed_id}\"'s collection failed; trying it by itself...", flush=True)
        failed_module = failed_id
        tests = None
    else:
        if trace: print()
        print(f"Test \"{failed_id}\" failed; trying it by itself...", flush=True)
        failed_module = get_module(failed_id)
        tests = [failed_id]

    solo = run_pytest(tests_path, pytest_args, modules=[failed_module], tests=tests, trace=trace)
    if solo.get_outcome(failed_id) != 'passed':
        print("That also fails by itself!", flush=True)
        return {
            'failed': failed_id,
            'error': f'{"Module" if failed_is_module else "Test"} also fails by itself',
        }

    tests = results.get_tests()

    if trace: print()
    modules = _reduce_modules(tests_path, tests, failed_id, results.get_modules(), failed_module,
                              trace=trace, pytest_args=pytest_args)

    if not failed_is_module:
        if trace: print()
        tests = _reduce_tests(tests_path, tests, failed_id, [*modules, failed_module],
                              trace=trace, pytest_args=pytest_args)

        # TODO if tests != [], see if it's enough to disable just them

    if trace: print()
    print("Reduced failure set:")
    print(f"    modules: {modules}")
    print(f"    tests: {tests}")
    print(flush=True)

    return {
        'failed': failed_id,
        'modules': modules,
        'tests': tests,
    }


def _parse_args():
    import argparse

    bool_action = argparse.BooleanOptionalAction if sys.version_info[:2] >= (3,9) else "store_true"

    ap = argparse.ArgumentParser()
    ap.add_argument('--trace', default=False, action=bool_action, help='show pytest outputs, etc.')
    ap.add_argument('--save-to', type=Path, help='file where to save results (JSON)')
    ap.add_argument('--pytest-args', type=str, default='', help='extra arguments to pass to pytest')
    ap.add_argument('--version', action='version',
                    version=f"%(prog)s v{__version__} (Python {'.'.join(map(str, sys.version_info[:3]))})")
    ap.add_argument('tests_path', type=Path, help='tests file or directory')

    return ap.parse_args()


def main():
    args = _parse_args()
    args.pytest_args = args.pytest_args.split()

    results = reduce(**vars(args))
    if args.save_to:
        with args.save_to.open("w") as f:
            json.dump(results, f)

    return 1 if 'error' in results else 0


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