from typing import Dict, List

from staff.headcounts.headcounts_summary.query_builder.query_params import RelatedEntity


class RelatedEntityChains:
    def __init__(self) -> None:
        self._chains: Dict[RelatedEntity, Dict[int, List[int]]] = {}

    def get_chain(self, grouping: RelatedEntity, entity_id: int) -> List[int]:
        return self._chains[grouping][entity_id]

    def set_chains(self, grouping: RelatedEntity, chains: Dict[int, List[int]]) -> None:
        self._chains[grouping] = chains
