from contextlib import contextmanager
from datetime import datetime

from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Session

from load.projects.cloud.loadtesting.db.tables import JobTable
from load.projects.cloud.loadtesting.db.tables import SignalTable, SignalStatus, SignalType
from load.projects.cloud.loadtesting.logan import lookup_logger


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

    def add(self, signal: SignalTable):
        try:
            signal.created_at = signal.created_at or datetime.utcnow()  # TODO: timezones
            self._session.add(signal)
            lookup_logger().debug(f'Added signal for job {signal.job_id}')
        except SQLAlchemyError:
            lookup_logger().exception(f'Exception while adding signal for job {signal.job_id}.')
            raise

    @contextmanager
    def send_to_tank(self, job: JobTable, signal_type: SignalType):
        signal: SignalTable = self._session.query(SignalTable).filter(
            SignalTable.job_id == job.id,
            SignalTable.status == SignalStatus.WAITING.value,
            SignalTable.type == signal_type.value,
        ).first()

        yield signal

        if signal is not None:
            signal.status = SignalStatus.EXECUTING.value
