"""Server-side HTML sanitization for agent output.

Provides a whitelist-based sanitizer that strips dangerous HTML tags and
attributes from markdown-rendered content before sending to clients.
This mitigates XSS risks when agent output is rendered in the browser.
"""

import re
from html import escape as html_escape

# Tags allowed in markdown output (subset safe for rendering)
ALLOWED_TAGS: set[str] = {
    # Block elements
    "p", "div", "br", "hr",
    "h1", "h2", "h3", "h4", "h5", "h6",
    "blockquote", "pre", "code",
    "ul", "ol", "li",
    "table", "thead", "tbody", "tr", "th", "td",
    "dl", "dt", "dd",
    "figure", "figcaption",
    # Inline elements
    "a", "strong", "b", "em", "i", "u", "s", "del", "ins",
    "sub", "sup", "mark", "small",
    "span", "abbr", "cite", "q",
    "kbd", "samp", "var",
    # Media (constrained by CSP)
    "img",
    # Details/summary
    "details", "summary",
}

# Attributes allowed per tag (whitelist approach)
ALLOWED_ATTRIBUTES: dict[str, set[str]] = {
    "a": {"href", "title", "target", "rel"},
    "img": {"src", "alt", "title", "width", "height"},
    "td": {"colspan", "rowspan", "align"},
    "th": {"colspan", "rowspan", "align", "scope"},
    "code": {"class"},  # language-* classes from syntax highlighting
    "pre": {"class"},
    "span": {"class"},
    "div": {"class"},
    "ol": {"start", "type"},
    "blockquote": {"cite"},
    "abbr": {"title"},
    "details": {"open"},
}

# Attributes that should never appear (event handlers, etc.)
DANGEROUS_ATTR_PATTERN = re.compile(
    r"\b(on\w+|style|formaction|xlink:href|data-bind)\s*=",
    re.IGNORECASE,
)

# Pattern for matching HTML tags
_TAG_PATTERN = re.compile(
    r"<(/?)(\w+)((?:\s+[^>]*?)?)(/?)>",
    re.DOTALL,
)

# Pattern for matching attributes within a tag
_ATTR_PATTERN = re.compile(
    r"""(\w[\w-]*)\s*=\s*(?:"([^"]*)"|'([^']*)'|(\S+))""",
)

# Dangerous protocols in href/src attributes
_DANGEROUS_PROTOCOLS = re.compile(
    r"^\s*(javascript|vbscript|data\s*:(?!image/))",
    re.IGNORECASE,
)


def sanitize_html(html: str) -> str:
    """Sanitize HTML content using a whitelist approach.

    Strips tags not in ALLOWED_TAGS, removes dangerous attributes,
    and neutralizes javascript: URIs.

    Args:
        html: Raw HTML string (typically from markdown rendering).

    Returns:
        Sanitized HTML string safe for browser rendering.
    """
    if not html:
        return html

    def _replace_tag(match: re.Match) -> str:
        closing_slash = match.group(1)  # "/" for closing tags
        tag_name = match.group(2).lower()
        attrs_str = match.group(3)
        self_closing = match.group(4)  # "/" for self-closing tags

        # Strip tags not in the whitelist
        if tag_name not in ALLOWED_TAGS:
            return ""

        # For closing tags, no attributes needed
        if closing_slash:
            return f"</{tag_name}>"

        # Filter attributes
        allowed_attrs = ALLOWED_ATTRIBUTES.get(tag_name, set())
        clean_attrs = []

        for attr_match in _ATTR_PATTERN.finditer(attrs_str):
            attr_name = attr_match.group(1).lower()
            attr_value = attr_match.group(2) or attr_match.group(3) or attr_match.group(4) or ""

            # Skip attributes not in the whitelist for this tag
            if attr_name not in allowed_attrs:
                continue

            # Block event handlers that might have slipped through
            if attr_name.startswith("on"):
                continue

            # Sanitize href and src attributes
            if attr_name in ("href", "src"):
                if _DANGEROUS_PROTOCOLS.match(attr_value):
                    continue

            # Escape the attribute value
            safe_value = html_escape(attr_value, quote=True)
            clean_attrs.append(f'{attr_name}="{safe_value}"')

        # Add rel="noopener noreferrer" to links with target
        if tag_name == "a":
            attr_dict = {a.split("=")[0]: a for a in clean_attrs}
            if "target" in attr_dict and "rel" not in attr_dict:
                clean_attrs.append('rel="noopener noreferrer"')

        attrs_part = " " + " ".join(clean_attrs) if clean_attrs else ""
        self_close = " /" if self_closing else ""

        return f"<{tag_name}{attrs_part}{self_close}>"

    # Remove any dangerous attribute patterns before tag processing
    html = DANGEROUS_ATTR_PATTERN.sub("", html)

    # Process tags
    sanitized = _TAG_PATTERN.sub(_replace_tag, html)

    return sanitized


def sanitize_agent_output(text: str) -> str:
    """Sanitize agent text output that may contain HTML.

    This is the main entry point for sanitizing agent responses before
    they are sent to WebSocket clients. It handles both plain text and
    HTML content from markdown rendering.

    Args:
        text: Agent output text, possibly containing HTML.

    Returns:
        Sanitized text safe for browser rendering.
    """
    if not text:
        return text

    # If the text looks like it contains HTML, sanitize it
    if "<" in text and ">" in text:
        return sanitize_html(text)

    return text
