#  Copyright (c) 2026 Tobias Erbsland - https://erbsland.dev
#  SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import re

from docutils import nodes
from docutils.parsers import rst
from docutils.parsers.rst.directives import single_char_or_unicode

from erbsland.sphinx.ansi.attribute import ANSIAttribute
from erbsland.sphinx.ansi.definition import definition_from_ansi_code

DEFAULT_THEME = "erbsland-ansi"


class ANSILiteralBlock(nodes.literal_block):
    """The literal_block node, for ANSI color codes."""

    pass


class ANSICodeParser(object):
    """Either remove ANSI formatting from a block, or colorize it for HTML output."""

    RE_ANSI_CSI_SEQUENCE: re.Pattern[str] = re.compile(r"\x1b\[([0-?]*)([ -/]*)([@-~])")

    def __call__(self, app, doctree, docname):
        if app.builder.name != "html":
            for ansi_block in doctree.traverse(ANSILiteralBlock):
                self._remove_ansi_formatting(ansi_block)
        else:
            for ansi_block in doctree.traverse(ANSILiteralBlock):
                self._colorize_block_contents(ansi_block)

    def _remove_ansi_formatting(self, block: ANSILiteralBlock):
        cleaned_text = self.RE_ANSI_CSI_SEQUENCE.sub("", block.rawsource)
        block.replace_self(nodes.literal_block(cleaned_text, cleaned_text))

    def _colorize_block_contents(self, block: ANSILiteralBlock):
        theme = block.get("ansi_theme", DEFAULT_THEME)
        new_literal = nodes.literal_block(block.rawsource, classes=[f"{theme}-block"])
        block.replace_self(new_literal)
        current_attributes: dict[ANSIAttribute, str] = {}
        last_end = 0
        nodes_with_formatting = []
        for match in self.RE_ANSI_CSI_SEQUENCE.finditer(block.rawsource):
            head = block.rawsource[last_end : match.start()]
            if head:
                nodes_with_formatting.append(self._create_formatting_node(head, current_attributes, theme))
            self._apply_csi_sequence(match.group(1), match.group(3), current_attributes)
            last_end = match.end()
        tail = block.rawsource[last_end:]

        if tail:
            nodes_with_formatting.append(self._create_formatting_node(tail, current_attributes, theme))
        new_literal.extend(nodes_with_formatting)

    def _create_formatting_node(
        self,
        text: str,
        current_attributes: dict[ANSIAttribute, str],
        theme: str,
    ):
        if current_attributes:
            classes = list([f"{theme}-{attr.to_class_name(value)}" for attr, value in current_attributes.items()])
            return nodes.inline(text=text, classes=classes)
        return nodes.Text(text)

    def _update_attributes(self, code: int, attributes: dict[ANSIAttribute, str]):
        definition = definition_from_ansi_code(code)
        if definition is None:
            return
        attributes_to_set = []
        if isinstance(definition.attribute, ANSIAttribute):
            attributes_to_set = [definition.attribute]
        elif isinstance(definition.attribute, list):
            attributes_to_set = definition.attribute
        for attr in attributes_to_set:
            if definition.value:
                attributes[attr] = definition.value
            elif attr in attributes:
                del attributes[attr]

    def _apply_csi_sequence(self, parameters: str, final_byte: str, attributes: dict[ANSIAttribute, str]):
        if final_byte != "m":
            return
        for code in self._parse_sgr_parameters(parameters):
            self._update_attributes(code, attributes)

    @staticmethod
    def _parse_sgr_parameters(parameters: str) -> list[int]:
        if not parameters:
            return [0]
        codes = []
        for parameter in parameters.split(";"):
            if not parameter:
                codes.append(0)
                continue
            if not parameter.isdigit():
                return []
            codes.append(int(parameter))
        return codes


class ANSIBlockDirective(rst.Directive):
    """
    A directive to include ANSI formatted output as a literal block

    The parameter ``escape-char`` can be used to replace the escape character with a different character.
    The parameter ``theme`` can be used to replace the CSS class prefix.
    """

    has_content = True
    option_spec = {
        "escape-char": single_char_or_unicode,
        "theme": rst.directives.unchanged,
    }

    def run(self):
        text = "\n".join(self.content)
        if "escape-char" in self.options:
            text = text.replace(self.options["escape-char"], "\x1b")
        block = ANSILiteralBlock(text, text)
        if "theme" in self.options:
            block["ansi_theme"] = self.options["theme"]
        return [block]
