import functools
import collections
import subprocess
import multiprocessing
import multiprocessing.connection
import tempfile
import shutil
import os
import sys
import inspect
import importlib
import concurrent.futures
import contextlib
import time
import random
import string

from tblib import pickling_support
pickling_support.install()

import forge
import anyio
from async_generator import aclosing


@contextlib.contextmanager
def compile_temp_proto(*relative_proto_paths):
    modules = []
    with tempfile.TemporaryDirectory() as temp_dir:
        sys.path.insert(0, temp_dir)
        try:
            for relative_proto_path in relative_proto_paths:
                proto_path = os.path.join(os.path.dirname(
                    inspect.currentframe().f_back.f_back.f_globals['__file__']),
                    relative_proto_path)
                proto_filename = os.path.basename(proto_path)
                proto_temp_path = os.path.join(temp_dir, proto_filename)
                shutil.copyfile(proto_path, proto_temp_path)
            for relative_proto_path in relative_proto_paths:
                proto_filename = os.path.basename(relative_proto_path)
                proto_temp_path = os.path.join(temp_dir, proto_filename)
                cmdline = [sys.executable, '-m', 'grpc_tools.protoc',
                           '--python_out=.', '--purerpc_out=.', '--grpc_python_out=.',
                           '-I' + temp_dir, proto_temp_path]
                subprocess.check_call(cmdline, cwd=temp_dir)

                pb2_module_name = proto_filename.replace(".proto", "_pb2")
                pb2_grpc_module_name = proto_filename.replace(".proto", "_pb2_grpc")
                grpc_module_name = proto_filename.replace(".proto", "_grpc")

                pb2_module = importlib.import_module(pb2_module_name)
                pb2_grpc_module = importlib.import_module(pb2_grpc_module_name)
                grpc_module = importlib.import_module(grpc_module_name)
                modules.extend((pb2_module, pb2_grpc_module, grpc_module))
            yield modules
        finally:
            sys.path.remove(temp_dir)


_WrappedResult = collections.namedtuple("_WrappedResult", ("result", "exc_info"))


def _wrap_gen_in_process(conn: multiprocessing.connection.Connection):
    def decorator(gen):
        @functools.wraps(gen)
        def new_func(*args, **kwargs):
            try:
                for elem in gen(*args, **kwargs):
                    conn.send(_WrappedResult(result=elem, exc_info=None))
            except:
                conn.send(_WrappedResult(result=None, exc_info=sys.exc_info()))
            finally:
                conn.close()
        return new_func
    return decorator


async def async_iterable_to_list(async_iterable):
    result = []
    async with aclosing(async_iterable) as async_iterable:
        async for value in async_iterable:
            result.append(value)
    return result


def random_payload(min_size=1000, max_size=100000):
    return "".join(random.choice(string.ascii_letters)
                   for _ in range(random.randint(min_size, max_size)))


@contextlib.contextmanager
def _run_context_manager_generator_in_process(cm_gen):
    parent_conn, child_conn = multiprocessing.Pipe(duplex=False)
    target_fn = _wrap_gen_in_process(child_conn)(cm_gen)

    process = multiprocessing.Process(target=target_fn)
    process.start()
    try:
        wrapped_result = parent_conn.recv()
        if wrapped_result.exc_info is not None:
            raise wrapped_result.exc_info[0].with_traceback(*wrapped_result.exc_info[1:])
        else:
            yield wrapped_result.result
    finally:
        try:
            if parent_conn.poll():
                exc_info = parent_conn.recv().exc_info
                if exc_info is not None:
                    raise exc_info[0].with_traceback(*exc_info[1:])
        finally:
            process.terminate()
            process.join()
            parent_conn.close()


