# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import html
import re
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Any,
    Generic,
    Literal,
    TypeVar,
    cast,
)

# Native to python
from xml.etree.ElementTree import Element, SubElement

# Markdown is a dependency of marimo, as such we utilize it as much as possible
# to parse markdown.
from markdown import Markdown
from markdown.blockparser import BlockParser
from markdown.blockprocessors import BlockProcessor
from markdown.preprocessors import Preprocessor
from markdown.util import HTML_PLACEHOLDER_RE, Registry

# As are extensions
from pymdownx.superfences import (  # type: ignore
    RE_NESTED_FENCE_START,
    SuperFencesCodeExtension,
)

from marimo import _loggers
from marimo._ast.cell import CellConfig
from marimo._ast.names import DEFAULT_CELL_NAME
from marimo._convert.common.format import markdown_to_marimo, sql_to_marimo
from marimo._convert.markdown.flavor import (
    _markdown_import_dialects,
    default_markdown_flavor,
)
from marimo._convert.markdown.flavor.base import (
    CodeCellBlock,
    MarkdownFlavor,
    MarkdownImportContext,
    MarkdownImportDialect,
)
from marimo._dependencies.dependencies import DependencyManager
from marimo._schemas.serialization import (
    AppInstantiation,
    CellDef,
    Header,
    NotebookSerializationV1,
    Violation,
)

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

LOGGER = _loggers.marimo_logger()

MARIMO_MD = "marimo-md"
MARIMO_CODE = "marimo-code"

ConvertKeys = Literal["marimo-ir"]


def backwards_compatible_sanitization(line: str) -> str:
    return line


def extract_attribs(
    line: str, fence_start: re.Match[str] | None = None
) -> dict[str, str]:
    # Extract attributes from the code block.
    # Blocks are expected to be like this:
    # python {.marimo disabled="true"}
    if fence_start is None:
        fence_start = RE_NESTED_FENCE_START.match(line)

    if fence_start:
        # attrs is a bit of a misnomer, matches
        # .python.marimo disabled="true"
        inner = fence_start.group("attrs")
        if inner:
            return {
                key: html.unescape(value)
                for key, value in re.findall(r'(\w+)="([^"]*)"', inner)
            }
    return {}


def _is_code_tag(text: str) -> bool:
    head = text.split("\n")[0].strip()
    legacy_format = bool(re.search(r"\{.*python.*\}", head))
    legacy_format |= bool(re.search(r"\{.*sql.*\}", head))
    if DependencyManager.new_superfences.has_required_version(quiet=True):
        supported_format = bool(re.search(r".*\{.*marimo.*\}", head))
        return legacy_format or supported_format
    return legacy_format


def _get_language(text: str) -> str:
    header = text.split("\n").pop(0)
    match = RE_NESTED_FENCE_START.match(header)
    if match and match.group("lang"):
        return str(match.group("lang"))
    if match and match.group("attrs"):
        attributes = str(match.group("attrs"))
        for language in ("python", "sql", "markdown"):
            if re.search(rf"(?:^|[.\s]){language}(?:[.\s]|$)", attributes):
                return language
    return "python"


def formatted_code_block(
    code: str,
    attributes: dict[str, str] | None = None,
    flavor: MarkdownFlavor | None = None,
) -> str:
    """Wraps code in a fenced code block with marimo attributes."""
    if flavor is None:
        flavor = default_markdown_flavor()
    attributes = dict(attributes or {})
    language = attributes.pop("language", "python")
    return flavor.render_code_cell(CodeCellBlock(code, language, attributes))


def app_config_from_root(root: Element) -> dict[str, Any]:
    """
    Extract app config from the root element.

    This may contains unknown keys.
    """

    # Extract meta data from root attributes.
    config_keys = {
        "title": "app_title",
    }
    config = {
        config_keys[key]: value
        for key, value in root.items()
        if key in config_keys
    }
    # Try to pass on other attributes as is
    config.update({k: v for k, v in root.items() if k not in config_keys})
    # Remove values particular to markdown saves.
    config.pop("marimo-version", None)

    return config


def get_source_from_tag(tag: Element) -> str:
    source = tag.text if tag.text else ""
    if tag.tag == MARIMO_MD:
        # Only check here to allow for empty code blocks.
        if not (source and source.strip()):
            return ""
        source = markdown_to_marimo(source)
    elif tag.attrib.get("language") == "sql":
        source = sql_to_marimo(
            source,
            tag.attrib.get("query", "_df"),
            str(tag.attrib.get("hide_output", "false")).lower() == "true",
            tag.attrib.get("engine", None),
        )
    else:
        assert tag.tag == MARIMO_CODE, f"Unknown tag: {tag.tag}"
    return source


