#
# Copyright 2023 Anders Kaplan
#
# This file is part of the Translate Toolkit.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.

"""
Module for parsing Markdown files for translation.

The principles for extraction of translation units are as follows:

1. Extract all content relevant for translation, at the cost of also
   including some formatting.
2. One translation unit per paragraph.
3. Keep formatting out of the translation units as much as possible.
   Exceptions include *phrase emphasis* and `inline code`.
   Use placeholders {1}, {2}, ..., as needed.
4. Avoid HTML entities in the translation units. Use Unicode
   equivalents if possible.

White space within translation units is normalized, because the PO format does
not preserve white space, and the translated Markdown content may have
to be reflowed anyway.
"""

from __future__ import annotations

import re
from itertools import chain
from typing import TYPE_CHECKING

import mistletoe.token
from mistletoe import block_token, span_token
from mistletoe.markdown_renderer import (
    Fragment,
    LinkReferenceDefinition,
    LinkReferenceDefinitionBlock,
    MarkdownRenderer,
)

from translate.storage import base

if TYPE_CHECKING:
    from collections.abc import Callable, Iterable


class MarkdownUnit(base.TranslationUnit):
    """A unit of translatable/localisable markdown content."""

    def __init__(self, source=None) -> None:
        super().__init__(source)
        self.locations = []

    def addlocation(self, location) -> None:
        self.locations.append(location)

    def getlocations(self):
        return self.locations


class MarkdownFile(base.TranslationStore[MarkdownUnit]):
    UnitClass = MarkdownUnit

    def __init__(
        self,
        inputfile=None,
        callback=None,
        max_line_length=None,
        extract_code_blocks=True,
        extract_frontmatter=True,
    ) -> None:
        """
        Construct a new object instance.

        :param inputfile: if specified, the content of this file is read and parsed.
        :param callback: a function which takes a chunk of untranslated content as
          input and returns the corresponding translated content. Defaults to
          a no-op.
        :param max_line_length: if specified, the document is word wrapped to the
          given line length when rendered.
        :param extract_code_blocks: if True (default), code blocks are extracted
          for translation. If False, code blocks are left as-is.
        :param extract_frontmatter: if True (default), front matter is extracted
          for translation. If False, it is preserved as-is.
        """
        base.TranslationStore.__init__(self)
        self.filename = getattr(inputfile, "name", None)
        self.callback = callback or self._dummy_callback
        self.max_line_length = max_line_length
        self.extract_code_blocks = extract_code_blocks
        self.extract_frontmatter = extract_frontmatter
        self.filesrc = ""
        if inputfile is not None:
            md_src = inputfile.read()
            inputfile.close()
            self.parse(md_src)

    def parse(self, data) -> None:
        """Process the given source string (binary)."""
        lines = data.decode().splitlines(keepends=False)
        front_matter_end = 0
        front_matter = ""
        has_front_matter = False
        for line_no, line in enumerate(lines):
            if not has_front_matter:
                if line and not line.startswith("---"):
                    # No front matter found
                    break
                has_front_matter = True
            elif line.startswith(("---", "...")):
                # End of front matter
                front_matter_end = line_no
                break

        if front_matter_end:
            # Include trailing space in the front matter
            if front_matter_end + 1 < len(lines) and (
                not lines[front_matter_end + 1] or lines[front_matter_end + 1].isspace()
            ):
                front_matter_end += 1
            front_matter = "\n".join(chain(lines[: front_matter_end + 1], [""]))
            if self.extract_frontmatter:
                # Keep front matter as a normal translation unit.
                unit = self.UnitClass(front_matter)
                if self.filename:
                    unit.addlocation(f"{self.filename}:1")
                unit.setdocpath("frontmatter[1]")
                self.addunit(unit)
                front_matter = self.callback(front_matter)
            lines = lines[front_matter_end + 1 :]

        with TranslatingMarkdownRenderer(
            self._translate_callback,
            block_token.Table,
            max_line_length=self.max_line_length,
            extract_code_blocks=self.extract_code_blocks,
            lookup_callback=self.callback,
        ) as renderer:
            document = block_token.Document(lines)
            self.filesrc = front_matter + renderer.render(document)

    @staticmethod
    def _dummy_callback(text: str) -> str:
        return text

    def _translate_callback(self, text: str, path: list[str], docpath: str = "") -> str:
        text = text.strip()
        if not text:
            return ""

        # emit a translation unit. The PO store takes care of the escaping.
        unit = self.addsourceunit(text)
        # Index path to avoid duplicate location on list items.
        unit.addlocation(f"{self.filename or ''}{''.join(path[0])}")
        if docpath:
            unit.setdocpath(docpath)

        # return translated text
        return self.callback(text)


