import abc
from typing import Iterator, Optional, Generator, Union
from .tree_traversals import pre_order_traversal
from .bounding_box import BoundingBox
from .ligatures import LIGATURES_TO_TEXT
import logging
from .constants import (
    LINE_DELIMITER,
    PAGE_DELIMITER,
    PARAGRAPH_DELIMITER,
    WORD_DELIMITER,
)


class TextBlock(metaclass=abc.ABCMeta):
    """
    Abstract class for all text blocks: Document, Page, Paragraph, Line, Word, Character.
    Every text box has a bounding box, a delimiter, a list of children(1) and a parent(2).
    The Document block is the root of the tree, while the Character block is the leaf of the tree.
    (1) The Document doesn't have a parent, (2) while the Character doesn't have children.
    """

    child_type = None

    def __init__(
        self,
        bounding_box: BoundingBox = None,
        delimiter="",
        children: list["TextBlock"] = None,
        parent: Optional["TextBlock"] = None,
    ):

        self.length: Optional[int] = None
        self.bounding_box: BoundingBox = bounding_box
        # self.children: list["TextBlock"] = children
        self.delimiter: str = delimiter
        self.set_children(children)
        self.set_parent(parent)
        self.relative_order: int = None
        self.start_index: int = None
        self.end_index: int = None

    def set_relative_order(self, index: int):
        """The index is the relative order of the block in the list of children of its parent."""
        self.relative_order = index

    def get_relative_order(self) -> int:
        if self.relative_order is None:
            raise ValueError("The relative order of the block is not set")
        return self.relative_order

    def set_bounding_box(self):
        """
        Setting the bounding box of a box based on the bounding boxes of its children.
        For more information see the documentation of the BoundingBox class.
        """
        if self.bounding_box is not None:
            return
        bounding_boxes = [child.get_bounding_box() for child in self.children]
        for bounding_box in bounding_boxes:
            if not isinstance(bounding_box, BoundingBox):
                raise TypeError(
                    f"When calling set_bounding_box for a box, all of it's children must have a bounding box"
                )
        self.bounding_box = BoundingBox.compose_bounding_boxes(bounding_boxes)

    def get_bounding_box(self) -> BoundingBox:
        """The bounding box of the box is the bounding box surrounding the text of the box."""
        if self.bounding_box is None:
            self.set_bounding_box()
        return self.bounding_box

    def set_children(self, children: Optional[list["TextBlock"]]):
        self.children: list["TextBlock"] = []
        if children is None:
            return
        for child in children:
            self.add_child(child)

    def get_children(self) -> list:
        return self.children

    def set_parent(self, parent: Optional["TextBlock"]):
        self.parent = parent

    def split_to_chunks(
        self, byte_threshold: int, encoding: str = "utf-8"
    ) -> list[list["TextBlock"]]:

        over_partitioned_chunks = self._split_to_chunks(byte_threshold, encoding)
        if not over_partitioned_chunks:
            return [[]]
        all_chunks = []

        current_type = type(over_partitioned_chunks[0])
        current_chunk = [over_partitioned_chunks[0]]
        current_chunk_size = over_partitioned_chunks[0].get_text_byte_size(encoding)

        for text_block in over_partitioned_chunks[1:]:

            if type(text_block) != current_type:
                current_type = type(text_block)
                current_chunk = [text_block]
                current_chunk_size = text_block.get_text_byte_size(encoding)

            elif (
                current_chunk_size + text_block.get_text_byte_size(encoding) + 1
                <= byte_threshold
            ):
                current_chunk.append(text_block)
                current_chunk_size += text_block.get_text_byte_size(encoding) + 1

        all_chunks.append(current_chunk)
        return all_chunks

    def _split_to_chunks(
        self, byte_threshold: int, encoding: str = "utf-8"
    ) -> list["TextBlock"]:
        if self.get_text_byte_size(encoding) <= byte_threshold:
            return [self]
        children_splitted_text = []
        for child in self.children:
            child_splitted_text = child._split_to_chunks(byte_threshold, encoding)
            children_splitted_text.extend(child_splitted_text)
        return children_splitted_text

    @staticmethod
    def pre_order_traversal(root: Optional["TextBlock"]) -> Generator:
        """
        Pre-order traversal is a depth-first traversal where the root is visited first, then the children.
        """
        if root is None:
            return
        yield root
        for child in root.children:
            yield from TextBlock.pre_order_traversal(child)

    @staticmethod
    def post_order_traversal(root: Optional["TextBlock"]) -> Generator:
        """
        Post-order traversal is a depth-first traversal where the root is visited last, after the children.
        """
        if root is None:
            return
        for child in root.children:
            yield from TextBlock.post_order_traversal(child)

        yield root

    @staticmethod
    def get_depth_by_type(_class: type):
        """
        Returns the depth of a block in the tree based on its type.
        """
        if _class == Document:
            return 0
        if _class == Page:
            return 1
        if _class == Paragraph:
            return 2
        if _class == Line:
            return 3
        if _class == Word:
            return 4
        if _class == Character:
            return 5
        raise TypeError("Type not recognized")

    def get_parent(self) -> "TextBlock":
        """
        Returns the parent of the block.
        """
        return self.parent

    def get_ancestor(self, _class: type) -> "TextBlock":
        """
        Returns the ancestor of the block based on the type of the ancestor.
        """
        self_depth = TextBlock.get_depth_by_type(type(self))
        parent_depth = TextBlock.get_depth_by_type(_class)
        if self_depth < parent_depth:
            raise TypeError(f"{_class} is not an ancestor of {type(self)}!")
        relative_depth = self_depth - parent_depth
        current = self
        for _ in range(relative_depth):
            current = current.get_parent()
        return current

    def add_child(self, child: "TextBlock"):
        """
        Adding a child to the list of children of the block.
        The parent and the index of the child is set automatically.
        """
        if not isinstance(child, self.child_type):
            raise TypeError(
                f"Child of type {type(self)} must be of type {self.child_type}"
            )
        self.children.append(child)
        child.set_parent(self)
        child.set_relative_order(len(self.children) - 1)

    def get_text(self):
        """
        Every block has an implicit text representation that can be retrieved by calling this method.
        """
        return self.delimiter.join([child.get_text() for child in self.children])

    def set_length(self):
        """
        self.length is equivalent to len(self.get_text()), but it is computed once and stored for later use.
        """
        if self.length is not None:
            return
        if not self.children:
            self.length = 0
            return
        lengths = [child.get_length() for child in self.children]
        self.length = sum(lengths) + len(self.delimiter) * (len(self.children) - 1)

    def set_text_byte_size(self, encoding: str = "utf-8"):
        if self.text_byte_size is not None:
            return

        if not self.children:
            self.text_byte_size = 0
            return

        byte_sizes = [child.get_text_byte_size(encoding) for child in self.children]

        self.text_byte_size = sum(byte_sizes) + len(self.children) - 1

    def get_text_byte_size(self, encoding: str = "utf-8") -> int:
        if self.text_byte_size is None:
            self.set_text_byte_size(encoding)
        return self.text_byte_size

    def get_all(self, _class: type) -> Generator:
        """
        Returns a generator of all the blocks of the given type that are descendants of the block.
        """
        if type(self) == _class:
            yield self
        for child in self.children:
            yield from child.get_all(_class)

    def get_length(self) -> int:
        if self.length is None:
            self.set_length()
        return self.length

    def __str__(self):
        return self.get_text()

    def __repr__(self):
        return f"{self.__class__.__name__}:\n{self.get_text()}"

    def __len__(self):
        """The length of a block is the number of children it has."""
        return self.length

    def __iter__(self) -> Iterator["TextBlock"]:
        """Iterating over the children of a block."""
        return iter(self.get_all(Character))

    def __getitem__(self, index: Union[int, slice]) -> Optional["Character"]:
        """Getting the child at the given index."""

        if isinstance(index, slice):
            start = index.start
            stop = index.stop
            step = index.step if index.step is not None else 1
            return [self.__getitem__(i) for i in range(start, stop, step)]

        if index < 0 or index >= len(self):
            raise IndexError("Index out of range of text length")
        return self._get_char(index)

    def _get_char(self, index: int) -> Optional["Character"]:
        for child in self.children:
            if 0 <= index < child.length:
                return child._get_char(index)
            index -= child.length + len(self.delimiter)
        return None

    def set_absolute_order(self, absolute_order: int):
        """
        Set the absolute order of the block in the document.
        """
        self.absolute_order = absolute_order

    def get_absolute_order(self) -> int:
        """
        Returns the absolute order of the block in the document.
        """
        return self.absolute_order

    def set_start_index(self) -> int:
        """
        Returns the index of the first character of the block in the document.
        """
        if self.parent is None:
            self.start_index = 0
            return
        start_index = self.parent.get_start_index()
        for child in self.parent.children[: self.relative_order]:
            start_index += child.get_length() + len(self.parent.delimiter)
        self.start_index = start_index

    def get_start_index(self):
        if self.start_index is None:
            self.set_start_index()
        return self.start_index

    def set_end_index(self):
        """
        Set the index of the last character of the block in the document.
        """
        self.end_index = self.get_start_index() + self.get_length()

    def get_end_index(self) -> int:
        if self.end_index is None:
            self.set_end_index()
        return self.end_index