def get_cell_config_from_tag(tag: Element, **defaults: bool) -> CellConfig:
    # Known boolean attributes.
    extracted_attrs: dict[str, bool | int] = {
        **defaults,
        **{
            k: v == "true"
            for k, v in tag.attrib.items()
            if k in ["hide_code", "disabled"]
        },
    }
    # "Column" is not a boolean attribute.
    for int_attr in ["column"]:
        if int_attr in tag.attrib:
            extracted_attrs[int_attr] = int(tag.attrib[int_attr])
    return CellConfig.from_dict(extracted_attrs)


T = TypeVar("T")


# TODO: Consider upstreaming some logic such that this isn't such a terrible
# hack. At some point rewriting / overriding the markdown parser would be a
# better idea than all these little work arounds.
@dataclass
class SafeWrap(Generic[T]):
    inner: T

    def strip(self) -> T:
        return self.inner


def _tree_to_ir(root: Element) -> SafeWrap[NotebookSerializationV1]:
    from marimo._ast.app_config import _AppConfig
    from marimo._ast.parse import NON_MARIMO_MARKDOWN_VIOLATION
    from marimo._utils import yaml
    from marimo._utils.scripts import wrap_script_metadata

    app_config = app_config_from_root(root)
    config_only = _AppConfig.sanitize(app_config)

    sources: list[str] = []
    names: list[str] = []
    cell_config: list[CellConfig] = []
    for child in root:
        names.append(child.get("name", DEFAULT_CELL_NAME))
        # Default to hiding markdown cells.
        cell_config.append(
            get_cell_config_from_tag(child, hide_code=child.tag == MARIMO_MD)
        )
        sources.append(get_source_from_tag(child))

    header_str = root.get("header", None)
    pyproject = root.get("pyproject", None)

    # Collect non-config frontmatter metadata
    frontmatter = {k: v for k, v in app_config.items() if k not in config_only}
    if pyproject:
        frontmatter["pyproject"] = pyproject
    if header_str:
        frontmatter["header"] = header_str

    # Build header: frontmatter YAML for md, or script preamble
    if frontmatter:
        header_value = yaml.dump(frontmatter, sort_keys=False)
    elif pyproject and not header_str:
        header_value = wrap_script_metadata(pyproject)
    elif header_str:
        header_value = header_str
    else:
        header_value = None

    # A markdown file is only recognized as a marimo notebook if it has at
    # least one marimo code cell (a `{python}`/`{sql}`/`{.marimo}` fence) or
    # marimo metadata in its frontmatter. Otherwise it is plain markdown, and
    # lint/fix should leave it untouched (see `is_non_marimo_markdown`).
    is_marimo = "marimo-version" in root.attrib or any(
        child.tag == MARIMO_CODE for child in root
    )
    violations = (
        [] if is_marimo else [Violation(NON_MARIMO_MARKDOWN_VIOLATION)]
    )

    notebook = NotebookSerializationV1(
        app=AppInstantiation(options=config_only),
        cells=[
            CellDef(
                name=name,
                code=source,
                options=config.asdict(),
            )
            for name, source, config in zip(
                names, sources, cell_config, strict=False
            )
        ],
        header=Header(value=header_value) if header_value else None,
        violations=violations,
    )
    return SafeWrap(notebook)


class IdentityParser(Markdown):
    """Leaves markdown unchanged."""

    # Considering how ubiquitous "markdown" is, it's a little surprising the
    # internal structure isn't cleaner/ more modular. This "monkey-patching"
    # is comparable to some of the code in markdown extensions- and given this
    # library has been around since 2004, the internals should be relatively
    # stable.
    output_formats: dict[Literal["identity"], Callable[[Element], str]] = {  # type: ignore[assignment, misc]
        "identity": lambda x: x.text if x.text else "",
    }

    def build_parser(self) -> IdentityParser:
        """
        Creates blank registries as a base.
        """
        self.preprocessors = Registry()
        self.parser = BlockParser(self)
        self.inlinePatterns = Registry()
        self.treeprocessors = Registry()
        self.postprocessors = Registry()
        return self

    def convert(self, text: str) -> str:
        """Override the convert method to return the parsed text.

        Note that evoked by itself, would create an infinite loop, since
        block-parsers will never dequeue the extracted blocks.
        """
        if len(self.parser.blockprocessors) == 0:
            self.parser.blockprocessors.register(
                IdentityProcessor(self.parser), "identity", 1
            )

        return super().convert(text)


