from datetime import datetime
from typing import List

from google.protobuf.json_format import MessageToJson
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Query

from load.projects.cloud.loadtesting.db.tables import JobTable, JobStatus, SignalTable, SignalStatus, \
    SignalType, OperationTable
from load.projects.cloud.loadtesting.logan import lookup_logger
from yandex.cloud.priv.loadtesting.v1 import tank_job_pb2


class JobQueries:
    def __init__(self, session):
        self._session = session

    def get(self, job_id) -> JobTable:
        return self._session.query(JobTable).filter(JobTable.id == job_id).first()

    def filter(self, *filters) -> List[JobTable]:
        return self._session.query(JobTable).filter(*filters).all()

    def get_waiting_for_tank(self, tank_id):
        return self._session.query(JobTable).filter(
            JobTable.tank_id == tank_id,
            JobTable.status == JobStatus.CREATED.value,
        ).order_by(JobTable.created_at).first()

    def get_by_folder(self, folder_id, offset=None, limit=None):
        query: Query = self._session.query(JobTable).filter(
            JobTable.folder_id == folder_id).order_by(JobTable.created_at.desc())
        if offset is not None:
            query = query.offset(offset)
        if limit is not None:
            query = query.limit(limit)
        return query.all()

    def count_for_folder(self, folder_id):
        return self._session.query(JobTable.id).filter(JobTable.folder_id == folder_id).count()

    def update_status(self, job: JobTable, status: JobStatus):
        job.status = status.value
        job.status_updated_at = datetime.utcnow()
        lookup_logger().debug(f"Updated job {status=}")

    def append_error(self, job: JobTable, error: str):
        if not job.errors:
            job.errors = [error]
        elif error not in job.errors:
            job.errors = job.errors + [error]

    def add(self, job: JobTable):
        try:
            self._session.add(job)
            lookup_logger().debug(f'Added job {job.id}')
        except SQLAlchemyError:
            lookup_logger().exception(f'Exception while updating job {job.id}.')
            raise

    def delete(self, job: JobTable):
        try:
            self._session.delete(job)
            lookup_logger().debug(f"Delete job: {job.id}")
        except SQLAlchemyError:
            lookup_logger().exception(f'Exception while deleting job {job.id}.')
            raise

    def close_pending_signals(self, job: JobTable, job_message: tank_job_pb2.TankJob, error: str = ""):
        for signal in self._session.query(SignalTable).filter(
                SignalTable.job_id == job.id,
                SignalTable.status != SignalStatus.DONE.value
        ).all():
            operation = self._session.query(OperationTable).filter(OperationTable.id == signal.operation_id).first()
            operation.done = True
            operation.done_resource_snapshot = MessageToJson(job_message)

            if signal.type == SignalType.STOP.value:
                if error:
                    operation.error = error
            else:
                lookup_logger().warn(f"Job {job.id} finished before operation {operation.id} proceeded.")
                operation.error = "Job finished before operation proceeded."
            self._session.delete(signal)
