import copy
from typing import Dict

from staff.headcounts.headcounts_summary.query_builder import Result


class HierarchySummaryCounter:
    def summarize(self, result: Result) -> None:
        for next_level_grouping in result.next_level_grouping.values():
            self.summarize(next_level_grouping)

        for child in result.children.values():
            self.summarize(child)

        for child in result.children.values():
            self._merge_next_level_grouping(result, child)

        if result.next_level_grouping:
            for next_level_grouping in result.next_level_grouping.values():
                result.summary = result.summary.plus(next_level_grouping.summary)
        else:
            for child in result.children.values():
                result.summary = result.summary.plus(child.summary)

    def _merge_dict(self, to_dict: Dict[int, Result], from_dict: Dict[int, Result]) -> None:
        for dep_id, result in from_dict.items():
            if dep_id in to_dict:
                merge_to = to_dict[dep_id]
                self._merge_tree(merge_to, result)
                merge_to.summary = merge_to.summary.plus(result.summary)
            else:
                to_dict[dep_id] = copy.deepcopy(result)

    def _merge_next_level_grouping(self, to_result: Result, from_result: Result) -> None:
        assert from_result.id in to_result.children
        self._merge_dict(to_result.next_level_grouping, from_result.next_level_grouping)

    def _merge_tree(self, to_result: Result, from_result: Result) -> None:
        assert to_result.id == from_result.id
        self._merge_dict(to_result.next_level_grouping, from_result.next_level_grouping)
        self._merge_dict(to_result.children, from_result.children)
