import io
import json
import uuid
from datetime import datetime

from google.protobuf.json_format import MessageToJson, Parse
from sqlalchemy.exc import SQLAlchemyError

import grpc
from load.projects.cloud.cloud_helper import aws
from load.projects.cloud.loadtesting.config import DEFAULT_PAGE_SIZE, ENV_CONFIG
from load.projects.cloud.loadtesting.db.job import JobTable, JobStatus
from load.projects.cloud.loadtesting.db.signals import SignalType, SignalStatus
from load.projects.cloud.loadtesting.db.tables import OperationTable, ResourceType, SignalTable, Generator, \
    AgentVersionStatus, StorageTable
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 yandex.cloud.priv.loadtesting.v1 import tank_job_pb2 as messages
from yandex.cloud.priv.loadtesting.v1 import tank_job_service_pb2 as job_service
from yandex.cloud.priv.loadtesting.v1 import tank_job_service_pb2_grpc as job_grpc
from load.projects.cloud.loadtesting.server.api.private_v1.ammo import create_ammo
from .create_form import get_form
from .job_config import JobConfig
from load.projects.cloud.loadtesting.server.api.private_v1.job_report import get_test_charts
from load.projects.cloud.loadtesting.server.api.private_v1.operation import create_operation_message, create_nested_message
from load.projects.cloud.loadtesting.server.api.common.utils import generate_cloud_id, authorize, ts_from_dt, parse_target

DELETING_STATUSES = [
    JobStatus.CREATED.value,
    JobStatus.FAILED.value,
    JobStatus.FINISHED.value,
    JobStatus.STOPPED.value,
    JobStatus.AUTOSTOPPED.value,
    JobStatus.NOT_FOUND.value
]


class InvalidAmmo(Exception):
    pass


class InvalidGenerator(Exception):
    pass


class InvalidPort(Exception):
    pass


def convert_status(status):
    if status == JobStatus.FINISHED.value:
        return messages.TankJob.Status.DONE
    if status == JobStatus.POST_PROCESS.value:
        return messages.TankJob.Status.POST_PROCESSING
    if status == JobStatus.WAITING_FOR_A_COMMAND_TO_RUN.value:
        return messages.TankJob.Status.PREPARING
    return status


def create_message(db_object: JobTable, errors=None) -> messages.TankJob:
    """
    !!! Не использовать в циклах
    !!! Гораздо медленнее чем create_message_for_list, из-за неявных поздапросов.
    """
    return messages.TankJob(
        id=db_object.id,
        folder_id=db_object.folder_id,
        name=db_object.name,
        description=db_object.description,
        labels=db_object.labels,
        created_at=ts_from_dt(db_object.created_at),
        started_at=ts_from_dt(db_object.started_at),
        finished_at=ts_from_dt(db_object.finished_at),
        updated_at=ts_from_dt(db_object.updated_at),
        generator=db_object.generator or None,
        tank_instance_id=db_object.tank.id if db_object.tank else None,  # TODO: fix lazy join
        target_address=db_object.target_address,
        target_port=db_object.target_port,
        target_version=db_object.version,
        config=json.dumps(db_object.config.content) if db_object.config else None,  # TODO: fix lazy join
        ammo_id=db_object.ammos[0].id if db_object.ammos else None,  # TODO: fix lazy join
        cases=db_object.cases,
        status=convert_status(db_object.status),
        errors=(errors or []) + (db_object.errors or []),
        favorite=db_object.favorite,
        imbalance_point=db_object.imbalance_point,
        imbalance_ts=None if not db_object.imbalance_ts else int(db_object.imbalance_ts.timestamp()),
        imbalance_at=None if not db_object.imbalance_ts else ts_from_dt(db_object.imbalance_ts),
    )


def create_message_for_list(db_object: JobTable, errors=None) -> messages.TankJob:
    return messages.TankJob(
        id=db_object.id,
        folder_id=db_object.folder_id,
        name=db_object.name,
        description=db_object.description,
        labels=db_object.labels,
        created_at=ts_from_dt(db_object.created_at),
        started_at=ts_from_dt(db_object.started_at),
        finished_at=ts_from_dt(db_object.finished_at),
        updated_at=ts_from_dt(db_object.updated_at),
        generator=db_object.generator or None,
        target_address=db_object.target_address,
        target_port=db_object.target_port,
        target_version=db_object.version,
        # cases=db_object.cases,
        status=convert_status(db_object.status),
        # errors=(errors or []) + (db_object.errors or []),
        favorite=db_object.favorite,
        tank_instance_id=db_object.tank_id,
    )