class Character(TextBlock):
    """
    The Character class represents a single character (not including delimiters).
    Notice this class play a special role in the hierarchy of TextBlock because it is the leaf of the document tree,
    and as such, it overrides some of the methods of its parent class.
    """

    child_type = None

    def __init__(self, bounding_box: BoundingBox, char: str):
        super().__init__(bounding_box=bounding_box)
        self.char = char
        self.verify_args()

    def verify_args(self):
        if not isinstance(self.char, str):
            raise ValueError("char must be a string")
        if len(self.char) != 1:
            logging.warning(
                "char %s is not a single character. It's length is %s",
                self.char,
                len(self.char),
            )

    def get_text(self) -> str:
        return self.char

    def set_length(self) -> int:
        self.length = 1

    def set_text_byte_size(self, encoding: str = "utf-8"):
        self.text_byte_size = len(self.char.encode(encoding))

    def _get_char(self, index: int) -> str:
        return self


class Word(TextBlock):
    """
    The Word class represents a single word.
    """

    child_type = Character

    def __init__(
        self, text: str, bounding_box: BoundingBox, children: list[Character] = None
    ):
        super().__init__(bounding_box=bounding_box, children=children)
        text = self.get_clean_text(text)
        chars_bboxes = BoundingBox.split_bounding_box(bounding_box, len(text))
        self.text = text
        for char, char_bbox in zip(text, chars_bboxes):
            self.add_child(Character(char_bbox, char))

    def get_clean_text(self, text: str) -> str:
        clean_text = [LIGATURES_TO_TEXT.get(char, char) for char in text]
        clean_text = "".join(clean_text)
        return clean_text


