import json
from datetime import datetime
from functools import cached_property

import grpc
import yandex.cloud.priv.loadtesting.agent.v1.job_service_pb2 as messages
import yandex.cloud.priv.loadtesting.agent.v1.job_service_pb2_grpc as job_grpc

from load.projects.cloud.cloud_helper import aws
from load.projects.cloud.loadtesting.config import ENV_CONFIG
from load.projects.cloud.loadtesting.db.tables import SignalType, TankTable
from load.projects.cloud.loadtesting.db.tables.agent_version import Status as AgentVersionStatus
from load.projects.cloud.loadtesting.db.tables.job import Status, JobTable
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.api.common.utils import authorize
from load.projects.cloud.loadtesting.server.api.private_v1.job import create_message

CLOSING_STATUSES = [
    messages.ClaimJobStatusRequest.NOT_FOUND,
    messages.ClaimJobStatusRequest.FINISHED,
    messages.ClaimJobStatusRequest.FAILED,
    messages.ClaimJobStatusRequest.STOPPED,
    messages.ClaimJobStatusRequest.AUTOSTOPPED
]


class Jobber(job_grpc.JobServiceServicer):
    def __init__(self):
        self.logger = lookup_logger('JobPublic')

    def Get(self, get_job: messages.GetJobRequest, context: grpc.ServicerContext):
        return _GetJob(self.logger).handle(get_job, context)

    class _ClaimStatus(handler.BasePublicHandler):
        _handler_name = 'ClaimStatus'

        def _utcnow(self):
            return datetime.utcnow()

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

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

            tank = db.tank.get(job.tank_id)

            authorize(self.context, self.user_token, tank.folder_id, permissions.TESTS_RUN, self.request_id)

            tank.client_updated_at = self._utcnow()

            if not tank.current_job and job.status == Status.CREATED.value:  # it means, it is the very first status claiming
                tank.current_job = job.id
                job.started_at = self._utcnow()  # add new field or take from metadata
            elif not tank.current_job and job.status == Status.WAITING_FOR_A_COMMAND_TO_RUN:
                tank.current_job = job.id
            elif tank.current_job != claim_status.job_id:
                if tank.current_job:
                    message = f'Job {tank.current_job} is in progress.'
                else:
                    message = f'No active jobs for the agent {tank.current_job} {claim_status.job_id}'
                raise context.abort(grpc.StatusCode.FAILED_PRECONDITION, message)

            db.job.update_status(job, Status(
                messages.ClaimJobStatusRequest.JobStatus.Name(claim_status.status)
            ))
            if claim_status.error:
                if ('error-type', 'internal') in context.invocation_metadata():
                    job.internal_tank_error = claim_status.error
                    db.job.add(job)
                    db.job.append_error(job, 'Internal agent error')
                else:
                    db.job.append_error(job, claim_status.error)

            if claim_status.status in CLOSING_STATUSES:
                db.job.close_pending_signals(
                    job,
                    create_message(job),
                    error=claim_status.error
                )
                job.finished_at = self._utcnow()
                tank.current_job = None

            return messages.ClaimJobStatusResponse(code=0)

    def ClaimStatus(self, claim_status: messages.ClaimJobStatusRequest, context: grpc.ServicerContext):
        return self._ClaimStatus(self.logger).handle(claim_status, context)

    class _GetSignal(handler.BasePublicHandler):
        _handler_name = 'GetSignal'

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

            if (job := db.job.get(get_signal.job_id)) is None:
                return context.abort(grpc.StatusCode.NOT_FOUND, f'Job {get_signal.job_id} is not found')

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

            with db.signal.send_to_tank(job, SignalType.STOP) as signal:
                if signal:
                    return messages.JobSignalResponse(
                        signal=messages.JobSignalResponse.STOP
                    )
            if job.status == Status.WAITING_FOR_A_COMMAND_TO_RUN.value:
                if job.run_at:
                    return messages.JobSignalResponse(
                        signal=messages.JobSignalResponse.RUN_IN,
                        run_in=(job.run_at - datetime.utcnow()).total_seconds(),
                    )
                else:
                    return messages.JobSignalResponse(
                        signal=messages.JobSignalResponse.WAIT
                    )

            return messages.JobSignalResponse(
                signal=messages.JobSignalResponse.SIGNAL_UNSPECIFIED
            )

    def GetSignal(self, get_signal: messages.JobSignalRequest, context):
        return self._GetSignal(self.logger).handle(get_signal, context)