def extract_job_attrs(request, db, job, uploader):
    """
    Values passed in form are more important than values passed in config,
    so they overlap the last ones.
    """

    job.name = request.name or job.config.content[uploader].get('job_name', '')
    job.description = request.description or job.config.content[uploader].get('job_dsc', '')

    if request.generator:
        generator = messages.TankJob.Generator.Name(request.generator)
        job.generator = generator
    else:
        if job.config.content.get('pandora', {}).get('enabled', False):
            job.generator = 'PANDORA'
        elif job.config.content.get('phantom', {}).get('enabled', False):
            job.generator = 'PHANTOM'
        else:
            raise InvalidGenerator

    ammo = None
    if request.ammo_id:
        ammo = db.ammo.get(request.ammo_id)
    elif request.HasField('test_data'):
        ammo = db.ammo.get_by_name(request.test_data.object_storage_filename, request.test_data.object_storage_bucket, job.folder_id)

    if ammo:
        if ammo.folder_id != job.folder_id:
            raise InvalidAmmo
        job.ammos.append(ammo)

    try:
        target_address = request.target_address
        target_port = request.target_port
        if pools := job.config.content.get('pandora', {}).get('config_content', {}).get(
                'pools', []):
            if mixed_target := pools[0].get('gun', {}).get('target', ''):
                target_address, target_port = parse_target(mixed_target)
        elif mixed_target := job.config.content.get('phantom', {}).get('address', ''):
            target_address, target_port = parse_target(mixed_target)
        job.target_address = target_address or ''
        job.target_port = target_port
    except ValueError:
        raise InvalidPort

    job.version = request.target_version or job.config.content[uploader].get('ver', '')

    return job