class TranslatingMarkdownRenderer(MarkdownRenderer):
    def __init__(
        self,
        translate_callback: Callable[[str, list[str], str], str],
        *extras,
        max_line_length: int | None = None,
        extract_code_blocks: bool = True,
        lookup_callback: Callable[[str], str] | None = None,
    ) -> None:
        super().__init__(*extras, max_line_length=max_line_length)  # ty:ignore[invalid-argument-type]
        self.translate_callback = translate_callback
        self.lookup_callback = lookup_callback
        self.bypass = False
        self.path = []
        self.ignore_translation = False
        self.extract_code_blocks = extract_code_blocks
        # Docpath tracking: heading hierarchy and sibling counts
        # _heading_stack: list of (level, heading_index) for current heading nesting
        self._heading_stack: list[tuple[int, int]] = []
        # _section_counts: dict mapping element type to count at current section level
        self._section_counts: list[dict[str, int]] = [{}]
        # _container_stack: list of (container_type, index) for nested containers
        self._container_stack: list[tuple[str, int]] = []
        # _current_docpath: the current docpath string for the current block
        self._current_docpath = ""

    def render(self, token: mistletoe.token.Token) -> str:
        try:
            # set the root node here, because the rendering also involves some parsing,
            # and the parsing requires a valid root node.
            mistletoe.token._root_node = token  # ty:ignore[invalid-assignment]
            return super().render(token)
        finally:
            mistletoe.token._root_node = None

    # rendering of span tokens:
    # override to inject placeholders and translate content

    def _build_docpath(self, element_type: str) -> str:
        """Build the current docpath string for a block element."""
        parts = [f"h{level}[{idx}]" for level, idx in self._heading_stack]
        for container_type, container_idx in self._container_stack:
            parts.append(f"{container_type}[{container_idx}]")
        counts = self._section_counts[-1]
        counts[element_type] = counts.get(element_type, 0) + 1
        parts.append(f"{element_type}[{counts[element_type]}]")
        return "/".join(parts)

    def _enter_heading(self, level: int) -> str:
        """Update heading stack and return docpath for this heading."""
        # Pop headings of equal or greater level
        while self._heading_stack and self._heading_stack[-1][0] >= level:
            self._heading_stack.pop()
            self._section_counts.pop()
        # Count this heading at the current section level
        counts = self._section_counts[-1]
        heading_key = f"h{level}"
        counts[heading_key] = counts.get(heading_key, 0) + 1
        idx = counts[heading_key]
        # Push new section level
        self._heading_stack.append((level, idx))
        self._section_counts.append({})
        # Build path
        parts = [f"h{level}[{i}]" for level, i in self._heading_stack]
        for container_type, container_idx in self._container_stack:
            parts.append(f"{container_type}[{container_idx}]")
        return "/".join(parts)

    def _enter_container(self, container_type: str) -> None:
        """Push a container (list item, blockquote) onto the nesting stack."""
        counts = self._section_counts[-1]
        counts[container_type] = counts.get(container_type, 0) + 1
        self._container_stack.append((container_type, counts[container_type]))
        self._section_counts.append({})

    def _exit_container(self) -> None:
        """Pop a container from the nesting stack."""
        if self._container_stack:
            self._container_stack.pop()
            self._section_counts.pop()

    _leading_ws = re.compile(r"^(\s+)\S")
    _trailing_ws = re.compile(r"\S(\s+)$")

    def render_raw_text(self, token: span_token.RawText) -> Iterable[Fragment]:
        # make separate fragments for leading and trailing space, if applicable
        match = self._leading_ws.search(token.content)
        if match:
            yield Fragment(match.group(1), wordwrap=True)
            token.content = token.content.lstrip()
        match = self._trailing_ws.search(token.content)
        if match:
            yield Fragment(token.content.rstrip(), wordwrap=True)
            yield Fragment(match.group(1), wordwrap=True)
        else:
            yield Fragment(token.content, wordwrap=True)

    def render_auto_link(self, token: span_token.AutoLink) -> Iterable[Fragment]:
        # replace with placeholder, unless in bypass mode
        if self.bypass:
            yield from super().render_auto_link(token)
            return

        yield Fragment(None, placeholder_content=super().render_auto_link(token))  # ty:ignore[invalid-argument-type]

    def render_line_break(self, token: span_token.LineBreak) -> Iterable[Fragment]:
        if self.bypass:
            yield from super().render_line_break(token)
            return

        yield Fragment("\n", wordwrap=token.soft, hard_line_break=not token.soft)

    def render_html_span(self, token: span_token.HTMLSpan) -> Iterable[Fragment]:
        # replace with placeholder, unless in bypass mode
        if self.bypass:
            yield from super().render_html_span(token)
            return

        yield Fragment(None, placeholder_content=super().render_html_span(token))  # ty:ignore[invalid-argument-type]

    def render_link_or_image(
        self, token: span_token.SpanToken, target: str
    ) -> Iterable[Fragment]:
        if self.bypass:
            yield from super().render_link_or_image(token, target)
            return

        yield from self.embed_span(Fragment("["), token.children, Fragment("]"))  # ty:ignore[invalid-argument-type]

        if token.dest_type in {"uri", "angle_uri"}:  # ty:ignore[unresolved-attribute]
            # Markdown link format: "[" description "](" dest_part [" " title] ")"
            dest_part = f"<{target}>" if token.dest_type == "angle_uri" else target  # ty:ignore[unresolved-attribute]
            placeholder = Fragment(None, important=True)  # ty:ignore[invalid-argument-type]
            placeholder.placeholder_content = [  # ty:ignore[unresolved-attribute]
                Fragment("("),
                Fragment(dest_part),
            ]
            if token.title:  # ty:ignore[unresolved-attribute]
                translated_title = self.translate_callback(
                    token.title,  # ty:ignore[unresolved-attribute]
                    [*self.path, "link-title"],
                    self._current_docpath,
                )
                placeholder.placeholder_content.extend(  # ty:ignore[unresolved-attribute]
                    [
                        Fragment(" ", wordwrap=True),
                        Fragment(token.title_delimiter),  # ty:ignore[unresolved-attribute]
                        Fragment(translated_title, wordwrap=True),
                        Fragment(
                            ")"
                            if token.title_delimiter == "("  # ty:ignore[unresolved-attribute]
                            else token.title_delimiter  # ty:ignore[unresolved-attribute]
                        ),
                    ]
                )
            placeholder.placeholder_content.append(Fragment(")"))  # ty:ignore[unresolved-attribute]
            yield placeholder
        elif token.dest_type == "full":  # ty:ignore[unresolved-attribute]
            # Markdown link format: "[" description "][" label "]"
            translated_label = self.translate_callback(
                token.label,  # ty:ignore[unresolved-attribute]
                [*self.path, "link-label"],
                self._current_docpath,
            )
            placeholder = Fragment(None, important=True)  # ty:ignore[invalid-argument-type]
            placeholder.placeholder_content = [  # ty:ignore[unresolved-attribute]
                Fragment("["),
                Fragment(translated_label, wordwrap=True),
                Fragment("]"),
            ]
            yield placeholder
        elif token.dest_type == "collapsed":  # ty:ignore[unresolved-attribute]
            # Markdown link format: "[" description "][]"
            yield Fragment("[]")
        else:
            # Markdown link format: "[" description "]"
            pass

    def render_link_reference_definition(
        self, token: LinkReferenceDefinition
    ) -> Iterable[Fragment]:
        # note: these tokens will never be encountered in bypass mode.
        # Translate label and title unless we're in an ignore section
        if self.ignore_translation:
            label = token.label
            title = token.title
        else:
            label = self.translate_callback(
                token.label, [*self.path, "link-label"], self._current_docpath
            )
            title = (
                self.translate_callback(
                    token.title, [*self.path, "link-title"], self._current_docpath
                )
                if token.title
                else None
            )

        placeholder = Fragment(None)  # ty:ignore[invalid-argument-type]
        placeholder.placeholder_content = [  # ty:ignore[unresolved-attribute]
            Fragment("["),
            Fragment(label, wordwrap=True),
            Fragment("]: ", wordwrap=True),
            Fragment(
                f"<{token.dest}>" if token.dest_type == "angle_uri" else token.dest
            ),
        ]
        if title:
            placeholder.placeholder_content.extend(  # ty:ignore[unresolved-attribute]
                [
                    Fragment(" ", wordwrap=True),
                    Fragment(token.title_delimiter),
                    Fragment(title, wordwrap=True),
                    Fragment(
                        ")" if token.title_delimiter == "(" else token.title_delimiter
                    ),
                ]
            )
        yield placeholder

    # rendering of block tokens:
    # override to keep track of the content path

    def render_heading(
        self, token: block_token.Heading, max_line_length: int
    ) -> Iterable[str]:
        self.path.append(f":{token.line_number}")  # ty:ignore[unresolved-attribute]
        self._current_docpath = self._enter_heading(token.level)
        content = list(super().render_heading(token, max_line_length=max_line_length))
        self.path.pop()
        return content

    def render_setext_heading(
        self, token: block_token.SetextHeading, max_line_length: int
    ) -> Iterable[str]:
        self.path.append(f":{token.line_number}")  # ty:ignore[unresolved-attribute]
        self._current_docpath = self._enter_heading(token.level)
        content = list(
            super().render_setext_heading(token, max_line_length=max_line_length)
        )
        self.path.pop()
        return content

    def render_quote(
        self, token: block_token.Quote, max_line_length: int
    ) -> Iterable[str]:
        self.path.append(f":{token.line_number}")  # ty:ignore[unresolved-attribute]
        self._enter_container("blockquote")
        content = list(super().render_quote(token, max_line_length=max_line_length))
        self._exit_container()
        self.path.pop()
        return content

    def render_paragraph(
        self, token: block_token.Paragraph, max_line_length: int
    ) -> Iterable[str]:
        self.path.append(f":{token.line_number}")  # ty:ignore[unresolved-attribute]
        self._current_docpath = self._build_docpath("p")
        content = list(super().render_paragraph(token, max_line_length=max_line_length))
        self.path.pop()
        return content

    def render_html_block(
        self, token: block_token.HtmlBlock, max_line_length: int
    ) -> Iterable[str]:
        # Check if this is a translation control comment
        content = token.content.strip()
        if content == "<!-- translate:off -->":
            self.ignore_translation = True
        elif content == "<!-- translate:on -->":
            self.ignore_translation = False

        # Return the raw HTML block content
        return super().render_html_block(token, max_line_length=max_line_length)

    def render_block_code(
        self, token: block_token.BlockCode, max_line_length: int
    ) -> Iterable[str]:
        if not self.extract_code_blocks or self.ignore_translation:
            return super().render_block_code(token, max_line_length=max_line_length)

        self.path.append(f":{token.line_number}")  # ty:ignore[unresolved-attribute]
        self._current_docpath = self._build_docpath("code")
        code_content = token.content[:-1]  # strip trailing \n
        translated = self.translate_callback(
            code_content, self.path, self._current_docpath
        )
        lines = translated.split("\n")
        self.path.pop()
        return self.prefix_lines(lines, "    ")

    def render_fenced_code_block(  # ty:ignore[invalid-method-override]
        self,
        token: block_token.CodeFence,
        max_line_length: int,
    ) -> Iterable[str]:
        if not self.extract_code_blocks or self.ignore_translation:
            return super().render_fenced_code_block(
                token,  # ty:ignore[invalid-argument-type]
                max_line_length=max_line_length,
            )

        self.path.append(f":{token.line_number}")  # ty:ignore[unresolved-attribute]
        self._current_docpath = self._build_docpath("code")
        code_content = token.content[:-1]  # strip trailing \n
        translated = self.translate_callback(
            code_content, self.path, self._current_docpath
        )
        indentation = " " * token.indentation
        result = [indentation + token.delimiter + token.info_string]
        result.extend(self.prefix_lines(translated.split("\n"), indentation))
        result.append(indentation + token.delimiter)
        self.path.pop()
        return result

    def render_list_item(
        self, token: block_token.ListItem, max_line_length: int
    ) -> Iterable[str]:
        self.path.append(f":{token.line_number}")
        self._enter_container("li")
        content = list(super().render_list_item(token, max_line_length=max_line_length))
        self._exit_container()
        self.path.pop()
        return content

    def render_table(
        self, token: block_token.Table, max_line_length: int
    ) -> Iterable[str]:
        self.path.append(f":{token.line_number}")  # ty:ignore[unresolved-attribute]
        self._table_docpath = self._build_docpath("table")
        self._table_row_index = 0
        content = list(super().render_table(token, max_line_length=max_line_length))
        self.path.pop()
        return content

    def table_row_to_text(self, row) -> list[str]:
        """Render each table cell with a unique docpath."""
        self._table_row_index += 1
        result = []
        for col_index, col in enumerate(row.children):
            self._current_docpath = (
                f"{self._table_docpath}/r[{self._table_row_index}]/c[{col_index + 1}]"
            )
            result.append(
                next(self.span_to_lines(col.children, max_line_length=None), "")  # ty:ignore[invalid-argument-type]
            )
        return result

    def render_link_reference_definition_block(
        self, token: LinkReferenceDefinitionBlock, max_line_length: int
    ) -> Iterable[str]:
        self.path.append(f":{token.line_number}")  # ty:ignore[unresolved-attribute]
        self._current_docpath = self._build_docpath("link-ref-def")
        content = list(
            super().render_link_reference_definition_block(
                token, max_line_length=max_line_length
            )
        )
        self.path.pop()
        return content

    # translation and placeholder functions

    def span_to_lines(
        self, tokens: Iterable[span_token.SpanToken], max_line_length: int
    ) -> Iterable[str]:
        """Renders a sequence of span tokens to markdown, with translation."""
        # If we're in an ignore section, skip translation
        if self.ignore_translation:
            original_bypass = self.bypass
            try:
                self.bypass = True
                fragments = self.make_fragments(tokens)
                # Expand placeholders before rendering
                expanded = list(self.expand_placeholders(fragments))
                return super().fragments_to_lines(
                    expanded, max_line_length=max_line_length
                )
            finally:
                self.bypass = original_bypass

        # turn the span into fragments, which may include placeholders.
        # list-ify the iterator because we may need to traverse it more than once
        fragments = list(self.make_fragments(tokens))
        original_bypass = self.bypass
        try:
            self.bypass = True

            # pre-process placeholders
            merged = self.merge_adjacent_placeholders(fragments)
            leader, content, trailer = self.trim_flanking_placeholders(merged)
            placeholders = self.insert_placeholder_markers(content)

            # render the translatable content (with placeholders) to markdown
            content_md = "\n".join(
                self.fragments_to_lines(content, max_line_length=float("inf"))  # ty:ignore[invalid-argument-type]
            )

            # translate and parse into new fragments. handle hard line breaks.
            if content_md:
                translated_md = self.translate_callback(
                    content_md, self.path, self._current_docpath
                )
                # If translation with placeholders didn't match and there are
                # placeholders, try looking up with expanded placeholders
                # (full markdown links, etc.) to support PO files that contain
                # the original markdown syntax instead of placeholder markers.
                # Note: if the translation is intentionally identical to the
                # source, this fallback is harmless — it will also not find a
                # different translation.
                if (
                    translated_md == content_md
                    and placeholders
                    and self.lookup_callback
                ):
                    expanded_content_md = self.remove_placeholder_markers(
                        content_md, list(placeholders)
                    )
                    expanded_translated_md = self.lookup_callback(expanded_content_md)
                    if expanded_translated_md != expanded_content_md:
                        translated_md = expanded_translated_md
                        # Clear placeholders since the expanded translation
                        # already contains full markdown syntax (e.g. links)
                        # and does not need placeholder replacement.
                        placeholders = []
                translated_md = translated_md.replace("\n", "\\\n").strip(" \t")
                translated_md = self.remove_placeholder_markers(
                    translated_md, placeholders
                )
                translated = self.make_fragments(
                    span_token.tokenize_inner(translated_md)
                )
                fragments = chain(leader, translated, trailer)

            # expand placeholders and render into final markdown.
            # list-ify to let all generators run before exiting the try/finally block
            expanded = list(self.expand_placeholders(fragments))
            return super().fragments_to_lines(expanded, max_line_length=max_line_length)
        finally:
            self.bypass = original_bypass

    @classmethod
    def merge_adjacent_placeholders(
        cls, fragments: Iterable[Fragment]
    ) -> Iterable[Fragment]:
        """Replaces sequences of placeholders and whitespace with larger placeholders."""
        fragments = list(fragments)
        start = 0
        while start < len(fragments):
            if getattr(fragments[start], "placeholder_content", None):
                end = 0
                for j in range(start + 1, len(fragments)):
                    if getattr(fragments[j], "placeholder_content", None):
                        end = j + 1
                    elif fragments[j].text.isspace():
                        pass
                    else:
                        break
                if end:
                    chunk = fragments[start:end]
                    placeholder = Fragment(None, placeholder_content=chunk)  # ty:ignore[invalid-argument-type]
                    placeholder.important = any(  # ty:ignore[unresolved-attribute]
                        getattr(fragment, "important", False) for fragment in chunk
                    )
                    fragments[start:end] = (placeholder,)
            start += 1

        return fragments

    @classmethod
    def trim_flanking_placeholders(
        cls, fragments: Iterable[Fragment]
    ) -> tuple[Iterable[Fragment], Iterable[Fragment], Iterable[Fragment]]:
        """
        Splits leading and trailing placeholders and whitespace, and the main
        content, into separate lists. Placeholders marked as important are kept
        with the main content.
        """
        fragments = list(fragments)

        pos = 0
        while pos < len(fragments):
            if getattr(fragments[pos], "placeholder_content", None):
                if getattr(fragments[pos], "important", False):
                    break
            elif not fragments[pos].text.isspace():
                break
            pos += 1

        t = len(fragments)
        while t - 1 >= pos:
            if getattr(fragments[t - 1], "placeholder_content", None):
                if getattr(fragments[t - 1], "important", False):
                    break
            elif not fragments[t - 1].text.isspace():
                break
            t -= 1

        return fragments[:pos], fragments[pos:t], fragments[t:]

    @classmethod
    def insert_placeholder_markers(
        cls, fragments: Iterable[Fragment]
    ) -> Iterable[Fragment]:
        """
        Sets the text of the (top-level) placeholder fragments to "{n}".
        Returns an ordered list of placeholders.
        """
        placeholders = []
        for fragment in fragments:
            content = getattr(fragment, "placeholder_content", None)
            if content:
                placeholders.append(fragment)
                fragment.text = f"{{{len(placeholders)}}}"

        return placeholders

    def remove_placeholder_markers(
        self, markdown: str, placeholders: Iterable[Fragment]
    ) -> str:
        """Replaces placeholder markers in the given markdown with placeholder content."""
        for index, placeholder in enumerate(placeholders):
            content = self.expand_placeholders(placeholder.placeholder_content)  # ty:ignore[unresolved-attribute]
            content_md = "\n".join(
                self.fragments_to_lines(content, max_line_length=float("inf"))  # ty:ignore[invalid-argument-type]
            )
            markdown = markdown.replace(f"{{{index + 1}}}", content_md)

        return markdown

    def expand_placeholders(self, fragments: Iterable[Fragment]) -> Iterable[Fragment]:
        """Expands placeholder fragments, recursively."""
        for fragment in fragments:
            content = getattr(fragment, "placeholder_content", None)
            if content:
                yield from self.expand_placeholders(content)
            else:
                yield fragment