class _GetJob(handler.BasePublicHandler):
    _handler_name = 'Get'

    def proceed(self):
        authorize(self.context, self.user_token, self.tank.folder_id, permissions.TESTS_RUN, self.request_id)

        if agent_version := self.db.agent_version.get(self.tank.agent_version):
            if agent_version.status == AgentVersionStatus.OUTDATED.value:
                error_message = f'The agent {self.tank.name} [{self.tank.id}] is OUTDATED.'
                if self.job:
                    self.db.job.update_status(self.job, Status.FAILED)
                    self.db.job.append_error(self.job, error_message)
                raise self.context.abort(grpc.StatusCode.INVALID_ARGUMENT, error_message)

        if self.job is None:
            return messages.Job()

        config = self.db.config.get(self.job.config.id)

        response = messages.Job(
            id=self.job.id,
            config=json.dumps(config.content).encode('utf-8'),
            logging_log_group_id=self.job.logging_log_group_id
        )
        if not self.ammo:
            return response

        if self.ammo.storage:
            response.MergeFrom(
                messages.Job(test_data=self.ammo_test_data)
            )
        else:
            response.MergeFrom(
                messages.Job(ammo=self.ammo_file)
            )
        return response

    @cached_property
    def tank(self) -> TankTable:
        tank_ = None
        # TODO: del with old agents CLOUDLOAD-328
        if self.request.compute_instance_id:
            tank_ = self.db.tank.get_by_compute_instance_id(self.request.compute_instance_id)
            if not tank_:
                raise self.context.abort(grpc.StatusCode.NOT_FOUND,
                                         f'Compute instance {self.request.compute_instance_id} is not found.')
        elif self.request.agent_instance_id:
            tank_ = self.db.tank.get(self.request.agent_instance_id)
            if not tank_:
                raise self.context.abort(grpc.StatusCode.NOT_FOUND,
                                         f'Agent {self.request.agent_instance_id} is not found.')
        return tank_

    @cached_property
    def job(self) -> JobTable:
        if job_id := self.tank.current_job:
            self.tank.current_job = None
            job = self.db.job.get(job_id)
            if job:
                self.db.job.update_status(job, Status.FAILED)
                self.db.job.append_error(job, 'The job was not finished')
            self.logger.error(
                f"There was {job_id=} as current_job of the agent={self.tank.id}. But the agent wants a new one")

        if not (waiting_job := self.db.job.get_waiting_for_tank(self.tank.id)):
            self.logger.info(f'No job for the agent {self.tank.id}')
            return
        return waiting_job

    @cached_property
    def ammo(self):
        if not self.job.ammos:
            return

        return self.db.ammo.get(self.job.ammos[0].id)

    @cached_property
    def ammo_test_data(self):
        return messages.StorageObject(
            object_storage_bucket=self.ammo.storage.bucket,
            object_storage_filename=self.ammo.s3_name
        )

    @cached_property
    def ammo_file(self):
        if (buffer := aws.download_file_to_buffer(s3_name=self.ammo.s3_name, bucket=ENV_CONFIG.OBJECT_STORAGE_DEFAULT_BUCKET)) is None:
            self.db.job.update_status(self.job, Status.FAILED)
            self.db.job.append_error(self.job, 'Failed to download ammo')
            raise self.context.abort(grpc.StatusCode.INTERNAL,
                                     f'Ammo {self.ammo.id} for job {self.job.id} is unavailable.')

        return messages.File(
            name=self.ammo.s3_name,
            content=buffer.getvalue(),
        )
