from functools import cached_property
from typing import List, Tuple

from sqlalchemy.orm import Session

from load.projects.cloud.loadtesting.db.ammo import AmmoQueries
from load.projects.cloud.loadtesting.db.connection import get_pg_session
from load.projects.cloud.loadtesting.db.job import JobQueries
from load.projects.cloud.loadtesting.db.job_config import JobConfigQueries
from load.projects.cloud.loadtesting.db.operation import OperationQueries
from load.projects.cloud.loadtesting.db.preset import PresetQueries
from load.projects.cloud.loadtesting.db.signals import SignalQueries
from load.projects.cloud.loadtesting.db.tank import TankQueries
from load.projects.cloud.loadtesting.db.agent_version import AgentVersionQueries
from load.projects.cloud.loadtesting.db.storage import StorageQueries

from load.projects.cloud.loadtesting.db.tables import TankTable, AgentVersionTable


class DB:
    def __init__(self):
        self._session: Session = None

    def __enter__(self):
        self._session = get_pg_session().__enter__()
        self._transaction = self._session.begin()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._transaction.__exit__(exc_type, exc_val, exc_tb)
        self._session.__exit__(exc_type, exc_val, exc_tb)

    def commit(self):
        self._session.commit()

    @cached_property
    def tank(self) -> TankQueries:
        return TankQueries(self._session)

    @cached_property
    def job(self) -> JobQueries:
        return JobQueries(self._session)

    @cached_property
    def config(self) -> JobConfigQueries:
        return JobConfigQueries(self._session)

    @cached_property
    def ammo(self) -> AmmoQueries:
        return AmmoQueries(self._session)

    @cached_property
    def signal(self) -> SignalQueries:
        return SignalQueries(self._session)

    @cached_property
    def operation(self) -> OperationQueries:
        return OperationQueries(self._session)

    @cached_property
    def preset(self) -> PresetQueries:
        return PresetQueries(self._session)

    @cached_property
    def agent_version(self) -> AgentVersionQueries:
        return AgentVersionQueries(self._session)

    @cached_property
    def storage(self) -> StorageQueries:
        return StorageQueries(self._session)

    def tanks_with_version_by_folder(self, folder_id, limit, offset) -> List[Tuple[TankTable, AgentVersionTable]]:
        query = self._session.query(TankTable, AgentVersionTable)
        query = query.filter(TankTable.folder_id == folder_id)
        query = query.join(AgentVersionTable, TankTable.agent_version == AgentVersionTable.image_id, isouter=True)

        query = query.order_by(TankTable.created_at.desc())
        query = query.limit(limit).offset(offset)

        return query.all()
