import io
import grpc
import uuid
import json
from datetime import datetime
from sqlalchemy.exc import SQLAlchemyError

from google.protobuf.json_format import MessageToJson
from load.projects.cloud.cloud_helper import aws
import load.projects.cloud.loadtesting.server.api.common.permissions as permissions
from load.projects.cloud.loadtesting.config import ENV_CONFIG, DEFAULT_PAGE_SIZE
from load.projects.cloud.loadtesting.db.tables import StorageTable, OperationTable, ResourceType
from load.projects.cloud.loadtesting.logan import lookup_logger
from load.projects.cloud.loadtesting.server.api.common import handler
from yandex.cloud.priv.loadtesting.v1 import storage_service_pb2_grpc, storage_service_pb2, \
    storage_pb2
from load.projects.cloud.loadtesting.server.api.private_v1.operation import create_operation_message, \
    create_nested_message
from load.projects.cloud.loadtesting.server.api.private_v1.ammo import create_ammo
from load.projects.cloud.loadtesting.server.api.common.utils import generate_cloud_id, authorize, ts_from_dt
from load.projects.cloud.loadtesting.server.pager import Pager


def _create_operation(logger, db, folder_id, description, storage_id, user_id, metadata_message,
                      foreign_operation_id=None, error=None, done=False, message=None):
    db_operation = OperationTable(
        id=generate_cloud_id(),
        folder_id=folder_id,
        foreign_operation_id=foreign_operation_id,
        description=description,
        created_at=datetime.utcnow(),
        created_by=user_id,
        target_resource_id=storage_id,
        target_resource_type=ResourceType.STORAGE.value,
        error=error,
        done=done
    )
    db_operation.resource_metadata = create_nested_message(metadata_message)
    if done:
        db_operation.done_resource_snapshot = MessageToJson(message)
    db.operation.add(db_operation)
    return create_operation_message(db_operation, logger)


class _StorageStatMock:
    def __init__(self):
        self.used_size = 0
        self.storage_class_counters = []


def _get_storage(logger, storage, user_token, request_id, mock_stat_on_aws_error=False):
    try:
        stat = aws.get_bucket_stats(storage.bucket, user_token, request_id)
    except grpc.RpcError as e:
        logger.warning('Failed to get bucket %s stat: %s', storage.bucket, e)
        if mock_stat_on_aws_error:
            stat = _StorageStatMock()
        else:
            raise

    simple_object_count = 0
    for class_counter in stat.storage_class_counters:
        simple_object_count += class_counter.counters.simple_object_count
    return storage_pb2.Storage(
        id=storage.id,
        folder_id=storage.folder_id,
        name=storage.name,
        description=storage.description,
        labels=storage.labels,
        created_at=ts_from_dt(storage.created_at),
        object_storage_bucket=storage.bucket,

        used_size=stat.used_size,
        simple_object_count=simple_object_count,
    )


