from __future__ import annotations

from collections import defaultdict
from typing import Optional, Dict, List, Type, TypeVar

from clarabot_nano.indexer.extractor.base import ALL_SPACE_SEARCH_RE, StopExtraction
from clarabot_nano.indexer.extractor.xml.base import AbstractCollector, CollectorHandler

# NOTE: big, strike and font are not HTML5 compatible
INLINE_TEXT_TAGS = {"a", "b", "i", "u", "s", "em", "strong", "small", "img", "big", "strike", "font"}
HEADER_TEXT_TAGS = {"h1", "h2", "h3", "h4"}


class BodyTagProcessor(object):
    __slots__ = ("tag",)

    def __init__(self, h: BodyHandler, tag: str, attrs: dict):
        self.tag = tag
        self._check_separator(h)

    def close(self, h: BodyHandler):
        self._check_separator(h)

    def _check_separator(self, h: BodyHandler):
        if not h.text_ends_with_separator:
            if self.tag not in INLINE_TEXT_TAGS:
                if h.headers and h.headers[-1] != "\n":
                    h.headers.append("\n")
                if h.texts and h.texts[-1] != "\n":
                    h.texts.append("\n")
                h.text_ends_with_separator = True

    def process(self, h: BodyHandler, text: str):
        # Spaces and irrelevant formatting is collapsed during the collection to avoid counting those towards the
        # length of the extracted text. This is how <p> tags are intended to behave anyways for example.
        # NOTE: the redundant formatting left by this will be normalized for the extractor as a final step.
        if text:  # ignore empty
            text = ALL_SPACE_SEARCH_RE.sub(" ", text)
            self._collect_text(h, text)

    def _collect_text(self, h: BodyHandler, text: str):
        if self.tag in HEADER_TEXT_TAGS:
            h.headers.append(text)
        h.texts.append(text)
        h.text_ends_with_separator = False
        h.texts_len += len(text)
        if h.extract_text_limit < h.texts_len:
            raise StopExtraction("enough-content")


class ListTagProcessor(BodyTagProcessor):
    __slots__ = ("count",)

    def __init__(self, h: BodyHandler, tag: str, attrs: dict):
        BodyTagProcessor.__init__(self, h, tag, attrs)
        self.count = 0

    def li(self):
        if self.tag == "ul":
            return "\u2022 "
        self.count += 1
        return f"{self.count}. "


class ListItemTagProcessor(BodyTagProcessor):
    __slots__ = ()

    def __init__(self, h: BodyHandler, tag: str, attrs: dict):
        BodyTagProcessor.__init__(self, h, tag, attrs)
        lp = h.get_parent_processor_by_type(ListTagProcessor)
        if lp:
            self._collect_text(h, lp.li())


class ImageTagProcessor(BodyTagProcessor):
    __slots__ = ()

    def __init__(self, h: BodyHandler, tag: str, attrs: dict):
        BodyTagProcessor.__init__(self, h, tag, attrs)
        if alt := attrs.get("alt"):
            self._collect_text(h, alt)


_T_BodyTagProcessor = TypeVar("_T_BodyTagProcessor", bound=BodyTagProcessor)


class BodyHandler(CollectorHandler):
    TAG = "body"

    TAG_PROCESSORS = {
        "ul": ListTagProcessor,
        "ol": ListTagProcessor,
        "li": ListItemTagProcessor,
        "img": ImageTagProcessor,
    }

    def __init__(self, extract_text_limit: int, sub_handlers: Dict[str, AbstractCollector] = None):
        CollectorHandler.__init__(self, sub_handlers)
        self.extract_text_limit = extract_text_limit
        self.headers = []
        self.texts = []
        self.texts_len = 0
        self.text_ends_with_separator = True
        self._tag_stack: List[BodyTagProcessor] = []
        self._tag_type_stacks: Dict[Type[BodyTagProcessor], List[BodyTagProcessor]] = defaultdict(list)

    def init(self, attrs: dict):
        self._tag_stack.append(tp := BodyTagProcessor(self, self.TAG, attrs))
        self._tag_type_stacks[BodyTagProcessor].append(tp)

    def start(self, tag: str, attrs: dict):
        tpt = self.TAG_PROCESSORS.get(tag, BodyTagProcessor)
        self._tag_stack.append(tp := tpt(self, tag, attrs))
        self._tag_type_stacks[tpt].append(tp)

    def end(self, tag: str):
        tp = self._tag_stack.pop()
        self._tag_type_stacks[tp.__class__].pop()

    def data(self, text: str):
        self._tag_stack[-1].process(self, text)

    def close(self):
        tp = self._tag_stack.pop()
        self._tag_type_stacks[tp.__class__].pop()

    def get_parent_processor_by_type(self, cls: Type[_T_BodyTagProcessor]) -> Optional[_T_BodyTagProcessor]:
        if type_stack := self._tag_type_stacks.get(cls):
            return type_stack[-1]
