import random
from collections.abc import Callable
from concurrent.futures import ProcessPoolExecutor

from .base import BatchMatrixData, TopologySkeleton


class TopologyPool:

    def __init__(self, topologies: list[TopologySkeleton]):
        if not topologies:
            raise ValueError("TopologyPool requires at least one topology")
        self.topologies = topologies

    def sample(self) -> TopologySkeleton:
        return random.choice(self.topologies)

    @staticmethod
    def build(
        generate_fn: Callable[..., TopologySkeleton],
        pool_size: int,
        num_workers: int = 1,
        **kwargs,
    ) -> "TopologyPool":
        if num_workers <= 1:
            topologies = [generate_fn(**kwargs) for _ in range(pool_size)]
        else:
            with ProcessPoolExecutor(max_workers=num_workers) as executor:
                futures = [executor.submit(generate_fn, **kwargs) for _ in range(pool_size)]
                topologies = [f.result() for f in futures]
        return TopologyPool(topologies)


class PooledGenerator:

    def __init__(
        self,
        pool: TopologyPool,
        fill_fn: Callable[[TopologySkeleton, int], BatchMatrixData],
    ):
        self.pool = pool
        self.fill_fn = fill_fn

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        assert batch_size == 1, "PooledGenerator only supports batch_size=1"
        skeleton = self.pool.sample()
        return self.fill_fn(skeleton, num_context_pairs)


class MultiSizePooledGenerator:

    def __init__(self, pools: dict[int, TopologyPool],
                 fill_fn: Callable[[TopologySkeleton, int], BatchMatrixData]):
        self.pools = pools
        self.sizes = list(pools.keys())
        self.fill_fn = fill_fn

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        assert batch_size == 1, "MultiSizePooledGenerator only supports batch_size=1"
        n = random.choice(self.sizes)
        skeleton = self.pools[n].sample()
        return self.fill_fn(skeleton, num_context_pairs)
