from typing import Dict, Generator, ValuesView, TYPE_CHECKING

from collections import defaultdict

from .filtering import ElementList
from .exceptions import InvalidSectionError, SectionNotFoundError

if TYPE_CHECKING:
    from .components import PDFDocument, PDFElement


class Section:
    """
    A continuous group of elements within a document.

    A section is intended to label a group of elements. Said elements must be continuous
    in the document.

    Warning:
        You should not instantiate a Section class yourself, but should call
        `create_section` from the `Sectioning` class below.

    Args:
        document (PDFDocument): A reference to the document.
        name (str): The name of the section.
        unique_name (str): Multiple sections can have the same name, but a unique name
            will be generated by the Sectioning class.
        start_element (PDFElement): The first element in the section.
        end_element (PDFElement): The last element in the section.
    """

    document: "PDFDocument"
    name: str
    unique_name: str
    start_element: "PDFElement"
    end_element: "PDFElement"

    def __init__(self, document, name, unique_name, start_element, end_element):
        if start_element._index > end_element._index:
            raise InvalidSectionError("end_element must come after start_element")
        self.document = document
        self.name = name
        self.unique_name = unique_name
        self.start_element = start_element
        self.end_element = end_element

    def __contains__(self, element: "PDFElement") -> bool:
        return element in self.elements

    @property
    def elements(self) -> "ElementList":
        """
        All the elements in the section.

        Returns:
            ElementList: All the elements in the section.
        """
        return self.document.elements.between(
            self.start_element, self.end_element, inclusive=True
        )

    def __eq__(self, other: object) -> bool:
        """
        Returns True if the two sections have the same unique name and are from the
        same document
        """
        if not isinstance(other, Section):
            raise NotImplementedError(f"Can't compare Section with {type(other)}")
        return all(
            [
                self.document == other.document,
                self.unique_name == other.unique_name,
                self.start_element == other.start_element,
                self.end_element == other.end_element,
                self.__class__ == other.__class__,
            ]
        )

    def __len__(self):
        """
        Returns the number of elements in the section.
        """
        return len(self.elements)

    def __repr__(self):
        return (
            f"<Section name: '{self.name}', unique_name: '{self.unique_name}', "
            f"number of elements: {len(self)}>"
        )


class Sectioning:
    """
    A sectioning utilities class, made available on all PDFDocuments as ``.sectioning``.
    """

    document: "PDFDocument"
    name_counts: Dict[str, int]
    sections_dict: Dict[str, Section]

    def __init__(self, document: "PDFDocument"):
        self.sections_dict = {}
        self.name_counts = defaultdict(int)
        self.document = document

    def create_section(
        self,
        name: str,
        start_element: "PDFElement",
        end_element: "PDFElement",
        include_last_element: bool = True,
    ):
        """
        Creates a new section with the specified name.

        Creates a new section with the specified name, starting at `start_element` and
        ending at `end_element` (inclusive). The unique name will be set to name_<idx>
        where <idx> is the number of existing sections with that name.

        Args:
            name (str): The name of the new section.
            start_element (PDFElement): The first element in the section.
            end_element (PDFElement): The last element in the section.
            include_last_element (bool): Whether the end_element should be included in
                the section, or only the elements which are strictly before the end
                element. Default: True (i.e. include end_element).

        Returns:
            Section: The created section.

        Raises:
            InvalidSectionError: If a the created section would be invalid. This is
                usually because the end_element comes after the start element.
        """
        current_count = self.name_counts[name]
        unique_name = f"{name}_{current_count}"
        self.name_counts[name] += 1

        if not include_last_element:
            if end_element._index == 0:
                raise InvalidSectionError(
                    "Section would contain no elements as end_element is the first "
                    "element in the document and include_last_element is False"
                )
            # We simply drop the index by one to get the element before
            end_element = self.document._element_list[end_element._index - 1]
        section = Section(self.document, name, unique_name, start_element, end_element)
        self.sections_dict[unique_name] = section
        return section

    def get_sections_with_name(self, name: str) -> Generator[Section, None, None]:
        """
        Returns a list of all sections with the given name.
        """
        return (
            self.sections_dict[f"{name}_{idx}"]
            for idx in range(0, self.name_counts[name])
        )

    def get_section(self, unique_name: str) -> Section:
        """
        Returns the section with the given unique name.

        Raises:
            SectionNotFoundError: If there is no section with the given unique_name.
        """
        try:
            return self.sections_dict[unique_name]
        except KeyError as err:
            raise SectionNotFoundError(
                f"Could not find section with name {unique_name}"
            ) from err

    @property
    def sections(self) -> ValuesView[Section]:
        """
        Returns the list of all created Sections.
        """
        return self.sections_dict.values()
