import json
from abc import ABC
from datetime import datetime, timedelta
from functools import cached_property
from multiprocessing.pool import ThreadPool
from threading import Thread
import time

from google.protobuf.json_format import MessageToJson

import grpc
import load.projects.cloud.loadtesting.server.api.common.permissions as permissions
from load.projects.cloud.cloud_helper import compute
from load.projects.cloud.loadtesting.config import DEFAULT_PAGE_SIZE, ENV_CONFIG, EnvType
from load.projects.cloud.loadtesting.db import DB
from load.projects.cloud.loadtesting.db.tables import TankTable, TankStatus, OperationTable, ResourceType, \
    AgentVersionTable, AgentVersionStatus, STATUS_COMMENTS
from load.projects.cloud.loadtesting.logan import lookup_logger
from load.projects.cloud.loadtesting.server.api.common import handler
from yandex.cloud.priv.compute.v1 import instance_pb2
from yandex.cloud.priv.loadtesting.v1 import tank_instance_service_pb2, tank_instance_service_pb2_grpc, \
    tank_instance_pb2
from load.projects.cloud.loadtesting.server.api.private_v1.operation import create_nested_message, create_operation_message
from load.projects.cloud.loadtesting.server.pager import Pager
from load.projects.cloud.loadtesting.server.api.common.utils import generate_cloud_id, authorize, ts_from_dt, waiting_for_operation, WAITING_TIMEOUT

TANK_IMAGE_FAMILY_ATTR = 'tank_image_family'
TANK_IMAGE_ID_ATTR = 'tank_image_id'
UPDATE_DELAY_TIME = timedelta(minutes=1)
PALTFORM_ID = "standard-v2"
METADATA_HOST_ATTR = 'server-host'
METADATA_PORT_ATTR = 'server-port'
METADATA_AGENT_VERSION_ATTR = 'agent-version'
METADATA_REQUEST_FREQUENCY = 'request-frequency'
METADATA_LOGGING_HOST_ATTR = 'cloud-helper-logging-host'
METADATA_LOGGING_PORT_ATTR = 'cloud-helper-logging-port'
METADATA_OBJECT_STORAGE_URL_ATTR = 'cloud-helper-object-storage-url'
METADATA_LT_CREATED_ATTR = 'loadtesting-created'

NOT_SET = object()


class ComputeOperationError(Exception):
    pass


