from yandex.cloud.priv.loadtesting.agent.v1 import agent_registration_service_pb2, agent_registration_service_pb2_grpc
from yandex.cloud.priv.loadtesting.v2 import agent_instance_pb2

import grpc
import json
from datetime import datetime
from functools import cached_property
from google.protobuf.json_format import MessageToJson
from load.projects.cloud.loadtesting.server.api.common import permissions, handler
from load.projects.cloud.loadtesting.server.api.common.utils import authorize, generate_cloud_id
from load.projects.cloud.loadtesting.logan import lookup_logger
from load.projects.cloud.loadtesting.server.api.private_v1.tank import METADATA_AGENT_VERSION_ATTR
from load.projects.cloud.loadtesting.server.api.private_v1.operation import create_nested_message, create_operation_message
from load.projects.cloud.loadtesting.db.tables import TankTable, OperationTable, ResourceType
from load.projects.cloud.cloud_helper import compute


class AgentRegistration(agent_registration_service_pb2_grpc.AgentRegistrationServiceServicer):
    def __init__(self):
        self.logger = lookup_logger('AgentRegistrationPublic')

    def Register(self, request, context):
        return Register(self.logger).handle(request, context)

    def ExternalAgentRegister(self, request, context):
        return ExternalAgentRegister(self.logger).handle(request, context)


class Register(handler.BasePublicHandler):
    _handler_name = 'Register'

    def proceed(self):
        compute_instance_id = self.request.compute_instance_id
        self.logger.bind('agent_compute_instance_id', compute_instance_id)
        if not self.db_tank:
            self.logger.warn(f'An agent with {compute_instance_id=} is not found')
            raise self.context.abort(grpc.StatusCode.NOT_FOUND,
                                     f'There is no compute instance {compute_instance_id} in loadtesting database')

        authorize(self.context, self.user_token, self.db_tank.folder_id, permissions.TESTS_RUN, self.request_id)
        self.validate_or_set_agent_version()
        return agent_registration_service_pb2.RegisterResponse(agent_instance_id=self.db_tank.id)

    @cached_property
    def db_tank(self):
        return self.db.tank.get_by_compute_instance_id(self.request.compute_instance_id)

    def validate_or_set_agent_version(self):
        for k, version_from_agent in self.context.invocation_metadata():
            if k == METADATA_AGENT_VERSION_ATTR:
                if self.db_tank.agent_version and self.db_tank.agent_version != version_from_agent:
                    self.logger.error('Agent version inconsistency! '
                                      f'Version in DB "{self.db_tank.agent_version}". '
                                      f'Version from agent "{version_from_agent}".')
                    raise self.context.abort(grpc.StatusCode.INVALID_ARGUMENT,
                                             'Agent version inconsistency.')
                self.db.tank.update_agent_version(self.db_tank, version_from_agent)


class ExternalAgentRegister(handler.BasePublicHandler):
    _handler_name = 'ExternalAgentRegister'

    def proceed(self):
        if self.request.compute_instance_id:
            db_tank = self.db.tank.get_by_compute_instance_id(self.request.compute_instance_id)
            if db_tank:
                raise self.context.abort(grpc.StatusCode.ALREADY_EXISTS,
                                         f'The agent with compute_instance_id={self.request.compute_instance_id} already exists')
            try:
                compute_instance = compute.get_instance(self.request.compute_instance_id, self.user_token, self.request_id, self.logger)
            except grpc.RpcError as error:
                if error.code() == grpc.StatusCode.NOT_FOUND:
                    raise self.context.abort(grpc.StatusCode.NOT_FOUND,
                                             f'Compute instance {self.request.compute_instance_id} is not found')
                raise error
            service_account_id = authorize(self.context, self.user_token, compute_instance.folder_id, permissions.AGENTS_REGISTER, self.request_id)
            tank_id = generate_cloud_id()
            db_tank = TankTable(
                id=tank_id,
                status=None,
                folder_id=compute_instance.folder_id,
                created_at=datetime.utcnow(),
                compute_instance_updated_at=datetime.utcnow(),
                client_updated_at=None,
                description="External compute agent",
                service_account_id=service_account_id,
                compute_instance_id=compute_instance.id,
                name=compute_instance.name,
                labels=json.dumps(list(compute_instance.labels)),
                # agent_version='',  # TODO
            )
        else:
            if not self.request.folder_id:
                raise self.context.abort(grpc.StatusCode.INVALID_ARGUMENT,
                                         'There is neither compute_instance_id nor folder_id')
            service_account_id = authorize(self.context, self.user_token, self.request.folder_id,
                                           permissions.AGENTS_REGISTER, self.request_id)
            tank_id = generate_cloud_id()
            db_tank = TankTable(
                id=tank_id,
                status=None,
                folder_id=self.request.folder_id,
                created_at=datetime.utcnow(),
                compute_instance_updated_at=datetime.utcnow(),
                client_updated_at=None,
                description="External agent",
                service_account_id=service_account_id,
                name=self.request.name,
                # agent_version='',  # TODO
            )

        self.db.tank.add(db_tank)

        metadata_message = agent_registration_service_pb2.ExternalAgentRegisterMetadata(agent_instance_id=db_tank.id)
        resource_metadata = create_nested_message(metadata_message)
        register_operation = OperationTable(
            id=generate_cloud_id(),
            folder_id=db_tank.folder_id,
            target_resource_id=db_tank.id,
            target_resource_type=ResourceType.AGENT.value,
            description='Register an agent',
            resource_metadata=resource_metadata,
            created_by=service_account_id,
            done=True,
            done_resource_snapshot=MessageToJson(agent_instance_pb2.AgentInstance(id=db_tank.id)),  # TODO
        )
        self.db.operation.add(register_operation)
        return create_operation_message(register_operation, self.logger)
