from utils.db_connectors.base_connector import BaseConnector
from utils.db_objects.direct_delivery import DeliveryBatch, \
    DeliveryBatchDetail, EncodingQueueDetail


class DeliveryBatchConnector(BaseConnector):

    def get_batch_id(self, dms_id):
        session = self.Session()
        query = session.query(
            DeliveryBatch.delivery_batch_id
        ).filter(DeliveryBatch.dms_master_master_id == dms_id).limit(1)
        results = query.all()
        session.close()
        delivery_batch_id = results[0][0]
        self.log.info(f'Returning batch id = {delivery_batch_id}')
        return delivery_batch_id

    def get_batch_encoding_queue_detail_records(self, batch_id):
        session = self.Session()
        query = session.query(
            EncodingQueueDetail.encoding_queue_detail_id).join(
            DeliveryBatchDetail,
            DeliveryBatchDetail.encoding_queue_detail_id ==
            EncodingQueueDetail.encoding_queue_detail_id).filter(
            DeliveryBatchDetail.delivery_batch_id == batch_id)
        results = query.all()
        session.close()
        return results

    def update_batch_status(self, batch_id, status):
        self.log.info(f'Updating batch id={batch_id} with status {status}')
        session = self.Session()
        session.query(DeliveryBatch).filter(
            DeliveryBatch.delivery_batch_id == batch_id
        ).update(
            {"status": status}
        )
        session.commit()
        session.close()

    def assert_status(self, batch_id, status):
        row = self.get_row_by_id(DeliveryBatch, 'delivery_batch_id', batch_id)
        assert row.status == status, 'Batch status was {}'.format(row.status)

    def get_batch_detail_rows(self, eqd_id):
        session = self.Session()
        results = session.query(DeliveryBatchDetail).filter_by(
            encoding_queue_detail_id=eqd_id).all()
        session.close()
        return results

    def assert_batch_created(self, eqd_id):
        batch_detail_row = self.get_batch_detail_rows(eqd_id)[0]
        batch_id = batch_detail_row.delivery_batch_id
        session = self.Session()
        results = session.query(DeliveryBatch).filter_by(
            delivery_batch_id=batch_id).all()
        assert len(results) == 1, 'Expected to find delivery batch row' \
                                  ' for id={}'.format(batch_id)

    def delete_batch_if_exists(self, eqd_id):
        session = self.Session()
        results = self.get_batch_detail_rows(eqd_id)
        if results:
            batch_detail_row = results[0]
            batch_id = batch_detail_row.delivery_batch_id
            session.delete(batch_detail_row)
            session.commit()
            session.query(DeliveryBatch).filter(
                DeliveryBatch.delivery_batch_id == batch_id).delete()
            session.commit()
        else:
            self.log.info(f'No batch exists for encoding_queue_detail_id = '
                          f'{eqd_id} no need to delete any records')
        session.close()