class TankHandler(handler.BasePrivateHandler, ABC):
    @property
    def tank_id(self):
        return self.request.id

    @cached_property
    def db_tank(self):
        db_tank = self.db.tank.get(self.tank_id)
        if not db_tank:
            self.context.abort(grpc.StatusCode.NOT_FOUND,
                               f"The agent {self.tank_id} is not found")
        return db_tank

    def _get(self, db_tank, compute_instance=None, db_version=NOT_SET):
        # None - допустимое значение для db_version,
        # поэтому для указания, что db_version не задан, заведён отдельный объект.

        if compute_instance is not None:
            final_status = TankHandler._tank_status(db_tank.client_status, db_tank.client_updated_at,
                                                    compute_instance.status)
            self.logger.debug(
                f"Current status of the agent instance {db_tank.id} is {final_status}")
            if db_tank.name != compute_instance.name:
                db_tank.name = compute_instance.name
            if db_tank.description != compute_instance.description:
                db_tank.description = compute_instance.description
            if db_tank.service_account_id != compute_instance.service_account_id:
                db_tank.service_account_id = compute_instance.service_account_id
            if db_tank.status != final_status:
                db_tank.status = final_status
            self.logger.debug("%s db tank has been synchronized with compute instance: %s", db_tank, compute_instance)
        else:
            if db_tank.status != db_tank.client_status:
                db_tank.status = db_tank.client_status
        if db_version is NOT_SET:
            db_version = self.db.agent_version.get(db_tank.agent_version)

        return DbToGrpcTranslator.tank(db_tank, db_version, self.lang)

    def _create_operation(self, user_id, db_tank, compute_operation, meta_message, description, is_deleting=False):
        self.db.tank.set_update_time(db_tank)
        # TODO error status?
        db_operation = OperationTable(
            id=generate_cloud_id(),
            folder_id=db_tank.folder_id,
            foreign_operation_id=compute_operation.id,
            target_resource_id=db_tank.id,
            target_resource_type=ResourceType.TANK.value,
            resource_metadata=meta_message,
            description=description,
            created_at=compute_operation.created_at.ToDatetime(),
            created_by=user_id,
            modified_at=compute_operation.modified_at.ToDatetime(),
            done=compute_operation.done,
            error=compute_operation.error.message)
        self.db.operation.add(db_operation)
        self.db.commit()
        operation = create_operation_message(db_operation, self.logger)

        # waiting for operation finishing
        # TODO join??
        t = Thread(target=TankHandler._waiting_for_operation,
                   args=(self.logger, self.user_token, db_tank.id, db_operation.id, compute_operation, self.request_id,
                         is_deleting,))
        t.start()
        return operation

    @staticmethod
    def _tank_status(db_tank_client_status, client_updated_time, compute_instance_status):
        if compute_instance_status != instance_pb2.Instance.Status.Value('RUNNING'):
            return instance_pb2.Instance.Status.Name(compute_instance_status)
        if db_tank_client_status is None:
            return TankStatus.INITIALIZING_CONNECTION.value
        if (datetime.utcnow() - client_updated_time) < UPDATE_DELAY_TIME:
            return db_tank_client_status
        return TankStatus.LOST_CONNECTION_WITH_TANK.value

    @staticmethod
    def _waiting_for_operation(parent_logger, user_token, tank_id, operation_id, compute_operation, request_id,
                               is_deleting=False):
        logger = parent_logger.getChild('deattached_waiting_for_operation')
        logger.debug('starting to wait')
        try:
            with DB() as db:
                db_operation = db.operation.get(operation_id)
                done_compute_operation = waiting_for_operation(db, db_operation, compute.get_operation, compute_operation, user_token, request_id)
                db_tank = db.tank.get(tank_id)
                if is_deleting and db_tank and not done_compute_operation.HasField('error'):
                    db.tank.delete(db_tank)
                else:
                    compute_instance = compute.get_instance(db_tank.compute_instance_id, user_token, request_id, logger)
                    db_tank.status = TankHandler._tank_status(db_tank.client_status, db_tank.client_updated_at,
                                                              compute_instance.status)
                    tank_instance = DbToGrpcTranslator.tank(db_tank, db.agent_version.get(db_tank.agent_version))
                    db.operation.add_snapshot(db_operation, MessageToJson(tank_instance))
        finally:
            logger.debug('waiting completed')

    def get_compute_instance(self, compute_instance_id):
        if compute_instance_id is None:
            return
        try:
            return compute.get_instance(compute_instance_id, self.user_token, self.request_id, self.logger)
        except grpc.RpcError as error:
            if error.code() == grpc.StatusCode.NOT_FOUND:
                return
            raise error

    def check_for_external_agent(self):
        if self.db_tank.compute_instance_id is None:
            raise self.context.abort(grpc.StatusCode.INVALID_ARGUMENT,
                                     "The operation is not supported for external agents")


