from __future__ import annotations

import logging
import re
import warnings
from copy import deepcopy
from enum import Enum
from html import unescape
from io import BytesIO
from pathlib import Path
from typing import Literal, Optional, Union, cast

from docling_core.types.doc import (
    DocItemLabel,
    DoclingDocument,
    DocumentOrigin,
    Formatting,
    ImageRef,
    ListItem,
    NodeItem,
    TableCell,
    TableData,
    TextItem,
)
from pydantic import AnyUrl, BaseModel, Field, TypeAdapter
from typing_extensions import Annotated, override

from docling.backend.abstract_backend import (
    DeclarativeDocumentBackend,
)
from docling.backend.html_backend import HTMLDocumentBackend
from docling.backend.utils.image_resource_loader import ImageResourceLoader
from docling.datamodel.backend_options import (
    HTMLBackendOptions,
    MarkdownBackendOptions,
)
from docling.datamodel.base_models import InputFormat
from docling.datamodel.document import InputDocument
from docling.exceptions import DocumentLoadError

# marko is only installed by the `format-markdown` extra, but DocumentConverter
# imports every backend eagerly. Importing it at module load would therefore
# break `import docling` on installs that omit the extra (the slim packages in
# particular). Guard the imports like the opendocument and xbrl backends do, and
# surface the failure only when Markdown is actually parsed.
# See https://github.com/docling-project/docling/issues/3613.
_MARKO_AVAILABLE: bool = False
_MARKO_IMPORT_ERROR: ImportError | None = None
try:  # pragma: no cover - import-time guard
    import marko
    import marko.element
    import marko.inline
    from marko import Markdown

    _MARKO_AVAILABLE = True
except ImportError as e:  # pragma: no cover - import-time guard
    _MARKO_IMPORT_ERROR = e

_log = logging.getLogger(__name__)

_INSTALL_HINT = (
    "The 'marko' package is required to process Markdown files. "
    "Install it with `pip install 'docling[format-markdown]'`."
)

_MARKER_BODY = "DOCLING_DOC_MD_HTML_EXPORT"
_START_MARKER = f"#_#_{_MARKER_BODY}_START_#_#"
_STOP_MARKER = f"#_#_{_MARKER_BODY}_STOP_#_#"


class _PendingCreationType(str, Enum):
    """CoordOrigin."""

    HEADING = "heading"
    LIST_ITEM = "list_item"


class _HeadingCreationPayload(BaseModel):
    kind: Literal["heading"] = "heading"
    level: int


class _ListItemCreationPayload(BaseModel):
    kind: Literal["list_item"] = "list_item"
    enumerated: bool


_CreationPayload = Annotated[
    Union[
        _HeadingCreationPayload,
        _ListItemCreationPayload,
    ],
    Field(discriminator="kind"),
]


