from __future__ import annotations

from abc import ABC, abstractmethod
import itertools as it
import re
from scipy.optimize import linear_sum_assignment
from scipy.spatial.distance import cdist

from manimlib.constants import WHITE
from manimlib.logger import log
from manimlib.mobject.svg.svg_mobject import SVGMobject
from manimlib.mobject.types.vectorized_mobject import VGroup
from manimlib.utils.color import color_to_rgb
from manimlib.utils.color import rgb_to_hex
from manimlib.utils.config_ops import digest_config

from typing import TYPE_CHECKING, Callable

if TYPE_CHECKING:
    from colour import Color
    from typing import Iterable, TypeVar, Union

    ManimColor = Union[str, Color]
    Span = tuple[int, int]
    Selector = Union[
        str,
        re.Pattern,
        tuple[Union[int, None], Union[int, None]],
        Iterable[Union[
            str,
            re.Pattern,
            tuple[Union[int, None], Union[int, None]]
        ]]
    ]
    T = TypeVar("T")


class StringMobject(SVGMobject, ABC):
    """
    An abstract base class for `MTex` and `MarkupText`

    This class aims to optimize the logic of "slicing submobjects
    via substrings". This could be much clearer and more user-friendly
    than slicing through numerical indices explicitly.

    Users are expected to specify substrings in `isolate` parameter
    if they want to do anything with their corresponding submobjects.
    `isolate` parameter can be either a string, a `re.Pattern` object,
    or a 2-tuple containing integers or None, or a collection of the above.
    Note, substrings specified cannot *partly* overlap with each other.

    Each instance of `StringMobject` generates 2 svg files.
    The additional one is generated with some color commands inserted,
    so that each submobject of the original `SVGMobject` will be labelled
    by the color of its paired submobject from the additional `SVGMobject`.
    """
    CONFIG = {
        "height": None,
        "stroke_width": 0,
        "stroke_color": WHITE,
        "path_string_config": {
            "should_subdivide_sharp_curves": True,
            "should_remove_null_curves": True,
        },
        "base_color": WHITE,
        "isolate": (),
        "protect": (),
    }

    def __init__(self, string: str, **kwargs):
        self.string = string
        digest_config(self, kwargs)
        if self.base_color is None:
            self.base_color = WHITE
        #self.base_color_hex = self.color_to_hex(self.base_color)

        self.parse()
        super().__init__(**kwargs)
        self.labels = [submob.label for submob in self.submobjects]

    def get_file_path(self) -> str:
        original_content = self.get_content(is_labelled=False)
        return self.get_file_path_by_content(original_content)

    @abstractmethod
    def get_file_path_by_content(self, content: str) -> str:
        return ""

    def generate_mobject(self) -> None:
        super().generate_mobject()

        labels_count = len(self.labelled_spans)
        if labels_count == 1:
            for submob in self.submobjects:
                submob.label = 0
            return

        labelled_content = self.get_content(is_labelled=True)
        file_path = self.get_file_path_by_content(labelled_content)
        labelled_svg = SVGMobject(file_path)
        #print(len(self.submobjects), len(labelled_svg.submobjects))  # ????
        if len(self.submobjects) != len(labelled_svg.submobjects):
            log.warning(
                "Cannot align submobjects of the labelled svg "
                "to the original svg. Skip the labelling process."
            )
            for submob in self.submobjects:
                submob.label = labels_count - 1
            return

        self.rearrange_submobjects_by_positions(labelled_svg)
        unrecognizable_colors = []
        for submob, labelled_svg_submob in zip(
            self.submobjects, labelled_svg.submobjects
        ):
            color_int = self.hex_to_int(self.color_to_hex(
                labelled_svg_submob.get_fill_color()
            ))
            if color_int >= labels_count:
                unrecognizable_colors.append(color_int)
                color_int = labels_count
            submob.label = color_int - 1
        if unrecognizable_colors:
            log.warning(
                "Unrecognizable color labels detected (%s). "
                "The result could be unexpected.",
                ", ".join([
                    self.int_to_hex(color)
                    for color in unrecognizable_colors
                ])
            )

    def rearrange_submobjects_by_positions(
        self, labelled_svg: SVGMobject
    ) -> None:
        # Rearrange submobjects of `labelled_svg` so that
        # each submobject is labelled by the nearest one of `labelled_svg`.
        # The correctness cannot be ensured, since the svg may
        # change significantly after inserting color commands.
        if not labelled_svg.submobjects:
            return

        bb_0 = self.get_bounding_box()
        bb_1 = labelled_svg.get_bounding_box()
        scale_factor = abs((bb_0[2] - bb_0[0]) / (bb_1[2] - bb_1[0]))
        labelled_svg.move_to(self).scale(scale_factor)

        distance_matrix = cdist(
            [submob.get_center() for submob in self.submobjects],
            [submob.get_center() for submob in labelled_svg.submobjects]
        )
        _, indices = linear_sum_assignment(distance_matrix)
        labelled_svg.set_submobjects([
            labelled_svg.submobjects[index]
            for index in indices
        ])

    # Toolkits

    def find_spans_by_selector(self, selector: Selector) -> list[Span]:
        def find_spans_by_single_selector(sel):
            if isinstance(sel, str):
                return [
                    match_obj.span()
                    for match_obj in re.finditer(re.escape(sel), self.string)
                ]
            if isinstance(sel, re.Pattern):
                return [
                    match_obj.span()
                    for match_obj in sel.finditer(self.string)
                ]
            if isinstance(sel, tuple) and len(sel) == 2 and all(
                isinstance(index, int) or index is None
                for index in sel
            ):
                l = len(self.string)
                span = tuple(
                    default_index if index is None else
                    min(index, l) if index >= 0 else max(index + l, 0)
                    for index, default_index in zip(sel, (0, l))
                )
                return [span]
            return None

        result = find_spans_by_single_selector(selector)
        if result is None:
            result = []
            for sel in selector:
                spans = find_spans_by_single_selector(sel)
                if spans is None:
                    raise TypeError(f"Invalid selector: '{sel}'")
                result.extend(spans)
        return list(filter(lambda span: span[0] <= span[1], result))

    @staticmethod
    def span_contains(span_0: Span, span_1: Span) -> bool:
        return span_0[0] <= span_1[0] and span_0[1] >= span_1[1]

    @staticmethod
    def color_to_hex(color: ManimColor) -> str:
        return rgb_to_hex(color_to_rgb(color))

    @staticmethod
    def hex_to_int(rgb_hex: str) -> int:
        return int(rgb_hex[1:], 16)

    @staticmethod
    def int_to_hex(rgb_int: int) -> str:
        return f"#{rgb_int:06x}".upper()

    # Parsing

    def parse(self) -> None:
        def get_substr(span: Span) -> str:
            return self.string[slice(*span)]

        def get_neighbouring_pairs(vals: Iterable[T]) -> list[tuple[T, T]]:
            val_list = list(vals)
            return list(zip(val_list[:-1], val_list[1:]))

        #def get_complement_spans(
        #    universal_span: Span, interval_spans: list[Span]
        #) -> list[Span]:
        #    if not interval_spans:
        #        return [universal_span]

        #    span_ends, span_starts = zip(*interval_spans)
        #    return list(zip(
        #        (universal_span[0], *span_starts),
        #        (*span_ends, universal_span[1])
        #    ))

        def join_strs(strs: list[str], inserted_strs: list[str]) -> str:
            return "".join(it.chain(*zip(strs, (*inserted_strs, ""))))

        command_matches = self.get_command_matches(self.string)
        #command_spans = [match_obj.span() for match_obj in command_matches]
        configured_items = self.get_configured_items()
        #configured_spans = [span for span, _ in configured_items]
        #configured_attr_dicts = [d for _, d in configured_items]
        categorized_spans = [
            [(0, len(self.string))],  # TODO
            [span for span, _ in configured_items],
            self.find_spans_by_selector(self.isolate),
            self.find_spans_by_selector(self.protect),
            [match_obj.span() for match_obj in command_matches]  # TODO
        ]

        sorted_items = sorted([
            (category, category_index, flag, *span[::flag])
            for category, spans in enumerate(categorized_spans)
            for category_index, span in enumerate(spans)
            for flag in (1, -1)
        ], key=lambda t: (
            t[3], t[2] * (2 if t[3] != t[4] else -1), -t[4],
            t[2] * t[0], t[2] * t[1]
        ))  # TODO

        labelled_spans = []
        attr_dicts = []
        inserted_items = []
        #labelled_items = []

        count = 0
        region_index = 0
        protect_level = 0
        region_levels = [0]
        open_command_stack = []
        open_stack = []

        #protect_level_stack = []
        #bracket_level_stack = []
        #inserted_position_stack = []
        #index_items_len = 0  # count * 2
        for category, i, flag, _, _ in sorted_items:
            if category >= 3:
                if flag == 1:
                    protect_level += 1
                    continue
                protect_level -= 1
                if category == 3:
                    continue
                region_index += 1
                command_match = command_matches[i]
                command_flag = self.get_command_flag(command_match)
                region_levels.append(region_levels[-1] + command_flag)
                if command_flag == 1:
                    open_command_stack.append(
                        (command_match, region_index, count)
                    )
                    continue
                elif command_flag == 0:
                    continue
                command_match_, region_index_, count_ = open_command_stack.pop()
                attr_dict = self.get_attr_dict_from_command_pair(
                    command_match_, command_match
                )
                if attr_dict is None:
                    continue
                span = (command_match_.end(), command_match.start())
                region_span = (region_index_, region_index - 1)
            else:
                if flag == 1:
                    open_stack.append(
                        (category, i, protect_level, region_index, count)
                    )
                    continue
                category_, i_, protect_level_, region_index_, count_ \
                    = open_stack.pop()
                span = categorized_spans[category][i]
                if (category_, i_) != (category, i):
                    log.warning(
                        "Partly overlapping substrings detected: '%s' and '%s'",
                        get_substr(categorized_spans[category_][i_]),
                        get_substr(span)
                    )
                    continue
                if protect_level_ or protect_level:
                    continue
                ls = region_levels[region_index_:region_index + 1]
                if ls and (any(ls[0] > l for l in ls) or ls[0] < ls[-1]):
                    log.warning(
                        "Cannot handle substring '%s'", get_substr(span)
                    )
                    continue
                attr_dict = configured_items[i][1] if category == 1 else {}
                region_span = (region_index_, region_index)
            #labelled_items.append(
            #    (span, region_span, (count_, count), attr_dict)
            #)
            pos = count_ * 2
            labelled_spans.append(span)
            attr_dicts.append(attr_dict)
            inserted_items.insert(pos, (count, 1, span[0], region_span[0]))
            inserted_items.append((count, -1, span[1], region_span[1]))
            count += 1

        #labelled_spans = []
        #attr_dicts = []
        #inserted_items = []

        #inserted_items.insert(0, (-1, 1, 0, 0))
        #inserted_items.append((-1, -1, len(self.string), region_index))

        inserted_label_items = [
            (label, flag)
            for label, flag, _, _ in inserted_items
        ]
        #inserted_interval_spans = []
        #command_matches_lists = []
        #subpieces_lists = []
        content_pieces = []
        matching_pieces = []
        for (_, _, prev_index, prev_region_index), (_, _, next_index, next_region_index) in get_neighbouring_pairs(inserted_items):
            region_matches = command_matches[prev_region_index:next_region_index]
            #command_matches_lists.append(region_matches)
            subpieces = [
                get_substr((start, end))
                for start, end in zip(
                    [prev_index, *(m.end() for m in region_matches)],
                    [*(m.start() for m in region_matches), next_index]
                )
            ]
            content_pieces.append(join_strs(subpieces, [
                self.replace_for_content(m) for m in region_matches
            ]))
            matching_pieces.append(join_strs(subpieces, [
                self.replace_for_matching(m) for m in region_matches
            ]))
            #subpieces_lists.append([
            #    get_substr(s)
            #    for s in get_complement_spans(
            #        (prev_index, next_index),
            #        [m.span() for m in region_matches]
            #    )
            #])



        #inserted_interval_spans = get_neighbouring_pairs([
        #    index
        #    for _, _, index, _ in inserted_items
        #])
        #command_matches_lists = [
        #    command_matches[slice(*region_range)]
        #    for region_range in get_neighbouring_pairs([
        #        region_index
        #        for _, _, _, region_index in inserted_items
        #    ])
        #]

        #subpieces_lists = [
        #    [
        #        get_substr(s)
        #        for s in get_complement_spans(
        #            span, [m.span() for m in match_list]
        #        )
        #    ]
        #    for span, match_list in zip(inserted_interval_spans, command_matches_lists)
        #]

        #def get_replaced_pieces(replace_func: Callable[[re.Match], str]) -> list[str]:
        #    return [
        #        join_strs(subpieces, [
        #            replace_func(command_match)
        #            for command_match in match_list
        #        ])
        #        for subpieces, match_list in zip(subpieces_lists, command_matches_lists)
        #    ]

        #content_pieces = get_replaced_pieces(self.replace_for_content)
        #matching_pieces = get_replaced_pieces(self.replace_for_matching)

        def get_content(is_labelled: bool) -> str:
            inserted_strings = [
                self.get_command_string(
                    attr_dicts[label],
                    is_end=flag < 0,
                    label_hex=self.int_to_hex(label + 1) if is_labelled else None
                )
                for label, flag in inserted_label_items[1:-1]
            ]
            prefix, suffix = self.get_content_prefix_and_suffix(
                is_labelled=is_labelled
            )
            return "".join([
                prefix,
                join_strs(content_pieces, inserted_strings),
                suffix
            ])

        def get_group_part_items_by_labels(labels: list[int]) -> list[tuple[str, list[int]]]:
            if not labels:
                return []

            range_lens, group_labels = zip(*(
                (len(list(grouper)), val)
                for val, grouper in it.groupby(labels)
            ))
            submob_indices_lists = [
                list(range(*submob_range))
                for submob_range in get_neighbouring_pairs(
                    [0, *it.accumulate(range_lens)]
                )
            ]

            def get_region_index(label, flag):
                #if label == -1:
                #    if flag == 1:
                #        return 0
                #    return len(inserted_label_items) - 1
                return inserted_label_items.index((label, flag))

            def get_labelled_span(label):
                #if label == -1:
                #    return (0, len(self.string))
                return labelled_spans[label]

            def label_contains(label_0, label_1):
                return self.span_contains(
                    get_labelled_span(label_0), get_labelled_span(label_1)
                )

            piece_starts = [
                get_region_index(group_labels[0], 1),
                *(
                    get_region_index(curr_label, 1)
                    if label_contains(prev_label, curr_label)
                    else get_region_index(prev_label, -1)
                    for prev_label, curr_label in get_neighbouring_pairs(
                        group_labels
                    )
                )
            ]
            piece_ends = [
                *(
                    get_region_index(curr_label, -1)
                    if label_contains(next_label, curr_label)
                    else get_region_index(next_label, 1)
                    for curr_label, next_label in get_neighbouring_pairs(
                        group_labels
                    )
                ),
                get_region_index(group_labels[-1], -1)
            ]

            #piece_ranges = get_complement_spans(
            #    (get_region_index(group_labels[0], 1), get_region_index(group_labels[-1], -1)),
            #    [
            #        (
            #            get_region_index(next_label, 1)
            #            if label_contains(prev_label, next_label)
            #            else get_region_index(prev_label, -1),
            #            get_region_index(prev_label, -1)
            #            if label_contains(next_label, prev_label)
            #            else get_region_index(next_label, 1)
            #        )
            #        for prev_label, next_label in get_neighbouring_pairs(
            #            group_labels
            #        )
            #    ]
            #)
            group_substrs = [
                re.sub(r"\s+", "", "".join(
                    matching_pieces[start:end]
                ))
                for start, end in zip(piece_starts, piece_ends)
            ]
            return list(zip(group_substrs, submob_indices_lists))

        #print(labelled_spans)
        self.labelled_spans = labelled_spans
        self.get_content = get_content
        self.get_group_part_items_by_labels = get_group_part_items_by_labels

    @staticmethod
    @abstractmethod
    def get_command_matches(string: str) -> list[re.Match]:
        return []

    @staticmethod
    @abstractmethod
    def get_command_flag(match_obj: re.Match) -> int:
        return 0

    @staticmethod
    @abstractmethod
    def replace_for_content(match_obj: re.Match) -> str:
        return ""

    @staticmethod
    @abstractmethod
    def replace_for_matching(match_obj: re.Match) -> str:
        return ""

    @staticmethod
    @abstractmethod
    def get_attr_dict_from_command_pair(
        open_command: re.Match, close_command: re.Match,
    ) -> dict[str, str] | None:
        return None

    @abstractmethod
    def get_configured_items(self) -> list[tuple[Span, dict[str, str]]]:
        return []

    @staticmethod
    @abstractmethod
    def get_command_string(
        attr_dict: dict[str, str], is_end: bool, label_hex: str | None
    ) -> str:
        return ""

    @abstractmethod
    def get_content_prefix_and_suffix(
        self, is_labelled: bool
    ) -> tuple[str, str]:
        return "", ""

    # Selector

    def get_submob_indices_list_by_span(
        self, arbitrary_span: Span
    ) -> list[int]:
        return [
            submob_index
            for submob_index, label in enumerate(self.labels)
            if self.span_contains(arbitrary_span, self.labelled_spans[label])
        ]

    def get_specified_part_items(self) -> list[tuple[str, list[int]]]:
        return [
            (
                self.string[slice(*span)],
                self.get_submob_indices_list_by_span(span)
            )
            for span in self.labelled_spans[:-1]
        ]

    def get_group_part_items(self) -> list[tuple[str, list[int]]]:
        return self.get_group_part_items_by_labels(self.labels)

    def get_submob_indices_lists_by_selector(
        self, selector: Selector
    ) -> list[list[int]]:
        return list(filter(
            lambda indices_list: indices_list,
            [
                self.get_submob_indices_list_by_span(span)
                for span in self.find_spans_by_selector(selector)
            ]
        ))

    def build_parts_from_indices_lists(
        self, indices_lists: list[list[int]]
    ) -> VGroup:
        return VGroup(*[
            VGroup(*[
                self.submobjects[submob_index]
                for submob_index in indices_list
            ])
            for indices_list in indices_lists
        ])

    def build_groups(self) -> VGroup:
        return self.build_parts_from_indices_lists([
            indices_list
            for _, indices_list in self.get_group_part_items()
        ])

    def select_parts(self, selector: Selector) -> VGroup:
        return self.build_parts_from_indices_lists(
            self.get_submob_indices_lists_by_selector(selector)
        )

    def select_part(self, selector: Selector, index: int = 0) -> VGroup:
        return self.select_parts(selector)[index]

    def set_parts_color(self, selector: Selector, color: ManimColor):
        self.select_parts(selector).set_color(color)
        return self

    def set_parts_color_by_dict(self, color_map: dict[Selector, ManimColor]):
        for selector, color in color_map.items():
            self.set_parts_color(selector, color)
        return self

    def get_string(self) -> str:
        return self.string