from typing import Any, List, Optional, Type, Union, get_args

from rich.highlighter import RegexHighlighter
from rich.theme import Theme
from rich.tree import Tree
from typing_extensions import TYPE_CHECKING
from typing_inspect import is_optional_type, is_union_type

from docarray.base_doc.doc import BaseDoc
from docarray.display.tensor_display import TensorDisplay
from docarray.typing import ID
from docarray.typing.tensor.abstract_tensor import AbstractTensor
from docarray.utils._internal._typing import safe_issubclass

if TYPE_CHECKING:
    from rich.console import Console, ConsoleOptions, RenderResult


class DocumentSummary:
    table_min_width: int = 40

    def __init__(
        self,
        doc: Optional['BaseDoc'] = None,
    ):
        self.doc = doc

    def summary(self) -> None:
        """Print non-empty fields and nested structure of this Document object."""
        import rich

        t = self._plot_recursion(node=self)
        rich.print(t)

    @staticmethod
    def schema_summary(cls: Type['BaseDoc']) -> None:
        """Print a summary of the Documents schema."""
        from rich.console import Console
        from rich.panel import Panel

        panel = Panel(
            DocumentSummary._get_schema(cls),
            title='Document Schema',
            expand=False,
            padding=(1, 3),
        )
        highlighter = SchemaHighlighter()

        console = Console(highlighter=highlighter, theme=highlighter.theme)
        console.print(panel)

    @staticmethod
    def _get_schema(
        cls: Type['BaseDoc'],
        doc_name: Optional[str] = None,
        recursion_list: Optional[List] = None,
    ) -> Tree:
        """Get Documents schema as a rich.tree.Tree object."""
        import re

        from rich.tree import Tree

        from docarray import BaseDoc, DocList

        if recursion_list is None:
            recursion_list = []

        if cls in recursion_list:
            return Tree(cls.__name__)
        else:
            recursion_list.append(cls)

        root = cls.__name__ if doc_name is None else f'{doc_name}: {cls.__name__}'
        tree = Tree(root, highlight=True)

        for field_name, value in cls._docarray_fields().items():
            if field_name != 'id':
                field_type = value.annotation
                field_cls = str(field_type).replace('[', '\[')
                field_cls = re.sub('<class \'|\'>|[a-zA-Z_]*[.]', '', field_cls)

                node_name = f'{field_name}: {field_cls}'

                if is_union_type(field_type) or is_optional_type(field_type):
                    sub_tree = Tree(node_name, highlight=True)
                    for arg in get_args(field_type):
                        if safe_issubclass(arg, BaseDoc):
                            sub_tree.add(
                                DocumentSummary._get_schema(
                                    cls=arg, recursion_list=recursion_list
                                )
                            )
                        elif safe_issubclass(arg, DocList):
                            sub_tree.add(
                                DocumentSummary._get_schema(
                                    cls=arg.doc_type, recursion_list=recursion_list
                                )
                            )
                    tree.add(sub_tree)

                elif safe_issubclass(field_type, BaseDoc):
                    tree.add(
                        DocumentSummary._get_schema(
                            cls=field_type,
                            doc_name=field_name,
                            recursion_list=recursion_list,
                        )
                    )

                elif safe_issubclass(field_type, DocList):
                    sub_tree = Tree(node_name, highlight=True)
                    sub_tree.add(
                        DocumentSummary._get_schema(
                            cls=field_type.doc_type, recursion_list=recursion_list
                        )
                    )
                    tree.add(sub_tree)

                else:
                    tree.add(node_name)

        return tree

    def __rich_console__(
        self, console: 'Console', options: 'ConsoleOptions'
    ) -> 'RenderResult':
        kls = self.doc.__class__.__name__
        doc_id = getattr(self.doc, 'id')
        if doc_id is not None:
            yield f':page_facing_up: [b]{kls} [/b]: [cyan]{doc_id[:7]} ...[cyan]'
        else:
            yield f':page_facing_up: [b]{kls} [/b]'

        from rich import box, text
        from rich.table import Table

        from docarray import BaseDoc, DocList

        table = Table(
            'Attribute',
            'Value',
            min_width=self.table_min_width,
            box=box.ROUNDED,
            highlight=True,
        )

        for field_name, value in self.doc.__dict__.items():
            col_1 = f'{field_name}: {value.__class__.__name__}'
            if (
                isinstance(value, (ID, DocList, BaseDoc))
                or field_name.startswith('_')
                or value is None
            ):
                continue
            elif isinstance(value, (str, bytes)):
                col_2 = str(value)[:50]
                if len(value) > 50:
                    col_2 += f' ... (length: {len(value)})'
                table.add_row(col_1, text.Text(col_2))
            elif isinstance(value, AbstractTensor):
                table.add_row(col_1, TensorDisplay(tensor=value))
            elif isinstance(value, (tuple, list, set, frozenset)):
                value_list = list(value)
                col_2 = ''
                for i, x in enumerate(value_list):
                    if len(col_2) + len(str(x)) < 50:
                        col_2 = str(value_list[: i + 1])
                    else:
                        col_2 = f'{col_2[:-1]}, ...] (length: {len(value_list)})'
                        break

                if type(value) == tuple:
                    col_2 = col_2.replace('[', '(', 1).replace(']', ')', -1)
                if type(value) == set or type(value) == frozenset:
                    col_2 = col_2.replace('[', '{', 1).replace(']', '}', -1)

                table.add_row(col_1, text.Text(col_2))
            elif isinstance(value, dict):
                col_2 = f'{value}'
                if len(col_2) > 50:
                    col_2 = f'{col_2[: 50]}' + ' ... } ' + f'(length: {len(value)})'
                table.add_row(col_1, text.Text(col_2))
            else:
                col_2 = f'{value}'
                table.add_row(col_1, text.Text(col_2))

        if table.rows:
            yield table

    @staticmethod
    def _plot_recursion(
        node: Union['DocumentSummary', Any], tree: Optional[Tree] = None
    ) -> Tree:
        """
        Store node's children in rich.tree.Tree recursively.

        :param node: Node to get children from.
        :param tree: Append to this tree if not None, else use node as root.
        :return: Tree with all children.

        """
        from docarray import BaseDoc, DocList

        tree = Tree(node) if tree is None else tree.add(node)  # type: ignore

        if hasattr(node, '__dict__'):
            nested_attrs = [
                k
                for k, v in node.doc.__dict__.items()
                if isinstance(v, (DocList, BaseDoc))
            ]
            for attr in nested_attrs:
                value = getattr(node.doc, attr)
                attr_type = value.__class__.__name__
                icon = ':diamond_with_a_dot:'

                if isinstance(value, BaseDoc):
                    icon = ':large_orange_diamond:'
                    value = [value]

                match_tree = tree.add(f'{icon} [b]{attr}: ' f'{attr_type}[/b]')
                max_show = 2
                for i, d in enumerate(value):
                    if i == max_show:
                        doc_type = d.__class__.__name__
                        DocumentSummary._plot_recursion(
                            f'... {len(value) - max_show} more {doc_type} documents\n',
                            tree=match_tree,
                        )
                        break
                    DocumentSummary._plot_recursion(DocumentSummary(doc=d), match_tree)

        return tree


class SchemaHighlighter(RegexHighlighter):
    """Highlighter to apply colors to a Document's schema tree."""

    highlights = [
        r'(?P<class>^[A-Z][a-zA-Z]*)',
        r'(?P<attr>^.*(?=:))',
        r'(?P<attr_type>(?<=:).*$)',
        r'(?P<union_or_opt>Union|Optional)',
        r'(?P<other_chars>[\[\],:])',
    ]

    theme = Theme(
        {
            'class': 'orange3',
            'attr': 'green4',
            'attr_type': 'medium_orchid',
            'union_or_opt': 'medium_purple4',
            'other_chars': 'black',
        }
    )
