import multiprocessing
import time

import grpc
import pytest
import requests

from cmon import __version__
from cmon.constants import __cmon_env__
from cmon.proto import cmon_pb2, cmon_pb2_grpc
from cmon.serve.runtimes.asyncio import AsyncNewLoopRuntime
from cmon.serve.runtimes.servers import BaseServer
from cmon.serve.runtimes.worker.request_handling import WorkerRequestHandler
from tests.helper import _generate_pod_args

from .test_runtimes import _create_gateway_runtime, _create_head_runtime


def _create_worker_runtime(port, name='', executor=None):
    args = _generate_pod_args()
    args.port = [port]
    args.name = name
    if executor:
        args.uses = executor
    with AsyncNewLoopRuntime(args, req_handler_cls=WorkerRequestHandler) as runtime:
        runtime.run_forever()


def _create_worker(port):
    # create a single worker runtime
    p = multiprocessing.Process(target=_create_worker_runtime, args=(port,))
    p.start()
    time.sleep(0.1)
    return p


def _create_gateway(port, graph, pod_addr, protocol):
    # create a single worker runtime
    # create a single gateway runtime
    p = multiprocessing.Process(
        target=_create_gateway_runtime,
        args=(graph, pod_addr, port, protocol),
    )
    p.start()
    time.sleep(0.1)
    return p


def _create_head(port, connection_list_dict, polling='ANY'):
    p = multiprocessing.Process(
        target=_create_head_runtime, args=(port, connection_list_dict, 'head', polling)
    )
    p.start()
    time.sleep(0.1)
    return p


@pytest.mark.parametrize('runtime', ['head', 'worker', 'gateway'])
def test_cmon_info_grpc_based_runtimes(runtime, port_generator):
    port = port_generator()
    connection_list_dict = {}
    graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
    pod_addresses = f'{{"pod0": ["0.0.0.0:{port}"]}}'
    if runtime == 'head':
        p = _create_head(port, connection_list_dict)
    elif runtime == 'gateway':
        p = _create_gateway(port, graph_description, pod_addresses, 'grpc')
    else:
        p = _create_worker(port)

    try:
        BaseServer.wait_for_ready_or_shutdown(
            timeout=5.0,
            ctrl_address=f'0.0.0.0:{port}',
            ready_or_shutdown_event=multiprocessing.Event(),
        )

        channel = grpc.insecure_channel(f'localhost:{port}')
        stub = cmon_pb2_grpc.CmonInfoRPCStub(channel)
        res = stub._status(
            cmon_pb2.google_dot_protobuf_dot_empty__pb2.Empty(),
        )
        assert res.cmon['cmon'] == __version__
        for env_var in __cmon_env__:
            assert env_var in res.envs
    except Exception:
        assert False
    finally:
        p.terminate()
        p.join()


@pytest.mark.parametrize('protocol', ['http', 'websocket'])
def test_cmon_info_gateway_http(protocol, port_generator):
    port = port_generator()
    graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
    pod_addresses = f'{{"pod0": ["0.0.0.0:{port}"]}}'
    p = _create_gateway(port, graph_description, pod_addresses, protocol)

    try:
        BaseServer.wait_for_ready_or_shutdown(
            timeout=5.0,
            ctrl_address=f'0.0.0.0:{port}',
            ready_or_shutdown_event=multiprocessing.Event(),
        )

        x = requests.get(f'http://localhost:{port}/status')
        resp = x.json()
        assert 'cmon' in resp
        assert 'envs' in resp
        assert resp['cmon']['cmon'] == __version__
        for env_var in __cmon_env__:
            assert env_var in resp['envs']
    except Exception:
        assert False
    finally:
        p.terminate()
        p.join()