def run_purerpc_service_in_process(service):
    def target_fn():
        import purerpc
        server = purerpc.Server(port=0)
        server.add_service(service)
        socket = server._create_socket_and_listen()
        yield socket.getsockname()[1]

        async def sleep_10_seconds_then_die():
            await anyio.sleep(20)
            raise ValueError

        async def main():
            async with anyio.create_task_group() as tg:
                await tg.spawn(server._run_async_server, socket)
                await tg.spawn(sleep_10_seconds_then_die)
        # import cProfile
        anyio.run(server._run_async_server, socket)
        # cProfile.runctx("anyio.run(main)", globals(), locals(), sort="tottime")
    return _run_context_manager_generator_in_process(target_fn)


def run_grpc_service_in_process(add_handler_fn):
    def target_fn():
        import grpc
        server = grpc.server(concurrent.futures.ThreadPoolExecutor(max_workers=1))
        port = server.add_insecure_port('[::]:0')
        add_handler_fn(server)
        server.start()
        yield port
        while True:
            time.sleep(60)
    return _run_context_manager_generator_in_process(target_fn)


def run_tests_in_workers(*, target, num_workers):
    parent_conn, child_conn = multiprocessing.Pipe(duplex=False)

    @_wrap_gen_in_process(child_conn)
    def target_fn():
        target()
        yield

    processes = [multiprocessing.Process(target=target_fn) for _ in range(num_workers)]
    for process in processes:
        process.start()

    try:
        for _ in range(num_workers):
            wrapped_result = parent_conn.recv()
            if wrapped_result.exc_info is not None:
                raise wrapped_result.exc_info[0].with_traceback(*wrapped_result.exc_info[1:])
    finally:
        parent_conn.close()
        for process in processes:
            process.join()


def async_test(corofunc):
    if not inspect.iscoroutinefunction(corofunc):
        raise TypeError("Expected coroutine function")

    @functools.wraps(corofunc)
    def func(**kwargs):
        return anyio.run(functools.partial(corofunc, **kwargs))
    return func


def grpc_client_parallelize(num_workers):
    def decorator(func):
        @functools.wraps(func)
        def new_func(*args, **kwargs):
            def target():
                func(*args, **kwargs)
            run_tests_in_workers(target=target, num_workers=num_workers)

        new_func.__parallelized__ = True
        return new_func
    return decorator


def purerpc_client_parallelize(num_tasks):
    def decorator(corofunc):
        if not inspect.iscoroutinefunction(corofunc):
            raise TypeError("Expected coroutine function")

        @functools.wraps(corofunc)
        async def new_corofunc(**kwargs):
            async with anyio.create_task_group() as tg:
                for _ in range(num_tasks):
                    await tg.spawn(functools.partial(corofunc, **kwargs))
        return new_corofunc
    return decorator


def grpc_channel(port_fixture_name, channel_arg_name="channel"):
    def decorator(func):
        if hasattr(func, "__parallelized__") and func.__parallelized__:
            raise TypeError("Cannot pass gRPC channel to already parallelized test, grpc_client_parallelize should "
                            "be the last decorator in chain")

        @forge.compose(
            forge.copy(func),
            forge.modify(channel_arg_name, name=port_fixture_name, interface_name="port_fixture_value"),
        )
        def new_func(*, port_fixture_value, **kwargs):
            import grpc
            with grpc.insecure_channel('127.0.0.1:{}'.format(port_fixture_value)) as channel:
                func(**kwargs, channel=channel)

        return new_func
    return decorator


def purerpc_channel(port_fixture_name, channel_arg_name="channel"):
    def decorator(corofunc):
        if not inspect.iscoroutinefunction(corofunc):
            raise TypeError("Expected coroutine function")

        @forge.compose(
            forge.copy(corofunc),
            forge.modify(channel_arg_name, name=port_fixture_name, interface_name="port_fixture_value"),
        )
        async def new_corofunc(*, port_fixture_value, **kwargs):
            import purerpc
            async with purerpc.insecure_channel("127.0.0.1", port_fixture_value) as channel:
                await corofunc(**kwargs, channel=channel)

        return new_corofunc
    return decorator
