import logging
from concurrent import futures

import yandex.cloud.priv.loadtesting.agent.v1.job_service_pb2 as job_public
import yandex.cloud.priv.loadtesting.agent.v1.job_service_pb2_grpc as job_grpc
import yandex.cloud.priv.loadtesting.agent.v1.tank_service_pb2 as tank_public
import yandex.cloud.priv.loadtesting.agent.v1.tank_service_pb2_grpc as tank_grpc
import yandex.cloud.priv.loadtesting.agent.v1.trail_service_pb2 as trail_public
import yandex.cloud.priv.loadtesting.agent.v1.trail_service_pb2_grpc as trail_grpc
# import yandex.cloud.priv.loadtesting.agent.v1.monitoring_service_pb2 as monitoring_public
# import yandex.cloud.priv.loadtesting.agent.v1.monitoring_service_pb2_gr   pc as monitoring_grpc
import load.projects.cloud.tank_client.proto.monitoring_service_pb2 as monitoring_public
import load.projects.cloud.tank_client.proto.monitoring_service_pb2_grpc as monitoring_grpc
import yandex.cloud.priv.loadtesting.agent.v1.test_service_pb2 as test_public
import yandex.cloud.priv.loadtesting.agent.v1.test_service_pb2_grpc as test_public_grpc
import yandex.cloud.priv.loadtesting.agent.v1.agent_service_pb2 as agent_public
import yandex.cloud.priv.loadtesting.agent.v1.agent_service_pb2_grpc as agent_public_grpc
import yandex.cloud.priv.loadtesting.agent.v1.agent_registration_service_pb2 as agent_registration_public
import yandex.cloud.priv.loadtesting.agent.v1.agent_registration_service_pb2_grpc as agent_registration_grpc
from grpc_reflection.v1alpha import reflection
from yandex.cloud.priv.loadtesting.v1 import operation_service_pb2 as operation
from yandex.cloud.priv.loadtesting.v1 import operation_service_pb2_grpc as operation_grpc
from yandex.cloud.priv.loadtesting.v1 import resource_preset_service_pb2 as preset
from yandex.cloud.priv.loadtesting.v1 import resource_preset_service_pb2_grpc as preset_grpc
from yandex.cloud.priv.loadtesting.v1 import storage_service_pb2 as storage
from yandex.cloud.priv.loadtesting.v1 import storage_service_pb2_grpc as storage_grpc
from yandex.cloud.priv.loadtesting.v1 import tank_instance_service_pb2 as tank_private
from yandex.cloud.priv.loadtesting.v1 import tank_instance_service_pb2_grpc as tank_private_grpc
from yandex.cloud.priv.loadtesting.v1 import tank_job_service_pb2 as job_private
from yandex.cloud.priv.loadtesting.v1 import tank_job_service_pb2_grpc as job_private_grpc
from yandex.cloud.priv.loadtesting.v2 import test_service_pb2 as test_private
from yandex.cloud.priv.loadtesting.v2 import test_service_pb2_grpc as test_private_grpc
from yandex.cloud.priv.loadtesting.v2 import agent_instance_service_pb2 as agent_private
from yandex.cloud.priv.loadtesting.v2 import agent_instance_service_pb2_grpc as agent_private_grpc
from yandex.cloud.priv.loadtesting.v2 import stats_service_pb2 as stats
from yandex.cloud.priv.loadtesting.v2 import stats_service_pb2_grpc as stats_grpc

import grpc
from load.projects.cloud.cloud_helper import metadata_compute
from load.projects.cloud.loadtesting import config
from load.projects.cloud.loadtesting import logan
from load.projects.cloud.loadtesting.server.api.agent_public.registration import AgentRegistration
from load.projects.cloud.loadtesting.server.api.private_v1.job import JobServicer
from load.projects.cloud.loadtesting.server.api.agent_public.job import Jobber
from load.projects.cloud.loadtesting.server.api.private_v1.preset import PresetServicer
from load.projects.cloud.loadtesting.server.api.agent_public.monitoring import Monitoring
from load.projects.cloud.loadtesting.server.api.agent_public.tank import Tanker
from load.projects.cloud.loadtesting.server.api.agent_public.agent import AgentServicer
from load.projects.cloud.loadtesting.server.api.agent_public.trail import Trailer
from load.projects.cloud.loadtesting.server.api.agent_public.test import TestServicer
from load.projects.cloud.loadtesting.server.api.private_v2.test import TestPrivateServicer
from load.projects.cloud.loadtesting.server.api.private_v1.storage import StorageServicer
from load.projects.cloud.loadtesting.server.api.common.utils import PushMonitoringData
from load.projects.cloud.loadtesting.server.api.common.handler import registry
from load.projects.cloud.cloud_helper.grpc_options import COMMON_CHANNEL_OPTIONS
from load.projects.cloud.loadtesting.server.api.private_v1.operation import OperationServicer
from load.projects.cloud.loadtesting.server.api.private_v1.tank import TankServicer
from load.projects.cloud.loadtesting.server.api.private_v2.agent import AgentServicer as AgentPrivateServicer
from load.projects.cloud.loadtesting.server.api.private_v2.stats import StatsServicer


