from sqlalchemy import select

from app.models.server import Server, ServerProtocol
from app.repositories.base import BaseRepository


class ServerRepository(BaseRepository[Server]):
    def __init__(self, session):
        super().__init__(session, Server)

    async def get_active(self) -> list[Server]:
        stmt = select(Server).where(Server.is_active.is_(True))
        result = await self.session.execute(stmt)
        return list(result.scalars().all())

    async def get_by_protocol(self, protocol: ServerProtocol) -> list[Server]:
        stmt = (
            select(Server)
            .where(Server.protocol == protocol)
            .where(Server.is_active.is_(True))
        )
        result = await self.session.execute(stmt)
        return list(result.scalars().all())

    async def get_by_location(self, country_code: str) -> list[Server]:
        stmt = (
            select(Server)
            .where(Server.country_code == country_code.upper())
            .where(Server.is_active.is_(True))
        )
        result = await self.session.execute(stmt)
        return list(result.scalars().all())

    async def get_by_name(self, name: str) -> Server | None:
        stmt = select(Server).where(Server.name == name)
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def get_least_loaded(self, protocol: ServerProtocol) -> Server | None:
        stmt = (
            select(Server)
            .where(Server.protocol == protocol)
            .where(Server.is_active.is_(True))
            .order_by(Server.load_percent.asc())
            .limit(1)
        )
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def count_active(self) -> int:
        stmt = select(Server).where(Server.is_active.is_(True))
        result = await self.session.execute(stmt)
        return len(result.scalars().all())