class TankServicer(tank_instance_service_pb2_grpc.TankInstanceServiceServicer):

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

    class _Get(TankHandler):
        _handler_name = 'Get'

        def proceed(self):
            authorize(self.context, self.user_token, self.db_tank.folder_id, permissions.AGENTS_GET, self.request_id)
            compute_instance = None
            if self.db_tank.compute_instance_id:
                compute_instance = self.get_compute_instance(self.db_tank.compute_instance_id)
                if compute_instance is None:
                    self.logger.warn(f"Compute doesn't have instance {self.db_tank.compute_instance_id}")
                    self.db.tank.delete(self.db_tank)
                    raise self.context.abort(grpc.StatusCode.NOT_FOUND, f"Instance {self.db_tank.id} is not found")

            tank = self._get(self.db_tank, compute_instance)
            return tank_instance_service_pb2.GetTankInstanceResponse(tank_instance=tank)

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

    class _List(TankHandler):
        _handler_name = 'List'
        # Общий pool чтобы ограничить кол-во соединений к Compute, иначе он начинает капризничать.
        # Как капризничает Compute: https://paste.yandex-team.ru/6874070
        parallel_requests_compute_can_handle = ThreadPool(100)

        def _get_compute_instances(self, compute_ids):
            compute_instances = self.parallel_requests_compute_can_handle.map(self.get_compute_instance, compute_ids)
            return compute_instances

        def proceed(self):
            request = self.request

            # TODO check every resource CLOUDLOAD-95
            authorize(self.context, self.user_token, request.folder_id, permissions.AGENTS_GET, self.request_id)
            pager = Pager(request.page_token,
                          request.page_size or DEFAULT_PAGE_SIZE,
                          request.filter)

            tanks_with_version = self.db.tanks_with_version_by_folder(request.folder_id,
                                                                      offset=pager.offset,
                                                                      limit=pager.page_size)

            pager.set_shift(len(tanks_with_version))

            compute_instances = self._get_compute_instances(tank.compute_instance_id
                                                            for (tank, _) in tanks_with_version)

            tank_instances = []
            for (tank, version), compute_instance in zip(tanks_with_version, compute_instances):
                if tank.compute_instance_id is not None and compute_instance is None:
                    self.logger.warn(f"Compute doesn't have instance {tank.compute_instance_id}")
                    self.db.tank.delete(tank)
                    continue
                tank_instances.append(self._get(tank, compute_instance, version))

            return tank_instance_service_pb2.TankInstancesList(folder_id=request.folder_id,
                                                               tank_instances=tank_instances,
                                                               next_page_token=pager.next_page_token)

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

    class _Create(TankHandler):
        _handler_name = 'Create'

        @cached_property
        def _get_image(self):
            if ENV_CONFIG.TARGET_AGENT_VERSION:
                if self.db.agent_version.get(ENV_CONFIG.TARGET_AGENT_VERSION) is None:
                    raise self.context.abort(grpc.StatusCode.INTERNAL,
                                             f"Agent versions from config {ENV_CONFIG.TARGET_AGENT_VERSION}"
                                             " has not been registered in database.")
                return ENV_CONFIG.TARGET_AGENT_VERSION

            if (version := self.db.agent_version.get_target()) is None:
                raise self.context.abort(grpc.StatusCode.NOT_FOUND, "Failed to define target agent version.")
            return version.image_id

        @cached_property
        def agent_metadata(self):
            host, port = ENV_CONFIG.SERVER_URL.split(':')
            metadata = {
                METADATA_HOST_ATTR: host,
                METADATA_PORT_ATTR: port,
                METADATA_REQUEST_FREQUENCY: ENV_CONFIG.AGENT_REQUEST_FREQUENCY,
                METADATA_AGENT_VERSION_ATTR: self._get_image,
                METADATA_LOGGING_HOST_ATTR: ENV_CONFIG.LOGGING_HOST,
                METADATA_LOGGING_PORT_ATTR: ENV_CONFIG.LOGGING_PORT,
                METADATA_OBJECT_STORAGE_URL_ATTR: ENV_CONFIG.OBJECT_STORAGE_URL,
                METADATA_LT_CREATED_ATTR: 'True'
            }
            if ssh_keys := self.request.metadata.get('ssh-keys'):
                user, ssh_pub_key = ssh_keys.split(':', 1)
                metadata[
                    "user-data"] = f"#cloud-config\nssh_pwauth: no\nusers:\n  - name: {user}\n    groups: sudo\n    shell: /bin/bash\n    " \
                                   f"sudo: ['ALL=(ALL) NOPASSWD:ALL']\n    ssh-authorized-keys:\n      - {ssh_pub_key}"
            metadata.update(self.request.metadata)
            return metadata

        def _get_preset_resources(self, tank_preset_id):
            preset = self.db.preset.get(tank_preset_id)
            return preset.memory, preset.cores, preset.disk_size

        def proceed(self):
            request = self.request
            db = self.db
            user_id = authorize(self.context, self.user_token, request.folder_id, permissions.AGENTS_CREATE,
                                self.request_id)
            description = request.description.replace('\n', ' ')

            self.logger.debug(f"Agent image {self._get_image} will be used")
            memory, cores, disk_size = self._get_preset_resources(request.preset_id)
            network_interface_specs = compute.set_ipv6_to_network_interface(
                request.network_interface_specs) if ENV_CONFIG.ENV_TYPE == EnvType.PREPROD.value else request.network_interface_specs
            compute_operation, compute_metadata = compute.create_instance(
                request.folder_id,
                request.name,
                description,
                request.labels,
                self._get_image,
                PALTFORM_ID,
                request.service_account_id,
                request.zone_id,
                self.agent_metadata,
                network_interface_specs,
                memory,
                cores,
                disk_size,
                self.user_token,
                self.request_id)

            # set data to db
            tank_id = generate_cloud_id()
            db_tank = TankTable(
                id=tank_id,
                status=None,
                folder_id=request.folder_id,
                tank_folder_id=request.folder_id,
                created_at=compute_operation.created_at.ToDatetime(),
                compute_instance_updated_at=datetime.utcnow(),
                client_updated_at=None,
                description=description,
                service_account_id=request.service_account_id,
                preset_id=request.preset_id,
                tank_version='',  # FIXME
                compute_instance_id=compute_metadata.instance_id,
                name=request.name,
                labels=json.dumps(list(request.labels)),
                agent_version=self._get_image,
            )

            db.tank.add(db_tank)
            metadata_message = tank_instance_service_pb2.CreateTankInstanceMetadata(id=db_tank.id)
            resource_metadata = create_nested_message(metadata_message)
            return self._create_operation(user_id, db_tank, compute_operation, resource_metadata, 'Create agent')

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

    class _Delete(TankHandler):
        _handler_name = 'Delete'

        def proceed(self):
            user_id = authorize(self.context, self.user_token, self.db_tank.folder_id, permissions.AGENTS_DELETE,
                                self.request_id)

            metadata_message = tank_instance_service_pb2.DeleteTankInstanceMetadata(id=self.db_tank.id)
            resource_metadata = create_nested_message(metadata_message)
            if self.db_tank.compute_instance_id:
                compute_operation = compute.delete_instance(self.db_tank.compute_instance_id, self.user_token, self.request_id)
                return self._create_operation(user_id, self.db_tank, compute_operation, resource_metadata,
                                              'Delete agent', is_deleting=True)
            else:
                self.db.tank.delete(self.db_tank)
                db_operation = OperationTable(
                    id=generate_cloud_id(),
                    description='Delete agent',
                    folder_id=self.db_tank.folder_id,
                    target_resource_id=self.db_tank.id,
                    target_resource_type=ResourceType.TANK.value,
                    resource_metadata=resource_metadata,
                    created_at=datetime.utcnow(),
                    created_by=user_id,
                    modified_at=datetime.utcnow(),
                    done=True)
                self.db.operation.add(db_operation)
                return create_operation_message(db_operation, self.logger)

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

    class _Stop(TankHandler):
        _handler_name = 'Stop'

        def proceed(self):
            self.check_for_external_agent()

            user_id = authorize(self.context, self.user_token, self.db_tank.folder_id, permissions.AGENTS_STOP,
                                self.request_id)
            compute_operation = compute.stop_instance(self.db_tank.compute_instance_id, self.user_token, self.request_id)
            metadata_message = tank_instance_service_pb2.StopTankInstanceMetadata(id=self.db_tank.id)
            resource_metadata = create_nested_message(metadata_message)
            return self._create_operation(user_id, self.db_tank, compute_operation, resource_metadata, 'Stop agent')

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

    class _Start(TankHandler):
        _handler_name = 'Start'

        def proceed(self):
            self.check_for_external_agent()

            user_id = authorize(self.context, self.user_token, self.db_tank.folder_id, permissions.AGENTS_START,
                                self.request_id)
            compute_operation = compute.start_instance(self.db_tank.compute_instance_id, self.user_token, self.request_id)
            metadata = tank_instance_service_pb2.StartTankInstanceMetadata(id=self.db_tank.id)
            resource_metadata = create_nested_message(metadata)
            return self._create_operation(user_id, self.db_tank, compute_operation, resource_metadata, 'Start agent')

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

    class _Restart(TankHandler):
        _handler_name = 'Restart'

        def proceed(self):
            self.check_for_external_agent()

            user_id = authorize(self.context, self.user_token, self.db_tank.folder_id, permissions.AGENTS_RESTART,
                                self.request_id)
            compute_operation = compute.restart_instance(self.db_tank.compute_instance_id, self.user_token, self.request_id)
            metadata = tank_instance_service_pb2.RestartTankInstanceMetadata(id=self.db_tank.id)
            resource_metadata = create_nested_message(metadata)
            return self._create_operation(user_id, self.db_tank, compute_operation, resource_metadata, 'Restart agent')

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

    class _UpgradeImage(TankHandler):
        _handler_name = 'UpgradeImage'

        @staticmethod
        def _upgrade_image(parent_logger, tank_id, image_id, operation_id, user_token, request_id):

            def _waiting_for_operation(compute_operation, db, db_operation, user_token, request_id):
                while not compute_operation.done:  # TODO add limit
                    time.sleep(WAITING_TIMEOUT)
                    compute_operation = compute.get_operation(compute_operation.id, user_token, request_id)
                if compute_operation.HasField('error'):
                    db.operation.update(db_operation, error=compute_operation.error.message, done=True)
                    raise ComputeOperationError

            logger = parent_logger.getChild('deattached_waiting_for_operation')
            logger.debug('starting to wait')
            try:
                with DB() as db:
                    db_tank = db.tank.get(tank_id)
                    db_operation = db.operation.get(operation_id)
                    logger.info(db_operation)
                    preset = db.preset.get(db_tank.preset_id)

                    # stop instance
                    compute_stop_operation = compute.stop_instance(db_tank.compute_instance_id, user_token, request_id)
                    _waiting_for_operation(compute_stop_operation, db, db_operation, user_token, request_id)

                    # update image
                    compute_update_image_operation = compute.update_image_id(db_tank.compute_instance_id, preset.disk_size, image_id, user_token, request_id)
                    _waiting_for_operation(compute_update_image_operation, db, db_operation, user_token, request_id)

                    # update metadata
                    instance = compute.get_instance(db_tank.compute_instance_id, user_token, request_id, logger, full=True)
                    current_metadata = instance.metadata
                    current_metadata["agent-version"] = image_id
                    compute_update_metadata_operation = compute.update_metadata(db_tank.compute_instance_id, current_metadata, user_token, request_id)
                    _waiting_for_operation(compute_update_metadata_operation, db, db_operation, user_token, request_id)
                    db.tank.update_agent_version(db_tank, image_id)

                    # start instance
                    compute_start_operation = compute.start_instance(db_tank.compute_instance_id, user_token, request_id)
                    _waiting_for_operation(compute_start_operation, db, db_operation, user_token, request_id)

                    db.operation.update(db_operation, done=True)
            except ComputeOperationError:
                logger.error(f'Could not update version for the agent {db_tank.id}')
            except Exception as e:
                logger.info(f"Unexpected error: {e}")
            finally:
                logger.debug('waiting completed')

        def proceed(self):
            self.check_for_external_agent()

            user_id = authorize(self.context, self.user_token, self.db_tank.folder_id, permissions.AGENTS_UPDATE,
                                self.request_id)
            target_version = self.db.agent_version.get_target()
            if self.db_tank.agent_version == target_version.image_id:
                self.context.abort(grpc.StatusCode.FAILED_PRECONDITION, f"The agent {self.db_tank.id} already has the newest version")

            metadata = tank_instance_service_pb2.UpgradeImageTankInstanceMetadata(id=self.db_tank.id)
            resource_metadata = create_nested_message(metadata)
            self.db.tank.set_update_time(self.db_tank)
            db_operation = OperationTable(
                id=generate_cloud_id(),
                folder_id=self.db_tank.folder_id,
                target_resource_id=self.db_tank.id,
                target_resource_type=ResourceType.TANK.value,
                description='Upgrade agent',
                resource_metadata=resource_metadata,
                created_by=user_id)
            self.db.operation.add(db_operation)
            self.db.commit()

            # upgrade image in background
            t = Thread(target=self._upgrade_image, args=(self.logger, self.db_tank.id, target_version.image_id, db_operation.id, self.user_token, self.request_id,))
            t.start()
            return create_operation_message(db_operation, self.logger)

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