class API:
    def __init__(self, insecure_port, secure_port=None):
        permanent_log_level = logan.parse_level(config.ENV_CONFIG.PERMANENT_LOG_LEVEL)
        instance_name = metadata_compute.get_current_instance_name()
        self.logging_controller = logan.Controller(
            permanent_log_level,
            config.ENV_CONFIG.ENABLE_LOCAL_LOGGING == 'true',
            config.ENV_CONFIG.UNIFIED_AGENT_URI,
            config.ENV_CONFIG.TEST_MODE == 'true',
        )

        self.insecure_port = insecure_port
        self.secure_port = secure_port
        self.logger: logan.Logan = logging.getLogger('loadtesting')
        self.logger.bind_global('server_compute_id', metadata_compute.get_current_instance_id())
        self.logger.bind_global('server_name', instance_name)
        self.max_workers = int(config.ENV_CONFIG.MAX_WORKERS)
        self._server_thread_pool = None
        PushMonitoringData.setup(registry, config.ENV_CONFIG.PROMETHEUS_PUSHGATEWAY, self.logger, instance_name)

        self._server_kwargs = dict(options=COMMON_CHANNEL_OPTIONS)
        if config.ENV_CONFIG.MAX_QUEUE is not None and (max_queue := int(config.ENV_CONFIG.MAX_QUEUE)) >= 0:
            self._server_kwargs['maximum_concurrent_rpcs'] = max_queue + self.max_workers

    def __enter__(self):
        self._server_thread_pool = futures.ThreadPoolExecutor(max_workers=self.max_workers).__enter__()
        self.server = grpc.server(
            self._server_thread_pool,
            **self._server_kwargs,
        )

        tank_private_grpc.add_TankInstanceServiceServicer_to_server(TankServicer(), self.server)
        agent_private_grpc.add_AgentInstanceServiceServicer_to_server(AgentPrivateServicer(), self.server)
        job_private_grpc.add_TankJobServiceServicer_to_server(JobServicer(), self.server)
        tank_grpc.add_TankServiceServicer_to_server(Tanker(), self.server)
        agent_public_grpc.add_AgentServiceServicer_to_server(AgentServicer(), self.server)
        job_grpc.add_JobServiceServicer_to_server(Jobber(), self.server)
        monitoring_grpc.add_MonitoringServiceServicer_to_server(Monitoring(), self.server)
        trail_grpc.add_TrailServiceServicer_to_server(Trailer(), self.server)
        operation_grpc.add_OperationServiceServicer_to_server(OperationServicer(), self.server)
        preset_grpc.add_ResourcePresetServiceServicer_to_server(PresetServicer(), self.server)
        agent_registration_grpc.add_AgentRegistrationServiceServicer_to_server(AgentRegistration(), self.server)
        test_public_grpc.add_TestServiceServicer_to_server(TestServicer(), self.server)
        storage_grpc.add_StorageServiceServicer_to_server(StorageServicer(), self.server)
        test_private_grpc.add_TestServiceServicer_to_server(TestPrivateServicer(), self.server)
        stats_grpc.add_StatsServiceServicer_to_server(StatsServicer(), self.server)

        service_names = (
            tank_private.DESCRIPTOR.services_by_name['TankInstanceService'].full_name,
            agent_private.DESCRIPTOR.services_by_name['AgentInstanceService'].full_name,
            job_private.DESCRIPTOR.services_by_name['TankJobService'].full_name,
            test_private.DESCRIPTOR.services_by_name['TestService'].full_name,
            operation.DESCRIPTOR.services_by_name['OperationService'].full_name,
            preset.DESCRIPTOR.services_by_name['ResourcePresetService'].full_name,
            tank_public.DESCRIPTOR.services_by_name['TankService'].full_name,
            agent_public.DESCRIPTOR.services_by_name['AgentService'].full_name,
            job_public.DESCRIPTOR.services_by_name['JobService'].full_name,
            monitoring_public.DESCRIPTOR.services_by_name['MonitoringService'].full_name,
            trail_public.DESCRIPTOR.services_by_name['TrailService'].full_name,
            agent_registration_public.DESCRIPTOR.services_by_name['AgentRegistrationService'].full_name,
            test_public.DESCRIPTOR.services_by_name['TestService'].full_name,
            storage.DESCRIPTOR.services_by_name['StorageService'].full_name,
            stats.DESCRIPTOR.services_by_name['StatsService'].full_name,
            reflection.SERVICE_NAME,)
        reflection.enable_server_reflection(service_names, self.server)

        self.server.add_insecure_port(f'[::]:{self.insecure_port}')
        if self.secure_port:
            with open(config.ENV_CONFIG.CERT_KEY_PATH, 'rb') as f:
                private_key = f.read()
            with open(config.ENV_CONFIG.CERT_PATH, 'rb') as f:
                certificate_chain = f.read()
            server_credentials = grpc.ssl_server_credentials(((private_key, certificate_chain),))
            self.server.add_secure_port(f'[::]:{self.secure_port}', server_credentials)

        self.server.start()
        self.logger.send(f"Loadtesting server was started. Secure port {self.secure_port}."
                         f" Insecure port {self.insecure_port}.")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.server.stop(None)
        self._server_thread_pool.__exit__(exc_type, exc_val, exc_tb)