class StorageServicer(storage_service_pb2_grpc.StorageServiceServicer):

    def __init__(self):
        self.logger = lookup_logger('StoragePrivate')

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

        def proceed(self):
            storage = self.db.storage.get(self.request.storage_id)
            if storage is None:
                self.context.abort(grpc.StatusCode.NOT_FOUND, f'There is no storage {self.request.storage_id}')
            authorize(self.context, self.user_token, storage.folder_id, permissions.STORAGE_GET, self.request_id)
            return _get_storage(self.logger, storage, self.user_token, self.request_id)

    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

            authorize(self.context, self.user_token, request.folder_id, permissions.STORAGE_GET, self.request_id)
            pager = Pager(request.page_token,
                          request.page_size or DEFAULT_PAGE_SIZE,
                          request.filter)

            db_storages = self.db.storage.get_by_folder(request.folder_id,
                                                        offset=pager.offset,
                                                        limit=pager.page_size)

            pager.set_shift(len(db_storages))
            storages = []
            for db_storage in db_storages:
                storages.append(_get_storage(self.logger, db_storage, self.user_token, self.request_id,
                                             mock_stat_on_aws_error=True))

            return storage_service_pb2.ListStorageResponse(storages=storages, next_page_token=pager.next_page_token)

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

    class _Create(handler.BasePrivateHandler):
        _handler_name = 'Create'

        def proceed(self):
            authorize(self.context, self.user_token, self.request.folder_id, permissions.STORAGE_CREATE, self.request_id)
            storage = StorageTable(
                id=generate_cloud_id(),
                folder_id=self.request.folder_id,
                bucket=self.request.object_storage_bucket_name,
                name=self.request.name or self.request.object_storage_bucket_name,
                description=self.request.description,
                created_at=datetime.utcnow()
            )
            self.db.storage.add(storage)
            metadata_message = storage_service_pb2.CreateStorageMetadata(storage_id=storage.id)
            return _create_operation(self.logger, self.db, self.request.folder_id, 'Create storage', storage.id,
                                     self.user_id, metadata_message, done=True, message=DbToGrpcTranslator.storage(storage))

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

    class _Update(handler.BasePrivateHandler):
        _handler_name = 'Update'

        def proceed(self):
            storage = self.db.storage.get(self.request.storage_id)
            if storage is None:
                self.context.abort(grpc.StatusCode.NOT_FOUND,
                                   f'There is no storage {self.request.storage_id}')
            authorize(self.context, self.user_token, storage.folder_id, permissions.STORAGE_UPDATE, self.request_id)
            if self.request.name:
                storage.name = self.request.name
            if self.request.description:
                storage.description = self.request.description

            metadata_message = storage_service_pb2.UpdateStorageMetadata(storage_id=storage.id)
            try:
                self.db.storage.add(storage)
                error = None
            except SQLAlchemyError:
                error = f'Failed to update storage {storage.id}.'
                self.context.abort(grpc.StatusCode.INTERNAL, error)
            finally:
                update_operation = _create_operation(self.logger, self.db, storage.folder_id, 'Update storage', storage.id, self.user_id,
                                                     metadata_message, done=True, message=DbToGrpcTranslator.storage(storage), error=error)
            return update_operation

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

    class _Delete(handler.BasePrivateHandler):
        _handler_name = 'Delete'

        def proceed(self):
            storage = self.db.storage.get(self.request.storage_id)
            if storage is None:
                self.context.abort(grpc.StatusCode.NOT_FOUND,
                                   f'There is no storage {self.request.storage_id}')
            authorize(self.context, self.user_token, storage.folder_id, permissions.STORAGE_DELETE, self.request_id)
            metadata_message = storage_service_pb2.DeleteStorageMetadata(storage_id=storage.id)
            self.db.storage.delete(storage)
            return _create_operation(self.logger, self.db, storage.folder_id, 'Delete storage', storage.id,
                                     self.user_id, metadata_message, done=True,
                                     message=DbToGrpcTranslator.storage(storage))

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

    class _UploadObject(handler.BasePrivateHandler):
        _handler_name = 'UploadObject'

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

            authorize(self.context, self.user_token, request.folder_id, permissions.STORAGE_UPLOAD, self.request_id)

            if not request.test_data:
                context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'Test data file is empty')

            storage = None
            if request.storage_id:
                storage = db.storage.get(request.storage_id)
                if storage is None:
                    self.context.abort(grpc.StatusCode.NOT_FOUND,
                                       f'There is no storage {self.request.storage_id}')
                if storage.folder_id != request.folder_id:
                    context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'The received storage_id has different folder')
                linked_bucket = storage.bucket
            else:
                linked_bucket = None

            if linked_bucket is not None:
                s3_filename = request.filename  # save original name
            else:
                s3_filename = f'test_data_{uuid.uuid4()}'
            self.logger.debug(f'test data with original name "{request.filename}" will be stored as "{s3_filename}"')

            buffer = io.BytesIO()
            buffer.write(request.test_data)
            buffer.seek(0)  # read и write на io работают как курсор, так что после write - мы в конце "файла"

            if linked_bucket is not None:
                upload_result = aws.upload_by_presign_url(linked_bucket, s3_filename, buffer, self.user_token, self.request_id)
            else:
                # upload to our storage
                upload_result = aws.upload_fileobj(buffer, s3_filename, bucket=ENV_CONFIG.OBJECT_STORAGE_DEFAULT_BUCKET)

            if upload_result is False:
                context.abort(grpc.StatusCode.INTERNAL, f'Failed to upload test data {s3_filename}')

            try:
                ammo = create_ammo(db, s3_filename, request.folder_id, storage)
            except SQLAlchemyError:
                context.abort(grpc.StatusCode.INTERNAL, f'Failed to save test data {s3_filename}')

            return storage_pb2.StorageObject(
                object_storage_bucket=linked_bucket,
                object_storage_filename=ammo.s3_name
            )

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


class DbToGrpcTranslator:

    @staticmethod
    def storage(db_storage: StorageTable) -> storage_pb2.Storage:
        return storage_pb2.Storage(
            id=db_storage.id,
            folder_id=db_storage.folder_id,
            created_at=ts_from_dt(db_storage.created_at),
            name=db_storage.name,
            description=db_storage.description,
            labels=json.loads(db_storage.labels) if db_storage.labels else None,
            object_storage_bucket=db_storage.bucket
        )