class DbToGrpcTranslator:

    @staticmethod
    def tank(db_tank: TankTable, db_agent_version=None, lang='ru') -> tank_instance_pb2.TankInstance:
        return tank_instance_pb2.TankInstance(
            id=db_tank.id,
            folder_id=db_tank.folder_id,
            created_at=ts_from_dt(db_tank.created_at),
            compute_instance_updated_at=ts_from_dt(db_tank.compute_instance_updated_at),
            name=db_tank.name,
            description=db_tank.description,
            labels=json.loads(db_tank.labels) if db_tank.labels else None,
            service_account_id=db_tank.service_account_id,
            preset_id=db_tank.preset_id,
            tank_version=db_tank.tank_version,
            status=tank_instance_pb2.TankInstance.Status.Value(db_tank.status) if db_tank.status else None,
            errors=[db_tank.errors] if db_tank.errors else [],  # FIXME split??
            current_job=db_tank.current_job,
            compute_instance_id=db_tank.compute_instance_id,
            agent_version=get_ui_agent_version(db_tank, db_agent_version, lang),
        )


def get_ui_agent_version(db_tank: TankTable, db_agent_version: AgentVersionTable, lang='ru'):
    if not db_tank.agent_version:
        return tank_instance_pb2.AgentVersion(
            status=tank_instance_pb2.AgentVersion.UNSET,
            status_comment=STATUS_COMMENTS[lang][AgentVersionStatus.UNSET.value]
        )
    if db_agent_version is None:
        return tank_instance_pb2.AgentVersion(
            status=tank_instance_pb2.AgentVersion.UNKNOWN,
            status_comment=STATUS_COMMENTS[lang][AgentVersionStatus.UNKNOWN.value]
        )
    assert db_tank.agent_version == db_agent_version.image_id
    status = db_agent_version.status
    if status == AgentVersionStatus.TARGET.value:
        status = AgentVersionStatus.ACTUAL.value
    assert status in tank_instance_pb2.AgentVersion.VersionStatus.keys()

    return tank_instance_pb2.AgentVersion(
        status=tank_instance_pb2.AgentVersion.VersionStatus.Value(status),
        revision=db_agent_version.revision,
        id=db_agent_version.image_id,
        status_comment=STATUS_COMMENTS[lang][status]
    )