class MarimoMdParser(IdentityParser):
    """Parses Markdown to marimo notebook string."""

    meta: dict[str, Any]

    output_formats: dict[ConvertKeys, Callable[[Element], SafeWrap[Any]]] = {  # type: ignore[assignment, misc]
        "marimo-ir": _tree_to_ir,
    }

    def __init__(
        self,
        *args: Any,
        output_format: ConvertKeys = "marimo-ir",
        import_dialects: Sequence[MarkdownImportDialect] = (),
        **kwargs: Any,
    ) -> None:
        super().__init__(
            *args, output_format=cast(Any, output_format), **kwargs
        )
        self.meta = {}
        import_context = MarkdownImportContext()
        # Build here opposed to the parent class since there is intermediate
        # logic after the parser is built, and it is more clear here what is
        # registered.
        self.stripTopLevelTags = False

        # Note: MetaPreprocessor does not properly handle frontmatter yaml, so
        # cleanup occurs in the block-processor.
        self.preprocessors.register(
            FrontMatterPreprocessor(self), "frontmatter", 100
        )
        if import_dialects:
            self.preprocessors.register(
                MarkdownImportDialectPreprocessor(
                    self, import_dialects, import_context
                ),
                "markdown-import-dialects",
                99,
            )
        fences_ext = SuperFencesCodeExtension()
        fences_ext.extendMarkdown(self)
        # TODO: Consider adding the admonition extension, and integrating it
        # with mo.markdown callouts.

        block_processor = ExpandAndClassifyProcessor(self.parser)
        block_processor.stash = fences_ext.stash.stash
        self.parser.blockprocessors.register(
            block_processor, "marimo-processor", 10
        )


