from datetime import datetime

from load.projects.cloud.loadtesting.db.tables import TankTable, TankStatus
from load.projects.cloud.loadtesting.logan import lookup_logger


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

    def add(self, tank_instance):
        self._session.add(tank_instance)
        lookup_logger().debug(f"Insert tank data: {tank_instance}")

    def get(self, tank_id) -> TankTable:
        result = self._session.query(TankTable).filter(TankTable.id == tank_id).first()
        lookup_logger().debug(f"Get tank from PG: {result}")
        return result

    def get_by_compute_instance_id(self, compute_instance_id) -> TankTable:
        result = self._session.query(TankTable).filter(
            TankTable.compute_instance_id == compute_instance_id
        ).first()
        lookup_logger().debug(f"Get tank from PG: {result}")
        return result

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

    def delete(self, tank):
        self._session.delete(tank)
        lookup_logger().debug(f"Delete tank: {tank.id}")

    def set_update_time(self, tank):
        tank.compute_instance_updated_at = datetime.utcnow()
        lookup_logger().debug(
            f"Set updated compute instance time {tank.compute_instance_updated_at} for tank {tank.id}")

    def update_status(self, tank: TankTable, status: TankStatus):
        tank.client_updated_at = datetime.utcnow()
        tank.client_status = status.value
        lookup_logger().debug(f"Updated tank {status=}")

    def update_agent_version(self, tank: TankTable, agent_version):
        # agent_version is image_id
        tank.agent_version = agent_version
