from collections import defaultdict


class DigarIdAssigner:
    """
    Assign DIGAR IDs to sections based on the following rules
        - TITLE_SECTION: always assigned ID "1".
        - ARTICLE: each article receives the next sequential integer ID.
        - ILLUSTRATION: always a child of another section; ID = "<parent_int>.<counter>",
            where the counter restarts at 1 for each parent.
            If the parent integer cannot be determined,
            the illustration receives its own sequential integer instead.
        - All other types (ADVERTISEMENT, OBITUARY, MISCELLANEOUS, etc.):
            objects on the same page share an integer prefix,
            with a sequential suffix per page group (e.g. "5.1", "5.2").
            If a texts[] entry already exists for said object, the DIGAR ID from
            texts[] is used.
            Sections are sorted by (start_page, sequence_number) before ID assignment.

    Parameters
    ----------
    data : dict
        Output of extract_all_data.

    Returns
    -------
    dict
        The same dict with DIGAR IDs assigned.
    """

    def __init__(self) -> None:
        """
        Initialize assignment state.
        """
        self.counter = 1
        self.illustration_counters = defaultdict(int)
        self.ad_page_integers = {}
        self.ad_page_item_counters = defaultdict(int)
        self.section_digar_ints = {}
        self.text_image_connection_lookup = {}

    def assign(self, data: dict) -> dict:
        """
        Assign DIGAR IDs to input data.

        Parameters
        ----------
        data : dict

        Returns
        -------
        dict
        """
        texts = data.get("texts", [])
        images = data.get("images", [])

        texts = self._sort_texts(texts)

        self._assign_texts(texts)
        self._assign_images(images)
        self._cleanup_intermediates(texts)

        return data

    def _sort_texts(self, texts: list[dict]) -> list[dict]:
        """
        Sort texts deterministically.
        """
        return sorted(
            texts,
            key=lambda x: (x.get("start_page", 0), int(x.get("sequence_nr", 0))),
        )

    def _assign_texts(self, texts: list[dict]) -> None:
        """
        Assign IDs to text sections.
        """
        for section in texts:
            digar_id = self._resolve_text_id(section)
            section["digar_id"] = digar_id

            section_id = section.get("section_id")
            if section_id and digar_id.isdigit():
                self.section_digar_ints[section_id] = int(digar_id)

            text_image_connection_id = section.get("text_image_connection_id")
            if text_image_connection_id:
                self.text_image_connection_lookup[text_image_connection_id] = digar_id

    def _resolve_text_id(self, section: dict) -> str:
        """
        Resolve text DIGAR ID.
        """
        t = (section.get("section_type") or "").lower()
        page = section.get("start_page", 1)

        if t == "title_section":
            return "1"

        if t == "article":
            self.counter += 1
            return str(self.counter)

        return self._page_group_id(page)

    def _assign_images(self, images: list[dict]) -> None:
        """
        Assign IDs to images.
        """
        for image in images:
            image["digar_id"] = self._resolve_image_id(image)

    def _resolve_image_id(self, image: dict) -> str:
        """
        Resolve image DIGAR ID.
        """
        text_image_connection_id = image.get("image_id")

        # priority 1: match text
        if text_image_connection_id in self.text_image_connection_lookup:
            return self.text_image_connection_lookup[text_image_connection_id]

        # priority 2: article
        article_id = image.get("article")
        if article_id in self.section_digar_ints:
            return self._article_image_id(article_id)

        # fallback
        page = image.get("page", 1)
        return self._page_group_id(page)

    def _article_image_id(self, article_id: str) -> str:
        """
        Assign image ID under article.
        """
        article_int = self.section_digar_ints[article_id]
        self.illustration_counters[article_int] += 1
        return f"{article_int}.{self.illustration_counters[article_int]}"

    def _page_group_id(self, page: int) -> str:
        """
        Assign page-grouped ID.
        """
        if page not in self.ad_page_integers:
            self.counter += 1
            self.ad_page_integers[page] = self.counter

        self.ad_page_item_counters[page] += 1

        return (
            f"{self.ad_page_integers[page]}."
            f"{self.ad_page_item_counters[page]}"
        )

    def _cleanup_intermediates(self, texts: list[dict]) -> None:
        """
        Remove intermediate fields from text objects.

        Parameters
        ----------
        texts : list[dict]
            List of text objects to clean.
        """
        for text in texts:
            text.pop("text_image_connection_id", None)