"""Safe markdown-to-HTML rendering helpers for assistant-facing UIs."""

from __future__ import annotations

import html
import re
from html.parser import HTMLParser
from urllib.parse import urlsplit

try:
    import markdown as markdown_lib
except ModuleNotFoundError:  # pragma: no cover - optional dependency at runtime
    markdown_lib = None

_ALLOWED_TAGS = {
    "a",
    "blockquote",
    "br",
    "code",
    "em",
    "h1",
    "h2",
    "h3",
    "h4",
    "h5",
    "h6",
    "hr",
    "i",
    "li",
    "ol",
    "p",
    "pre",
    "strong",
    "table",
    "tbody",
    "td",
    "th",
    "thead",
    "tr",
    "ul",
}
_VOID_TAGS = {"br", "hr"}
_DROP_WITH_CONTENT_TAGS = {"script", "style", "iframe", "object", "embed", "svg", "math"}
_ALLOWED_ATTRS = {
    "a": {"href", "title"},
    "code": {"class"},
}
_SAFE_LINK_SCHEMES = {"http", "https", "mailto"}
_CLASS_TOKEN_RE = re.compile(r"^[A-Za-z0-9_-]{1,48}$")
_LIST_MARKER_RE = re.compile(r"^\s*(?:[-*+]\s+|\d+\.\s+)")


def _prepare_markdown(text: str) -> str:
    """Normalize chat markdown so common list patterns render predictably."""
    lines = text.splitlines()
    output: list[str] = []
    in_fence = False
    for line in lines:
        stripped = line.strip()
        if stripped.startswith("```") or stripped.startswith("~~~"):
            in_fence = not in_fence
            output.append(line)
            continue
        is_list_line = (not in_fence) and bool(_LIST_MARKER_RE.match(line))
        if is_list_line and output:
            previous = output[-1]
            previous_is_list = bool(_LIST_MARKER_RE.match(previous))
            if previous.strip() and not previous_is_list:
                output.append("")
        output.append(line)
    return "\n".join(output)


def _sanitize_href(value: str) -> str | None:
    href = value.strip()
    if not href:
        return None
    lower = href.lower()
    compact = re.sub(r"[\x00-\x20]+", "", lower)
    if compact.startswith(("javascript:", "vbscript:", "data:")):
        return None
    if href.startswith(("/", "#")):
        return href
    parsed = urlsplit(href)
    if parsed.scheme:
        if parsed.scheme.lower() not in _SAFE_LINK_SCHEMES:
            return None
        return href
    return href


def _sanitize_class_value(value: str) -> str | None:
    tokens = [token for token in value.split() if _CLASS_TOKEN_RE.fullmatch(token)]
    if not tokens:
        return None
    return " ".join(tokens[:8])


class _SafeHtmlSanitizer(HTMLParser):
    """Minimal allowlist sanitizer for markdown-generated HTML."""

    def __init__(self) -> None:
        super().__init__(convert_charrefs=False)
        self._parts: list[str] = []
        self._drop_depth = 0

    def _serialize_attrs(self, tag: str, attrs: list[tuple[str, str | None]]) -> str:
        allowed = _ALLOWED_ATTRS.get(tag, set())
        output_attrs: list[tuple[str, str]] = []
        has_safe_href = False
        for name, value in attrs:
            key = str(name or "").strip().lower()
            if not key or key.startswith("on") or key not in allowed:
                continue
            if value is None:
                continue
            raw_value = str(value)
            if tag == "a" and key == "href":
                safe_href = _sanitize_href(raw_value)
                if not safe_href:
                    continue
                output_attrs.append(("href", safe_href))
                has_safe_href = True
                continue
            if tag == "code" and key == "class":
                safe_class = _sanitize_class_value(raw_value)
                if not safe_class:
                    continue
                output_attrs.append(("class", safe_class))
                continue
            output_attrs.append((key, raw_value.strip()))
        if tag == "a" and has_safe_href:
            output_attrs.append(("rel", "noopener noreferrer nofollow"))
        if not output_attrs:
            return ""
        return "".join(f' {name}="{html.escape(value, quote=True)}"' for name, value in output_attrs)

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        tag_name = str(tag or "").strip().lower()
        if not tag_name:
            return
        if self._drop_depth > 0:
            if tag_name in _DROP_WITH_CONTENT_TAGS:
                self._drop_depth += 1
            return
        if tag_name in _DROP_WITH_CONTENT_TAGS:
            self._drop_depth = 1
            return
        if tag_name not in _ALLOWED_TAGS:
            return
        attr_html = self._serialize_attrs(tag_name, attrs)
        self._parts.append(f"<{tag_name}{attr_html}>")

    def handle_startendtag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        tag_name = str(tag or "").strip().lower()
        if not tag_name:
            return
        if self._drop_depth > 0:
            return
        if tag_name in _DROP_WITH_CONTENT_TAGS:
            return
        if tag_name not in _ALLOWED_TAGS:
            return
        attr_html = self._serialize_attrs(tag_name, attrs)
        self._parts.append(f"<{tag_name}{attr_html}>")
        if tag_name not in _VOID_TAGS:
            self._parts.append(f"</{tag_name}>")

    def handle_endtag(self, tag: str) -> None:
        tag_name = str(tag or "").strip().lower()
        if not tag_name:
            return
        if self._drop_depth > 0:
            if tag_name in _DROP_WITH_CONTENT_TAGS:
                self._drop_depth = max(0, self._drop_depth - 1)
            return
        if tag_name in _ALLOWED_TAGS and tag_name not in _VOID_TAGS:
            self._parts.append(f"</{tag_name}>")

    def handle_data(self, data: str) -> None:
        if self._drop_depth > 0:
            return
        self._parts.append(html.escape(data, quote=False))

    def handle_entityref(self, name: str) -> None:
        if self._drop_depth > 0:
            return
        self._parts.append(f"&{name};")

    def handle_charref(self, name: str) -> None:
        if self._drop_depth > 0:
            return
        self._parts.append(f"&#{name};")

    def handle_comment(self, data: str) -> None:
        return

    def get_html(self) -> str:
        return "".join(self._parts)


def sanitize_html(unsafe_html: str) -> str:
    """Sanitize untrusted HTML with an allowlist policy."""
    parser = _SafeHtmlSanitizer()
    parser.feed(str(unsafe_html or ""))
    parser.close()
    return parser.get_html()


def render_markdown(markdown_text: str) -> str:
    """Render markdown to safe HTML for UI insertion."""
    text = str(markdown_text or "")
    if not text.strip():
        return ""
    prepared = _prepare_markdown(text)
    if markdown_lib is None:
        return f"<pre><code>{html.escape(prepared)}</code></pre>"
    unsafe_html = markdown_lib.markdown(
        prepared,
        extensions=["extra", "sane_lists", "tables", "nl2br"],
        output_format="html5",
    )
    return sanitize_html(unsafe_html)
