import math
import random
from typing import Optional, List, Iterator

from ielts_be.configs.constants import EducationalContent


class DifficultyHelper:

    def __init__(self, difficulties: Optional[List[str]]):
        self.difficulties = difficulties
        self.distributed: Optional[Iterator[str]] = None

    def distribute_for_count(self, count: int) -> None:
        if not self.difficulties or count == 0:
            return

        result = []
        remaining = count
        difficulties_count = len(self.difficulties)

        for i, diff in enumerate(self.difficulties):
            if i == difficulties_count - 1:
                slots = remaining
            else:
                slots = math.ceil(remaining / (difficulties_count - i))

            result.extend([diff] * slots)
            remaining -= slots

        self.distributed = iter(result)

    def pick_difficulty(self, difficulty: Optional[str]) -> str:
        if difficulty:
            return difficulty if difficulty != "Random" else random.choice(EducationalContent.DIFFICULTIES)

        if self.distributed:
            return next(self.distributed)

        return random.choice(EducationalContent.DIFFICULTIES)