import typing
from datetime import datetime, timedelta

import google.protobuf.message
from google.protobuf.any_pb2 import Any
from google.protobuf.empty_pb2 import Empty
from google.protobuf.json_format import Parse, MessageToJson
from google.rpc.code_pb2 import Code as rpc_code
from google.rpc.status_pb2 import Status as rpc_status
from yandex.cloud.priv.loadtesting.v1 import operation_service_pb2, operation_service_pb2_grpc
from yandex.cloud.priv.loadtesting.v1 import tank_instance_pb2
from yandex.cloud.priv.loadtesting.v1 import tank_job_pb2
from yandex.cloud.priv.loadtesting.v1 import storage_pb2
from yandex.cloud.priv.loadtesting.v2 import agent_instance_pb2
from yandex.cloud.priv.operation import operation_pb2

import grpc
from load.projects.cloud.cloud_helper import compute
from load.projects.cloud.loadtesting.config import DEFAULT_PAGE_SIZE
from load.projects.cloud.loadtesting.db.tables import ResourceType
from load.projects.cloud.loadtesting.logan import lookup_logger
from load.projects.cloud.loadtesting.server.api.common import permissions, handler
from load.projects.cloud.loadtesting.server.pager import Pager
from load.projects.cloud.loadtesting.server.api.common.utils import authorize
from load.projects.cloud.loadtesting.server.api.common.utils import ts_from_dt
# from multiprocessing.pool import Pool

OPERATION_LIMIT_TIME = timedelta(minutes=10)


class OperationServicer(operation_service_pb2_grpc.OperationServiceServicer):
    def __init__(self):
        self.logger = lookup_logger('Operations')

    class _Get(handler.BasePrivateHandler):
        _handler_name = 'Get'

        def proceed(self):
            operation = self.db.operation.get(self.request.operation_id)
            if not operation:
                self.context.abort(grpc.StatusCode.NOT_FOUND, f"Operation {self.request.operation_id} is not found")
            _authorize(self.context, operation.folder_id, operation.target_resource_type, self.user_token, self.request_id)
            _update(operation, self.db, self.user_token, self.request_id, self.logger)
            return create_operation_message(operation, self.logger)

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

    class _List(handler.BasePrivateHandler):
        _handler_name = 'List'

        def proceed(self):
            request = self.request
            db = self.db

            pager = Pager(request.page_token, request.page_size or DEFAULT_PAGE_SIZE)
            operations = db.operation.get_by_folder(request.folder_id,
                                                    offset=pager.offset,
                                                    limit=pager.page_size)
            pager.set_shift(len(operations))
            operations_list = []
            # TODO
            # args = [(self.context, operation.folder_id, operation.target_resource_type, self.user_token, self.request_id) for operation in operations]
            # with Pool(len(operations)) as p:  # count??
            #     p.starmap_async(_authorize, args)
            authorize(self.context, self.user_token, request.folder_id, permissions.AGENTS_GET, self.request_id)
            authorize(self.context, self.user_token, request.folder_id, permissions.TESTS_GET, self.request_id)
            for operation in operations:
                _update(operation, self.db, self.user_token, self.request_id, self.logger)
                operations_list.append(create_operation_message(operation, self.logger))
            return operation_service_pb2.ListOperationsResponse(operations=operations_list,
                                                                next_page_token=pager.next_page_token)

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


def _authorize(context, folder_id, target_resource_type, user_token, request_id):
    if target_resource_type == ResourceType.TANK.value:
        authorize(context, user_token, folder_id, permissions.AGENTS_GET, request_id)
    elif target_resource_type == ResourceType.JOB.value:
        authorize(context, user_token, folder_id, permissions.TESTS_GET, request_id)
    else:
        context.abort(grpc.StatusCode.INTERNAL, "Unknown resource type for authorization")


def _update(operation, db, iam_user_token, request_id, logger):
    if operation.target_resource_type == ResourceType.TANK.value and not operation.done and operation.foreign_operation_id:
        try:
            compute_operation = compute.get_operation(operation.foreign_operation_id, iam_user_token, request_id)
        except grpc.RpcError as error:
            if error.code() == grpc.StatusCode.NOT_FOUND:
                logger.warn(f'There is no operation {operation.id} in compute')
                return
            raise error
        db.operation.update(
            operation,
            error=compute_operation.error.message,
            modified_at=compute_operation.modified_at.ToDatetime(),
            done=compute_operation.done)
        # TODO work out deleting and snapshot
    if not operation.done and (datetime.utcnow() - operation.created_at) > OPERATION_LIMIT_TIME:
        db.operation.update(
            operation,
            error='Unknown error',
            modified_at=datetime.utcnow(),
            done=True)


def create_operation_message(db_operation, logger):
    created_at = ts_from_dt(db_operation.created_at)
    modified_at = ts_from_dt(db_operation.modified_at)
    operation = operation_pb2.Operation(
        id=db_operation.id,
        description=db_operation.description,
        created_at=created_at,
        created_by=db_operation.created_by,
        modified_at=modified_at,
        done=db_operation.done)

    if db_operation.resource_metadata:
        operation.metadata.CopyFrom(Parse(db_operation.resource_metadata, Any()))

    if db_operation.error:
        # TODO: code depends on error
        error = rpc_status(message=db_operation.error, code=rpc_code.INVALID_ARGUMENT)
        operation.error.CopyFrom(error)
    elif db_operation.done_resource_snapshot:
        if db_operation.target_resource_type == ResourceType.TANK.value:
            message_type = tank_instance_pb2.TankInstance
        elif db_operation.target_resource_type == ResourceType.JOB.value:
            message_type = tank_job_pb2.TankJob
        elif db_operation.target_resource_type == ResourceType.STORAGE.value:
            message_type = storage_pb2.Storage
        elif db_operation.target_resource_type == ResourceType.AGENT.value:
            message_type = agent_instance_pb2.AgentInstance
        else:
            raise Exception('Unknown resource type')
        snapshot_instance = Parse(db_operation.done_resource_snapshot, message_type())
        response = Any()
        response.Pack(snapshot_instance)
        operation.response.CopyFrom(response)
    else:
        response = Any()
        response.Pack(Empty())
        operation.response.CopyFrom(response)

    logger.debug('Get operation: %s', operation)
    return operation


def create_nested_message(message: google.protobuf.message.Message) -> typing.Mapping[str, str]:
    nested_message = Any()
    nested_message.Pack(message)
    return MessageToJson(nested_message)
