import random
from typing import Optional

from pyscopg2.balancer_policy.base import BaseBalancerPolicy


class GreedyBalancerPolicy(BaseBalancerPolicy):
    async def _get_pool(
            self,
            read_only: bool,
            fallback_master: Optional[bool] = None,
            choose_master_as_replica: bool = False
    ):
        candidates = []
        if read_only:
            candidates.extend(
                await self._pool_manager.get_replica_pools(
                    fallback_master=fallback_master,
                ),
            )
        if (
                not read_only or
                (
                    choose_master_as_replica and
                    self._pool_manager.master_pool_count > 0
                )
        ):
            candidates.extend(await self._pool_manager.get_master_pools())

        fat_pool = max(candidates, key=self._pool_manager.get_pool_freesize)
        max_freesize = self._pool_manager.get_pool_freesize(fat_pool)
        return random.choice([
            candidate
            for candidate in candidates
            if self._pool_manager.get_pool_freesize(candidate) == max_freesize
        ])


__all__ = ["GreedyBalancerPolicy"]