class Line(TextBlock):
    """
    The Line class represents a single line.
    """

    child_type = Word

    def __init__(self, bounding_box: BoundingBox = None, children: list[Word] = None):
        super().__init__(
            bounding_box=bounding_box, delimiter=WORD_DELIMITER, children=children
        )


class Paragraph(TextBlock):
    """
    The Paragraph class represents a single paragraph.
    """

    child_type = Line

    def __init__(self, bounding_box: BoundingBox = None, children: list[Line] = None):
        super().__init__(
            bounding_box=bounding_box, delimiter=LINE_DELIMITER, children=children
        )


class Page(TextBlock):
    """
    The Page class represents a single page.
    """

    child_type = Paragraph

    def __init__(self, children: list[Paragraph] = None):
        super().__init__(
            bounding_box=BoundingBox(0, 1, 1, 0),
            delimiter=PARAGRAPH_DELIMITER,
            children=children,
        )

    def set_bounding_box(self):
        """The bounding box of the page is fixed to BoundingBox(left=0, top=1, right=1, bottom=0).
        The coordinates of the blocks are PDF coordinates (i.e. the origin is at the bottom left), and normalized to the page size."""
        return


class Document(TextBlock):
    """
    The Document class represents a single document.
    """

    child_type = Page

    def __init__(self, children: list[Page] = None):
        super().__init__(delimiter=PAGE_DELIMITER, children=children)
        self.set_length()

        for text_block in pre_order_traversal(self):
            text_block.set_start_index()
            text_block.set_end_index()
        block_types = [Page, Paragraph, Line, Word, Character]
        for block_type in block_types:
            text_blocks = self.get_all(block_type)
            for i, text_block in enumerate(text_blocks):
                text_block.set_absolute_order(i)

        self.text = self.get_text()

    def set_bounding_box(self):
        return

    def get_text_block_number(
        self, index: int, text_block_type: type, absolute: bool
    ) -> Optional[int]:
        """Get the absolute or relative order of the text block of type text_block_type that contains
        the character at index `index`."""
        character = self[index]
        if character is not None:
            text_block = character.get_ancestor(text_block_type)
            if absolute:
                return text_block.get_absolute_order()
            return text_block.get_relative_order()
        previous_char = self[index - 1]
        next_char = self[index + 1]
        if previous_char is None or next_char is None:
            return None
        previous_text_block = previous_char.get_ancestor(text_block_type)
        next_text_block = next_char.get_ancestor(text_block_type)
        if previous_text_block == next_text_block:
            if absolute:
                return previous_text_block.get_absolute_order()
            return previous_text_block.get_relative_order()
        return None

    def get_relative_page_number(self, index: int) -> int:
        return self.get_text_block_number(index, Page, absolute=False)

    def get_relative_paragraph_number(self, index: int) -> int:
        return self.get_text_block_number(index, Paragraph, absolute=False)

    def get_relative_line_number(self, index: int) -> int:
        return self.get_text_block_number(index, Line, absolute=False)

    def get_relative_word_number(self, index: int) -> int:
        return self.get_text_block_number(index, Word, absolute=False)

    def get_relative_character_number(self, index: int) -> int:
        return self.get_text_block_number(index, Character, absolute=False)

    def get_absolute_page_number(self, index: int) -> int:
        return self.get_text_block_number(index, Page, absolute=True)

    def get_absolute_paragraph_number(self, index: int) -> int:
        return self.get_text_block_number(index, Paragraph, absolute=True)

    def get_absolute_line_number(self, index: int) -> int:
        return self.get_text_block_number(index, Line, absolute=True)

    def get_absolute_word_number(self, index: int) -> int:
        return self.get_text_block_number(index, Word, absolute=True)

    def get_absolute_character_number(self, index: int) -> int:
        return self.get_text_block_number(index, Character, absolute=True)