class MarkdownDocumentBackend(DeclarativeDocumentBackend):
    _ENTITY_RE = re.compile(r"&(#\d+|#x[0-9a-fA-F]+|\w+);")

    @staticmethod
    def _unescape_except_pipe(text: str) -> str:
        def replace(match):
            entity = match.group(0)

            # entities that represent |
            if entity in ("&#124;", "&#x7C;", "&vert;"):
                return entity

            return unescape(entity)

        return MarkdownDocumentBackend._ENTITY_RE.sub(replace, text)

    def _shorten_underscore_sequences(self, markdown_text: str, max_length: int = 10):
        # This regex will match any sequence of underscores
        pattern = r"_+"

        def replace_match(match):
            underscore_sequence = match.group(
                0
            )  # Get the full match (sequence of underscores)

            # Shorten the sequence if it exceeds max_length
            if len(underscore_sequence) > max_length:
                return "_" * max_length
            else:
                return underscore_sequence  # Leave it unchanged if it is shorter or equal to max_length

        # Use re.sub to replace long underscore sequences
        shortened_text = re.sub(pattern, replace_match, markdown_text)

        if len(shortened_text) != len(markdown_text):
            warnings.warn("Detected potentially incorrect Markdown, correcting...")

        return shortened_text

    def _shorten_leading_dash_sequences(
        self, markdown_text: str, max_length: int = 10
    ) -> str:
        pattern = re.compile(
            rf"^([ \t]*)(?:-\s+){{{max_length + 1},}}-?(?=\S)", re.MULTILINE
        )
        shortened_text, count = pattern.subn(r"\1- ", markdown_text)

        if count > 0:
            warnings.warn("Detected potentially incorrect Markdown, correcting...")

        return shortened_text

    @override
    def __init__(
        self,
        in_doc: InputDocument,
        path_or_stream: Union[BytesIO, Path],
        options: Optional[MarkdownBackendOptions] = None,
    ):
        # Raised first so a missing optional dependency gives an actionable
        # message rather than a NameError when marko is dereferenced below.
        if not _MARKO_AVAILABLE:
            raise ImportError(_INSTALL_HINT) from _MARKO_IMPORT_ERROR
        if options is None:
            options = MarkdownBackendOptions()
        super().__init__(in_doc, path_or_stream, options)

        _log.debug("Starting MarkdownDocumentBackend...")

        # Markdown file:
        self.path_or_stream = path_or_stream
        self.valid = True
        self.markdown = ""  # To store original Markdown string

        self.in_table = False
        self.md_table_buffer: list[str] = []
        self._html_blocks: int = 0
        self._image_loader: Optional[ImageResourceLoader] = None

        try:
            if isinstance(self.path_or_stream, BytesIO):
                text_stream = self.path_or_stream.getvalue().decode("utf-8")
                # remove invalid sequences
                # very long sequences of underscores will lead to unnecessary long processing times.
                # In any proper Markdown files, underscores have to be escaped,
                # otherwise they represent emphasis (bold or italic)
                self.markdown = self._shorten_underscore_sequences(text_stream)
                self.markdown = self._shorten_leading_dash_sequences(self.markdown)
            if isinstance(self.path_or_stream, Path):
                with open(self.path_or_stream, encoding="utf-8") as f:
                    md_content = f.read()
                    # remove invalid sequences
                    # very long sequences of underscores will lead to unnecessary long processing times.
                    # In any proper Markdown files, underscores have to be escaped,
                    # otherwise they represent emphasis (bold or italic)
                    self.markdown = self._shorten_underscore_sequences(md_content)
                    self.markdown = self._shorten_leading_dash_sequences(self.markdown)
            self.valid = True

            _log.debug(self.markdown)
        except Exception as e:
            raise DocumentLoadError(
                f"Could not initialize MD backend for file with hash {self.document_hash}."
            ) from e
        return

    def _close_table(self, doc: DoclingDocument):
        if self.in_table:
            _log.debug("=== TABLE START ===")
            for md_table_row in self.md_table_buffer:
                _log.debug(md_table_row)
            _log.debug("=== TABLE END ===")
            tcells: list[TableCell] = []
            result_table = []
            for n, md_table_row in enumerate(self.md_table_buffer):
                data = []
                if n == 0:
                    header = [t.strip() for t in md_table_row.split("|")[1:-1]]
                    for value in header:
                        data.append(value)
                    result_table.append(data)
                if n > 1:
                    values = [t.strip() for t in md_table_row.split("|")[1:-1]]
                    for value in values:
                        data.append(value)
                    result_table.append(data)

            for trow_ind, trow in enumerate(result_table):
                for tcol_ind, cellval in enumerate(trow):
                    row_span = (
                        1  # currently supporting just simple tables (without spans)
                    )
                    col_span = (
                        1  # currently supporting just simple tables (without spans)
                    )
                    icell = TableCell(
                        text=unescape(cellval.strip()),
                        row_span=row_span,
                        col_span=col_span,
                        start_row_offset_idx=trow_ind,
                        end_row_offset_idx=trow_ind + row_span,
                        start_col_offset_idx=tcol_ind,
                        end_col_offset_idx=tcol_ind + col_span,
                        column_header=trow_ind == 0,
                        row_header=False,
                    )
                    tcells.append(icell)

            num_rows = len(result_table)
            num_cols = len(result_table[0])
            self.in_table = False
            self.md_table_buffer = []  # clean table markdown buffer
            # Initialize Docling TableData
            table_data = TableData(
                num_rows=num_rows, num_cols=num_cols, table_cells=tcells
            )
            # Populate
            for tcell in tcells:
                table_data.table_cells.append(tcell)
            if len(tcells) > 0:
                doc.add_table(data=table_data)
        return

    def _create_list_item(
        self,
        doc: DoclingDocument,
        parent_item: Optional[NodeItem],
        text: str,
        enumerated: bool,
        formatting: Optional[Formatting] = None,
        hyperlink: Optional[Union[AnyUrl, Path]] = None,
    ):
        item = doc.add_list_item(
            text=text,
            enumerated=enumerated,
            parent=parent_item,
            formatting=formatting,
            hyperlink=hyperlink,
        )
        return item

    def _create_heading_item(
        self,
        doc: DoclingDocument,
        parent_item: Optional[NodeItem],
        text: str,
        level: int,
        formatting: Optional[Formatting] = None,
        hyperlink: Optional[Union[AnyUrl, Path]] = None,
    ):
        if level == 1:
            item = doc.add_title(
                text=text,
                parent=parent_item,
                formatting=formatting,
                hyperlink=hyperlink,
            )
        else:
            item = doc.add_heading(
                text=text,
                level=level - 1,
                parent=parent_item,
                formatting=formatting,
                hyperlink=hyperlink,
            )
        return item

    def _flush_creation_stack(
        self,
        *,
        doc: DoclingDocument,
        creation_stack: list[_CreationPayload],
        snippet_text: str,
        parent_item: Optional[NodeItem],
        list_ordered_flag_by_ref: dict[str, bool],
        list_last_item_by_ref: dict[str, ListItem],
        formatting: Optional[Formatting],
        hyperlink: Optional[Union[AnyUrl, Path]],
    ) -> Optional[NodeItem]:
        """
        Lazily create list items / headings when we first see their inline content.

        Important: Marko list items/headings can contain inline nodes that are NOT RawText
        (e.g. CodeSpan, Link). If we only flush on RawText, pending payloads can leak to
        later nodes and attach to a wrong parent, producing a very deep tree.
        """
        while len(creation_stack) > 0:
            to_create = creation_stack.pop()
            if isinstance(to_create, _ListItemCreationPayload):
                enumerated = (
                    list_ordered_flag_by_ref.get(parent_item.self_ref, False)
                    if parent_item
                    else False
                )
                parent_ref = parent_item.self_ref if parent_item else None
                parent_item = self._create_list_item(
                    doc=doc,
                    parent_item=parent_item,
                    text=snippet_text,
                    enumerated=enumerated,
                    formatting=formatting,
                    hyperlink=hyperlink,
                )
                if parent_ref:
                    list_last_item_by_ref[parent_ref] = cast(ListItem, parent_item)

            elif isinstance(to_create, _HeadingCreationPayload):
                # Not keeping as parent_item as logic for correctly tracking
                # that not implemented yet (section components not captured
                # as heading children in marko)
                self._create_heading_item(
                    doc=doc,
                    parent_item=parent_item,
                    text=snippet_text,
                    level=to_create.level,
                    formatting=formatting,
                    hyperlink=hyperlink,
                )

        return parent_item

    def _iterate_elements(  # noqa: C901
        self,
        *,
        element: marko.element.Element,
        depth: int,
        doc: DoclingDocument,
        visited: set[marko.element.Element],
        creation_stack: list[
            _CreationPayload
        ],  # stack for lazy item creation triggered deep in marko's AST (on RawText)
        list_ordered_flag_by_ref: dict[str, bool],
        list_last_item_by_ref: dict[str, ListItem],
        parent_item: Optional[NodeItem] = None,
        formatting: Optional[Formatting] = None,
        hyperlink: Optional[Union[AnyUrl, Path]] = None,
    ):
        if element in visited:
            return

        # Iterates over all elements in the AST
        # Check for different element types and process relevant details
        if (
            isinstance(element, marko.block.Heading)
            or isinstance(element, marko.block.SetextHeading)
        ) and len(element.children) > 0:
            self._close_table(doc)
            _log.debug(
                " - Heading level %s, content: %s",
                element.level,
                element.children[0].children,  # type: ignore
            )

            if len(element.children) > 1:  # inline group will be created further down
                parent_item = self._create_heading_item(
                    doc=doc,
                    parent_item=parent_item,
                    text="",
                    level=element.level,
                    formatting=formatting,
                    hyperlink=hyperlink,
                )
            else:
                creation_stack.append(_HeadingCreationPayload(level=element.level))

        elif isinstance(element, marko.block.List):
            has_non_empty_list_items = False
            for child in element.children:
                if isinstance(child, marko.block.ListItem) and len(child.children) > 0:
                    has_non_empty_list_items = True
                    break

            self._close_table(doc)
            _log.debug(" - List %s", "ordered" if element.ordered else "unordered")
            if has_non_empty_list_items:
                parent_item = doc.add_list_group(name="list", parent=parent_item)
                list_ordered_flag_by_ref[parent_item.self_ref] = element.ordered

        elif (
            isinstance(element, marko.block.ListItem)
            and len(element.children) > 0
            and isinstance((child := element.children[0]), marko.block.Paragraph)
            and len(child.children) > 0
        ):
            self._close_table(doc)
            _log.debug(" - List item")

            enumerated = (
                list_ordered_flag_by_ref.get(parent_item.self_ref, False)
                if parent_item
                else False
            )
            non_list_children: list[marko.element.Element] = [
                item
                for item in child.children
                if not isinstance(item, marko.block.ListItem)
            ]
            if len(non_list_children) > 1:  # inline group will be created further down
                parent_ref: Optional[str] = (
                    parent_item.self_ref if parent_item else None
                )
                parent_item = self._create_list_item(
                    doc=doc,
                    parent_item=parent_item,
                    text="",
                    enumerated=enumerated,
                    formatting=formatting,
                    hyperlink=hyperlink,
                )
                if parent_ref:
                    list_last_item_by_ref[parent_ref] = cast(ListItem, parent_item)
            else:
                creation_stack.append(_ListItemCreationPayload(enumerated=enumerated))

        elif isinstance(element, marko.inline.Image):
            self._close_table(doc)
            _log.debug(" - Image with alt: %s, url: %s", element.title, element.dest)

            fig_caption: Optional[TextItem] = None
            if element.title is not None and element.title != "":
                title = unescape(element.title)
                fig_caption = doc.add_text(
                    label=DocItemLabel.CAPTION,
                    text=title,
                    formatting=formatting,
                    hyperlink=hyperlink,
                )

            image_ref = self._load_image_ref(element.dest)
            doc.add_picture(parent=parent_item, image=image_ref, caption=fig_caption)

        elif isinstance(element, marko.inline.Emphasis):
            _log.debug(" - Emphasis: %s", element.children)
            formatting = deepcopy(formatting) if formatting else Formatting()
            formatting.italic = True

        elif isinstance(element, marko.inline.StrongEmphasis):
            _log.debug(" - StrongEmphasis: %s", element.children)
            formatting = deepcopy(formatting) if formatting else Formatting()
            formatting.bold = True

        elif isinstance(element, marko.inline.Link):
            _log.debug(" - Link: %s", element.children)
            hyperlink = TypeAdapter(Optional[Union[AnyUrl, Path]]).validate_python(
                element.dest
            )

        elif isinstance(element, marko.inline.RawText | marko.inline.Literal):
            _log.debug(" - RawText/Literal: %s", element.children)
            original_text = (
                element.children if isinstance(element.children, str) else ""
            )
            snippet_text = unescape(original_text.strip())
            is_table_row = "|" in snippet_text and (
                self.in_table or original_text.lstrip().startswith("|")
            )
            if is_table_row:
                self.in_table = True
            if self.in_table and snippet_text:
                snippet_text = self._unescape_except_pipe(original_text.strip())
                # If we're in a table, keep adding text (for formatted content in cells)
                if self.md_table_buffer:
                    self.md_table_buffer[len(self.md_table_buffer) - 1] += snippet_text
                else:
                    self.md_table_buffer.append(snippet_text)
            elif snippet_text:
                # Not in table - close any pending table and process as regular text
                self._close_table(doc)

                if creation_stack:
                    parent_item = self._flush_creation_stack(
                        doc=doc,
                        creation_stack=creation_stack,
                        snippet_text=snippet_text,
                        parent_item=parent_item,
                        list_ordered_flag_by_ref=list_ordered_flag_by_ref,
                        list_last_item_by_ref=list_last_item_by_ref,
                        formatting=formatting,
                        hyperlink=hyperlink,
                    )
                else:
                    doc.add_text(
                        label=DocItemLabel.TEXT,
                        parent=parent_item,
                        text=snippet_text,
                        formatting=formatting,
                        hyperlink=hyperlink,
                    )

        elif isinstance(element, marko.inline.CodeSpan):
            self._close_table(doc)
            _log.debug(" - Code Span: %s", element.children)
            snippet_text = str(element.children).strip()
            # If this CodeSpan is the only content of a list item / heading, Marko won't
            # emit RawText. Flush pending creations here to avoid leaking payloads.
            if creation_stack and snippet_text:
                parent_item = self._flush_creation_stack(
                    doc=doc,
                    creation_stack=creation_stack,
                    snippet_text=snippet_text,
                    parent_item=parent_item,
                    list_ordered_flag_by_ref=list_ordered_flag_by_ref,
                    list_last_item_by_ref=list_last_item_by_ref,
                    formatting=formatting,
                    hyperlink=hyperlink,
                )
                # Represent CodeSpan as the container's text; avoid adding a duplicate CodeItem.
                return
            doc.add_code(
                parent=parent_item,
                text=snippet_text,
                formatting=formatting,
                hyperlink=hyperlink,
            )

        elif (
            isinstance(element, marko.block.CodeBlock | marko.block.FencedCode)
            and len(element.children) > 0
            and isinstance((child := element.children[0]), marko.inline.RawText)
            and len(snippet_text := (child.children.strip())) > 0
        ):
            self._close_table(doc)
            _log.debug(" - Code Block: %s", element.children)
            doc.add_code(
                parent=parent_item,
                text=snippet_text,
                formatting=formatting,
                hyperlink=hyperlink,
            )

        elif isinstance(element, marko.inline.LineBreak):
            if self.in_table:
                _log.debug("Line break in a table")
                self.md_table_buffer.append("")

        elif isinstance(element, marko.block.HTMLBlock):
            self._html_blocks += 1
            self._close_table(doc)
            _log.debug("HTML Block: %s", element)
            if (
                len(element.body) > 0
            ):  # If Marko doesn't return any content for HTML block, skip it
                html_block = element.body.strip()

                # wrap in markers to enable post-processing in convert()
                text_to_add = f"{_START_MARKER}{html_block}{_STOP_MARKER}"
                doc.add_code(
                    parent=parent_item,
                    text=text_to_add,
                    formatting=formatting,
                    hyperlink=hyperlink,
                )
        else:
            if not isinstance(element, str):
                self._close_table(doc)
                _log.debug("Some other element: %s", type(element).__name__)

        if (
            isinstance(element, marko.block.Paragraph | marko.block.Heading)
            and len(element.children) > 1
        ):
            parent_item = doc.add_inline_group(parent=parent_item)

        processed_block_types = (
            marko.block.CodeBlock,
            marko.block.FencedCode,
            marko.inline.RawText,
        )

        # Iterate through the element's children (if any)
        if hasattr(element, "children") and not isinstance(
            element, processed_block_types
        ):
            for child in element.children:
                if (
                    isinstance(element, marko.block.ListItem)
                    and isinstance(child, marko.block.List)
                    and parent_item
                    and list_last_item_by_ref.get(parent_item.self_ref, None)
                ):
                    _log.debug(
                        "walking into new List hanging from item of parent list %s",
                        parent_item.self_ref,
                    )
                    parent_item = list_last_item_by_ref[parent_item.self_ref]

                self._iterate_elements(
                    element=child,
                    depth=depth + 1,
                    doc=doc,
                    visited=visited,
                    creation_stack=creation_stack,
                    list_ordered_flag_by_ref=list_ordered_flag_by_ref,
                    list_last_item_by_ref=list_last_item_by_ref,
                    parent_item=parent_item,
                    formatting=formatting,
                    hyperlink=hyperlink,
                )

    def _get_image_loader(self) -> ImageResourceLoader:
        """Lazily build the shared image-resource loader.

        Resolving and decoding image sources (``data:`` URIs, local files, and
        remote URLs) together with the relevant safety limits is shared with the
        HTML backend through :class:`ImageResourceLoader`, so that logic is not
        duplicated here.
        """
        if self._image_loader is None:
            md_options = cast(MarkdownBackendOptions, self.options)
            self._image_loader = ImageResourceLoader(
                enable_local_fetch=md_options.enable_local_fetch,
                enable_remote_fetch=md_options.enable_remote_fetch,
                max_image_data_base64_bytes=md_options.max_image_data_base64_bytes,
            )
        return self._image_loader

    def _load_image_ref(self, dest: str) -> Optional[ImageRef]:
        """Resolve and decode a Markdown image source into an ``ImageRef``.

        Returns ``None`` when image loading is disabled, the source is empty, or
        the image cannot be loaded.
        """
        md_options = cast(MarkdownBackendOptions, self.options)
        if not md_options.fetch_images or not dest:
            return None
        base_path = (
            str(md_options.source_uri) if md_options.source_uri is not None else None
        )
        return self._get_image_loader().load_image_ref(dest, base_path)

    def is_valid(self) -> bool:
        return self.valid

    def unload(self):
        if isinstance(self.path_or_stream, BytesIO):
            self.path_or_stream.close()
        self.path_or_stream = None

    @classmethod
    def supports_pagination(cls) -> bool:
        return False

    @classmethod
    def supported_formats(cls) -> set[InputFormat]:
        return {InputFormat.MD}

    def convert(self) -> DoclingDocument:
        _log.debug("converting Markdown...")

        origin = DocumentOrigin(
            filename=self.file.name or "file",
            mimetype="text/markdown",
            binary_hash=self.document_hash,
        )

        doc = DoclingDocument(name=self.file.stem or "file", origin=origin)

        if self.is_valid():
            # Parse the markdown into an abstract syntax tree (AST)
            marko_parser = Markdown()
            parsed_ast = marko_parser.parse(self.markdown)
            # Start iterating from the root of the AST
            self._iterate_elements(
                element=parsed_ast,
                depth=0,
                doc=doc,
                parent_item=None,
                visited=set(),
                creation_stack=[],
                list_ordered_flag_by_ref={},
                list_last_item_by_ref={},
            )
            self._close_table(doc=doc)  # handle any last hanging table

            # if HTML blocks were detected, export to HTML and delegate to HTML backend
            if self._html_blocks > 0:
                # export to HTML
                html_backend_cls = HTMLDocumentBackend
                html_str = doc.export_to_html()

                def _restore_original_html(txt, regex):
                    _txt, count = re.subn(regex, "", txt)
                    if count != self._html_blocks:
                        raise RuntimeError(
                            "An internal error has occurred during Markdown conversion."
                        )
                    return _txt

                # restore original HTML by removing previously added markers
                for regex in [
                    rf"<pre>\s*<code>\s*{_START_MARKER}",
                    rf"{_STOP_MARKER}\s*</code>\s*</pre>",
                ]:
                    html_str = _restore_original_html(txt=html_str, regex=regex)
                self._html_blocks = 0
                # delegate to HTML backend
                stream = BytesIO(bytes(html_str, encoding="utf-8"))
                md_options = cast(MarkdownBackendOptions, self.options)
                html_options = HTMLBackendOptions(
                    enable_local_fetch=md_options.enable_local_fetch,
                    enable_remote_fetch=md_options.enable_remote_fetch,
                    fetch_images=md_options.fetch_images,
                    source_uri=md_options.source_uri,
                    infer_furniture=False,
                    add_title=False,
                )
                in_doc = InputDocument(
                    path_or_stream=stream,
                    format=InputFormat.HTML,
                    backend=html_backend_cls,
                    filename=self.file.name,
                    backend_options=html_options,
                )
                html_backend_obj = html_backend_cls(
                    in_doc=in_doc,
                    path_or_stream=stream,
                    options=html_options,
                )
                doc = html_backend_obj.convert()
        else:
            raise RuntimeError(
                f"Cannot convert md with {self.document_hash} because the backend failed to init."
            )
        return doc