class JobServicer(job_grpc.TankJobServiceServicer):

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

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

        def proceed(self):
            request = self.request
            context = self.context
            db = self.db
            job = db.job.get(request.id)
            if not job:
                raise context.abort(grpc.StatusCode.NOT_FOUND, f'Job {request.id} is not found.')

            authorize(self.context, self.user_token, job.folder_id, permissions.TESTS_GET, self.request_id)

            return create_message(job)

    def Get(self, request: job_service.GetTankJobRequest, context: grpc.ServicerContext) -> messages.TankJob:
        return self._Get(self.logger).handle(request, context)

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

        def proceed(self):
            request = self.request
            db = self.db
            # TODO check every item CLOUDLOAD-95
            authorize(self.context, self.user_token, request.folder_id, permissions.TESTS_GET, self.request_id)

            pager = Pager(request.page_token,
                          request.page_size or DEFAULT_PAGE_SIZE,
                          request.filter)
            jobs = db.job.get_by_folder(request.folder_id,
                                        offset=pager.offset,
                                        limit=pager.page_size)
            pager.set_shift(len(jobs))
            jobs_list = [create_message_for_list(job) for job in jobs]

            return job_service.ListTankJobsResponse(tank_jobs=jobs_list, next_page_token=pager.next_page_token)

    def List(self, request: job_service.GetTankJobRequest, context: grpc.ServicerContext) -> messages.TankJob:
        return self._List(self.logger).handle(request, context)

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

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

            user_id = authorize(self.context, self.user_token, request.folder_id, permissions.TESTS_CREATE,
                                self.request_id)

            tank = db.tank.get(request.tank_instance_id)
            agent_version_revision = 0
            if tank.agent_version:
                if version := db.agent_version.get(tank.agent_version):
                    agent_version_revision = version.revision or 0
                    if version.status == AgentVersionStatus.OUTDATED.value:
                        raise context.abort(grpc.StatusCode.INVALID_ARGUMENT,
                                            f'The agent {tank.name} [{tank.id}] is OUTDATED, '
                                            'please update the agent to the actual version.')

            job = JobTable()
            job.id = generate_cloud_id()
            job.tank = tank
            job.logging_log_group_id = request.logging_log_group_id
            if request.folder_id != job.tank.folder_id:
                context.abort(
                    grpc.StatusCode.INVALID_ARGUMENT,
                    f'Tank {job.tank.id} does not belong to folder {request.folder_id}'
                )

            # TODO: job is saved to DB with the operation,
            #  but neither _create_job nor _create_config are executed by this time.
            create_operation = OperationTable(
                id=generate_cloud_id(),
                folder_id=request.folder_id,
                description='Create Job',
                created_by=user_id,
                target_resource_id=job.id,
                target_resource_type=ResourceType.JOB.value
            )
            metadata_message = job_service.CreateTankJobMetadata(id=job.id)
            create_operation.resource_metadata = create_nested_message(metadata_message)
            db.operation.add(create_operation)

            if not job.tank:
                error = f'Tank instance {request.tank_instance_id} is not found.'
                db.operation.update(create_operation, error=error, done=True)
                context.abort(grpc.StatusCode.NOT_FOUND, error)

            # Validation should be moved to separate async function
            self._create_ammo()
            config = self._create_config(request, agent_version_revision)
            if not config.error:
                self.logger.debug(f"Test will be run with {config.content=}")
                job = self._create_job(job, config)
                db.job.add(job)
            else:
                error = f'Config validation errors: {config.error}'
                self.logger.error(error)
                db.operation.update(create_operation, error=error, done=True)
                context.abort(grpc.StatusCode.INVALID_ARGUMENT, error)

            db.operation.update(create_operation, done=True)
            job_message = create_message(job)
            create_operation.done_resource_snapshot = MessageToJson(job_message)
            db.operation.add(create_operation)

            return create_operation_message(create_operation, self.logger)

        def _create_ammo(self):
            # TODO: разобраться, почему не работает
            # ни bool(self.request.HasField('test_data'))
            # ни bool(self.request.test_data)
            if self.request.HasField('test_data') and str(self.request.test_data):
                bucket = self.request.test_data.object_storage_bucket
                s3_file_name = self.request.test_data.object_storage_filename
                aws.check_access_to_file(bucket, s3_file_name, self.user_token)
                ammo = self.db.ammo.get_by_name(s3_file_name, bucket, self.request.folder_id)
                if ammo is None:
                    storage = self.db.storage.get_by_bucket(bucket, self.request.folder_id)
                    if storage is None:
                        storage = StorageTable(
                            id=generate_cloud_id(),
                            folder_id=self.request.folder_id,
                            bucket=bucket,
                            name=bucket,
                            description="Created with a running test",
                            created_at=datetime.utcnow()
                        )
                        self.db.storage.add(storage)
                    create_ammo(self.db, self.request.test_data.object_storage_filename, self.request.folder_id, storage)

        def _create_config(self, request, agent_version_revision=0) -> JobConfig:
            config = JobConfig()

            if request.config != '':
                config.raw_content = request.config
            else:
                config.make_from_scratch(self.db, request, agent_version_revision)
            if config.raw_content:
                config.deserialize()

            if not config.error:
                config.add_generator_debug()
                config.rewrite_ammo(self.db, request)
                config.validate()
                config.check_sanity()

            return config

        def _create_job(self, job: JobTable, config: JobConfig) -> JobTable:
            job.config = config
            job.folder_id = self.request.folder_id
            try:
                enriched_job = extract_job_attrs(self.request, self.db, job, config.uploader)
            except InvalidAmmo:
                self.context.abort(
                    grpc.StatusCode.INVALID_ARGUMENT,
                    f'Ammo {self.request.ammo_id} does not belong to folder {self.request.folder_id}'
                )
            except InvalidGenerator:
                self.context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'No valid generator found')
            except InvalidPort:
                self.context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'Target port is invalid')

            enriched_job.created_at = datetime.utcnow()
            enriched_job.run_at = datetime.utcnow()
            enriched_job.status = JobStatus.CREATED.value
            return enriched_job

    def Create(self, request: job_service.CreateTankJobRequest, context: grpc.ServicerContext):
        return self._Create(self.logger).handle(request, context)

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

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

            job = db.job.get(request.id)
            if not job:
                context.abort(grpc.StatusCode.NOT_FOUND, f'Job {request.id} is not found.')

            user_id = authorize(self.context, self.user_token, job.folder_id, permissions.TESTS_UPDATE, self.request_id)

            update_operation = OperationTable(
                id=generate_cloud_id(),
                folder_id=job.folder_id,
                description='Update Job',
                target_resource_id=job.id,
                target_resource_type=ResourceType.JOB.value,
                created_by=user_id)

            metadata_message = job_service.UpdateTankJobMetadata(id=job.id)
            update_operation.resource_metadata = create_nested_message(metadata_message)

            db.operation.add(update_operation)

            current_imbalance_ts = 0
            try:
                current_imbalance_ts = job.imbalance_ts.timestamp()
            except AttributeError:
                self.logger.debug('Imbalance ts for job %s is None', job.id)

            if request.name != '' or job.name != '':
                job.name = request.name
            if request.description != '' or job.description != '':
                job.description = request.description
            if request.favorite is True or job.favorite is True:
                job.favorite = request.favorite
            if request.target_version != '' or job.version != '':
                job.version = request.target_version
            if request.imbalance_at and request.imbalance_at.ToMicroseconds != 0 or current_imbalance_ts != 0:
                job.imbalance_ts = request.imbalance_at.ToDatetime()
            if request.imbalance_point != 0 or job.imbalance_point != 0:
                job.imbalance_point = request.imbalance_point

            try:
                db.job.add(job)
            except SQLAlchemyError:
                error = f'Failed to update job {job.id}.'
                update_operation.error = error
                db.operation.update(update_operation, error=error, done=True)
                context.abort(grpc.StatusCode.INTERNAL, error)

            db.operation.update(update_operation, done=True)
            job_message = create_message(job)
            update_operation.done_resource_snapshot = MessageToJson(job_message)

            return create_operation_message(update_operation, self.logger)

    def Update(self, request: job_service.CreateTankJobRequest, context: grpc.ServicerContext):
        return self._Update(self.logger).handle(request, context)

    class _GetGenerators(handler.BasePrivateHandler):
        _handler_name = "GetGenerators"

        def proceed(self):
            request = self.request

            authorize(self.context, self.user_token, request.folder_id, permissions.TESTS_GET, self.request_id)
            generators = sorted(
                [s.value for s in Generator if s.value != 'GENERATOR_UNSPECIFIED'],
                reverse=False
            )
            return job_service.GetGeneratorsResponse(
                generator_name=generators
            )

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

    class _GetCreateForm(handler.BasePrivateHandler):
        _handler_name = 'GetCreateForm'

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

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

            if not request.generator:
                context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'Generator not specified')
            generator = messages.TankJob.Generator.Name(request.generator)
            create_form = get_form(generator, self.lang)
            create_form['generator_type'] = generator.lower()

            form_message = job_service.FormResponse()
            Parse(json.dumps(create_form), form_message)

            return form_message

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

    class _UploadAmmo(handler.BasePrivateHandler):
        _handler_name = 'UploadAmmo'

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

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

            if not request.ammo:
                context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'Ammo file is empty')

            # валидность имён: https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
            s3_filename = f'ammo_{uuid.uuid4()}'
            self.logger.debug(f'ammo with original name "{request.filename}" will be stored as "{s3_filename}"')

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

            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 ammo {s3_filename}')

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

            return job_service.UploadFileResponse(
                id=ammo.id,
                status=job_service.UploadFileResponse.OK,
                filename=ammo.s3_name
            )

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

    class _ValidateConfig(handler.BasePrivateHandler):
        _handler_name = 'ValidateConfig'

        def proceed(self):
            request = self.request

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

            config = JobConfig()
            config.raw_content = request.config
            config.deserialize()
            if not config.error:
                config.validate()

            if config.error:
                # TODO: check that Console shows errors with status=FAILED
                return job_service.ValidateConfigResponse(
                    status=job_service.ValidateConfigResponse.Status.FAILED,
                    errors=[json.dumps(config.error)]
                )

            return job_service.ValidateConfigResponse(
                status=job_service.ValidateConfigResponse.Status.OK
            )

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

    class _GetConfig(handler.BasePrivateHandler):
        _handler_name = 'GetConfig'

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

            job = db.job.get(request.job_id)
            if not job:
                context.abort(grpc.StatusCode.NOT_FOUND, f'Job {request.job_id} is not found.')
            authorize(self.context, self.user_token, job.folder_id, permissions.TESTS_GET, self.request_id)
            config = db.config.get_by_job_id(job.id)
            if not job:
                context.abort(grpc.StatusCode.NOT_FOUND, f'Config for job {request.job_id} is not found.')

            return job_service.GetConfigResponse(
                config=json.dumps(config.content).encode('utf-8')
            )

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

    class _GetFolderStats(handler.BasePrivateHandler):
        _handler_name = 'GetFolderStats'

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

            # TODO realize about authorization of this request
            try:
                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)
            except Exception:
                if self.context._state.code == grpc.StatusCode.PERMISSION_DENIED:
                    jobs_count, tanks_count = 0, 0
                else:
                    raise
            else:
                jobs_count = db.job.count_for_folder(request.folder_id)
                tanks_count = db.tank.count_for_folder(request.folder_id)
            return job_service.FolderStats(
                tanks_count=tanks_count,
                jobs_count=jobs_count,
            )

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

    class _Stop(handler.BasePrivateHandler):
        _handler_name = 'Stop'

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

            job = db.job.get(request.id)
            if not job:
                context.abort(grpc.StatusCode.NOT_FOUND, f'Job {request.id} is not found.')
            user_id = authorize(self.context, self.user_token, job.folder_id, permissions.TESTS_STOP, self.request_id)
            if job.status not in [
                JobStatus.CREATED.value, JobStatus.INITIATED.value,
                JobStatus.PREPARING.value, JobStatus.RUNNING.value
            ]:
                context.abort(grpc.StatusCode.FAILED_PRECONDITION, f'Can\'t stop job in status {job.status}.')

            stop_operation = OperationTable(
                id=generate_cloud_id(), folder_id=job.folder_id,
                description='Stop Job',
                target_resource_type=ResourceType.JOB.value,
                target_resource_id=job.id,
                created_by=user_id
            )
            metadata_message = job_service.StopTankJobMetadata(id=job.id)
            stop_operation.resource_metadata = create_nested_message(metadata_message)
            db.operation.add(stop_operation)

            job.status = JobStatus.STOPPING.value
            stop_signal = SignalTable(job_id=job.id,
                                      type=SignalType.STOP.value,
                                      status=SignalStatus.WAITING.value,
                                      operation_id=stop_operation.id)
            db.job.add(job)
            db.signal.add(stop_signal)

            return create_operation_message(stop_operation, self.logger)

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

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

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

            job = db.job.get(request.id)
            if not job:
                context.abort(grpc.StatusCode.NOT_FOUND, f'Job {request.id} is not found.')
            if job.status not in DELETING_STATUSES:
                context.abort(grpc.StatusCode.FAILED_PRECONDITION,
                              f'The job {job.id} has an invalid state "{job.status}" for this operation')
            user_id = authorize(self.context, self.user_token, job.folder_id, permissions.TESTS_DELETE, self.request_id)

            delete_operation = OperationTable(
                id=generate_cloud_id(),
                folder_id=job.folder_id,
                description='Delete Job',
                created_by=user_id,
                target_resource_type=ResourceType.JOB.value,
                target_resource_id=job.id
            )
            metadata_message = job_service.DeleteTankJobMetadata(id=job.id)
            delete_operation.resource_metadata = create_nested_message(metadata_message)
            db.operation.add(delete_operation)

            try:
                db.job.delete(job)
                error = None
            except SQLAlchemyError:
                error = f'Failed to update job {job.id}.'
                context.abort(grpc.StatusCode.INTERNAL, error)
            finally:
                db.operation.update(delete_operation, error=error, done=True)

            return create_operation_message(delete_operation, self.logger)

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

    class _GetReport(handler.BasePrivateHandler):
        _handler_name = 'GetReport'

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

            job = db.job.get(request.job_id)
            if not job:
                context.abort(grpc.StatusCode.NOT_FOUND, f'Job {request.job_id} is not found.')
            authorize(self.context, self.user_token, job.folder_id, permissions.TESTS_GET, self.request_id)

            charts = []
            if job.started_at:
                charts = get_test_charts(job, self.lang)

            return messages.TankReport(
                job_id=job.id,
                imbalance_point=job.imbalance_point,
                imbalance_ts=int(job.imbalance_ts.timestamp()) if job.imbalance_ts else None,  # this field is int, not Timestamp
                imbalance_at=ts_from_dt(job.imbalance_ts),
                charts=charts,
                finished=bool(job.finished_at)
            )

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