class SanitizeParser(IdentityParser):
    """Sanitizes Markdown to non-executable string."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        # Potentially no need for a separate sanitizer. We could use a
        # heuristic to determine if this block should be treated as code, but
        # to catch all edgecases, just run it through the similar superfence
        # logic.
        self.stripTopLevelTags = False

        fences_ext = SuperFencesCodeExtension()
        fences_ext.extendMarkdown(self)

        preprocessor = SanitizeProcessor(self)
        preprocessor.stash = fences_ext.stash.stash
        self.preprocessors.register(preprocessor, "marimo-sanitizer", 1)

        # Add back in identity to dequeue.
        self.parser.blockprocessors.register(
            IdentityProcessor(self.parser), "identity", 1
        )


class FrontMatterPreprocessor(Preprocessor):
    """Preprocessor for to extract YAML front matter.

    The built-in MetaPreprocessor does not handle frontmatter yaml properly, so
    this is a custom implementation.

    Like the built-in MetaPreprocessor, this preprocessor extracts yaml and
    stores it in the Markdown's metadata attribute. Inspired by conversation
    and linked project in github/Python-Markdown/markdown/497. See docdown
    (BSD-3) or python-frontmatter (MIT) for similar implementations.
    """

    def __init__(self, md: MarimoMdParser):
        super().__init__(md)
        self.md: MarimoMdParser = md

    def run(self, lines: list[str]) -> list[str]:
        if not lines:
            return lines

        doc = "\n".join(lines)

        meta, doc = extract_frontmatter(doc)
        if isinstance(meta, dict):
            self.md.meta.update(meta)

        return doc.split("\n")


class MarkdownImportDialectPreprocessor(Preprocessor):
    """Normalize dialect-specific markdown before SuperFences parses it."""

    def __init__(
        self,
        md: MarimoMdParser,
        import_dialects: Sequence[MarkdownImportDialect],
        context: MarkdownImportContext,
    ) -> None:
        super().__init__(md)
        self.md: MarimoMdParser = md
        self.import_dialects = import_dialects
        self.context = context

    def run(self, lines: list[str]) -> list[str]:
        for dialect in self.import_dialects:
            lines = dialect.preprocess(lines, self.context)
        self.md.meta.update(self.context.metadata)
        return lines


class SanitizeProcessor(Preprocessor):
    """Prevent unintended executable code block injection.

    Typically run on Markdown fragments (e.g. cells) to prevent code injection.
    **Note***: Must run after SuperFencesCodeExtension.
    """

    stash: dict[str, Any]

    def run(self, lines: list[str]) -> list[str]:
        # Note, an empty stash is not sufficient since partially open code
        # blocks could be in the text.
        if not lines:
            return lines

        is_code = False
        for i, line in enumerate(lines):
            # Still need to do all replacements
            if HTML_PLACEHOLDER_RE.match(line.strip()):
                lookup = line.strip()[1:-1]
                code = self.stash[lookup][0]
                lines[i] = code
                # This is a tag we would normally parse on.
                # So protect it from being parsed improperly, by just treating
                # it as code.
                is_code = is_code or _is_code_tag(code)
            # We also need to check for code block delimiters that superfences
            # did not catch, as this will break other code blocks.
            is_code = is_code or RE_NESTED_FENCE_START.match(line)

        if not is_code:
            return lines

        return formatted_code_block(
            markdown_to_marimo("\n".join(lines))
        ).split("\n")


class IdentityProcessor(BlockProcessor):
    """Leaves markdown unchanged."""

    def test(*_args: Any) -> bool:
        return True

    def run(self, parent: Element, blocks: list[str]) -> None:
        parent.text = "\n\n".join(blocks)
        blocks.clear()


class ExpandAndClassifyProcessor(BlockProcessor):
    """Separates code blocks and markdown blocks."""

    stash: dict[str, Any]

    def test(*_args: Any) -> bool:
        return True

    def run(self, parent: Element, blocks: list[str]) -> None:
        # Copy app metadata to the parent element.
        assert isinstance(self.parser.md, MarimoMdParser)
        for key, value in self.parser.md.meta.items():
            if isinstance(value, str):
                parent.set(key, value)

        text: list[str] = []

        def add_paragraph() -> None:
            if not text:
                return
            # An additional line break is added before code blocks.
            if text[-1].strip() == "":
                text.pop()
                if not text:
                    return
            paragraph = SubElement(parent, MARIMO_MD)
            paragraph.text = "\n".join(text).strip()
            text.clear()

        # Operate on line basis, not block basis, but use block processor
        # instead of preprocessor, because we still want to operate on the
        # xml tree.
        for line in "\n\n".join(blocks).split("\n"):
            # Superfences replaces code blocks with a placeholder,
            # Check for the placeholder, and ensure it is a marimo code block,
            # otherwise consider it as markdown.
            if not HTML_PLACEHOLDER_RE.match(line.strip()):
                # Use <!----> to indicate a separation between cells.
                if line.strip() == "<!---->":
                    add_paragraph()
                    continue
                text.append(line)
                continue

            lookup = line.strip()[1:-1]
            code = self.stash[lookup][0]
            if not _is_code_tag(code):
                text.extend(code.split("\n"))
                continue

            # Definitively a code block, so add the previous markdown.
            add_paragraph()

            code_block = SubElement(parent, MARIMO_CODE)
            block_lines = code.split("\n")
            code_block.text = "\n".join(block_lines[1:-1])

            attribs = extract_attribs(block_lines[0])
            if attribs:
                code_block.attrib = attribs

            # Set after to prevent lang being flushed.
            code_block.set("language", _get_language(code))

        add_paragraph()
        # Flush to indicate all blocks have been processed.
        blocks.clear()


def convert_from_md_to_marimo_ir(
    text: str, filepath: str | None = None
) -> NotebookSerializationV1:
    if not text.strip():
        return NotebookSerializationV1(
            app=AppInstantiation(options={}), filename=filepath
        )
    notebook = MarimoMdParser(
        output_format="marimo-ir",
        import_dialects=_markdown_import_dialects(text, filepath),
    ).convert(text)
    assert isinstance(notebook, NotebookSerializationV1)
    return NotebookSerializationV1(
        app=notebook.app,
        filename=filepath,
        cells=notebook.cells,
        violations=notebook.violations,
        valid=notebook.valid,
        header=notebook.header,
        version=notebook.version,
    )


def extract_frontmatter(text: str) -> tuple[dict[str, str], str]:
    from marimo._utils import yaml

    result = yaml.YAML_FRONT_MATTER_REGEX.match(text)

    if result:
        yaml_content = result.group(1)
        body = text[result.end() :].lstrip("\n")
        try:
            return yaml.load(yaml_content), body
        # If there's an error in parsing YAML, ignore the meta and proceed.
        except yaml.YAMLError:
            LOGGER.warning(
                "Error parsing frontmatter YAML. Ignoring frontmatter."
            )
            return {}, body
    return {}, text


def sanitize_markdown(text: str) -> str:
    return SanitizeParser(output_format="identity").convert(text)


def is_sanitized_markdown(text: str) -> bool:
    # "Unsanitized" markdown contains potentially unintended executatable code
    # block, which require backticks.
    return "```" not in text or sanitize_markdown(text) == text
