"""HTML sanitization with nh3 (SPEC §6.4, §10.4, invariant I-8).

Mail HTML is hostile: sanitized server-side at index time with a broad-but-safe
allowlist (structural + presentational tags, filtered inline CSS), and rendered
inside a sandboxed iframe. ``cid:`` images are rewritten to the per-message
endpoint; remote images (and CSS ``url()``) are blocked by default and flagged so
the UI can offer "load images" (SPEC §9.4).
"""

from __future__ import annotations

import re
import uuid
from dataclasses import dataclass

import nh3

# Bump when the sanitization policy changes; body_parts.sanitizer_v records it so a
# render re-sanitizes stale content from the raw blob (SPEC §10.4).
SANITIZER_VERSION = 3

_ALLOWED_TAGS: set[str] = {
    # structure
    "a",
    "abbr",
    "address",
    "article",
    "aside",
    "b",
    "bdi",
    "bdo",
    "blockquote",
    "br",
    "caption",
    "center",
    "cite",
    "code",
    "col",
    "colgroup",
    "dd",
    "del",
    "details",
    "dfn",
    "div",
    "dl",
    "dt",
    "em",
    "figcaption",
    "figure",
    "font",
    "footer",
    "h1",
    "h2",
    "h3",
    "h4",
    "h5",
    "h6",
    "header",
    "hr",
    "i",
    "img",
    "ins",
    "kbd",
    "li",
    "main",
    "mark",
    "nav",
    "ol",
    "p",
    "pre",
    "q",
    "s",
    "samp",
    "section",
    "small",
    "span",
    "strike",
    "strong",
    "style",
    "sub",
    "summary",
    "sup",
    "table",
    "tbody",
    "td",
    "tfoot",
    "th",
    "thead",
    "time",
    "tr",
    "tt",
    "u",
    "ul",
    "var",
    "wbr",
}

# Attributes allowed on any tag (presentational email formatting).
_GENERIC_ATTRS: set[str] = {
    "style",
    "class",
    "align",
    "valign",
    "dir",
    "title",
    "width",
    "height",
    "bgcolor",
    "color",
    "face",
    "size",
    "span",
    "nowrap",
    "lang",
}

_TAG_SPECIFIC: dict[str, set[str]] = {
    "a": {"href", "target"},
    "img": {"src", "alt"},
    "td": {"colspan", "rowspan"},
    "th": {"colspan", "rowspan", "scope"},
    "table": {"cellpadding", "cellspacing", "border"},
    "font": {"face", "color", "size"},
    "style": {"type", "media"},
}

# nh3 has no "generic attributes" option, so expand the generic set onto every tag.
_ALLOWED_ATTRS: dict[str, set[str]] = {
    tag: set(_GENERIC_ATTRS) | _TAG_SPECIFIC.get(tag, set()) for tag in _ALLOWED_TAGS
}

_URL_SCHEMES: set[str] = {"http", "https", "mailto", "cid", "data"}

_STYLE_DANGEROUS = re.compile(
    r"(expression\s*\(|javascript:|vbscript:|@import|-moz-binding|behavior\s*:)", re.IGNORECASE
)
_STYLE_URL = re.compile(r"url\s*\([^)]*\)", re.IGNORECASE)
# Script-like CSS that voids a whole <style> block (unlike @import, which we can strip).
_STYLE_BLOCK_DANGEROUS = re.compile(
    r"(expression\s*\(|javascript:|vbscript:|-moz-binding|behavior\s*:)", re.IGNORECASE
)
_STYLE_IMPORT = re.compile(r"@import[^;]*;?", re.IGNORECASE)
_STYLE_ELEMENT = re.compile(r"(<style\b[^>]*>)(.*?)(</style>)", re.IGNORECASE | re.DOTALL)


@dataclass(frozen=True)
class SanitizeResult:
    html: str
    has_remote_images: bool
    sanitizer_v: int


def _sanitize_style(value: str, allow_remote: bool) -> str | None:
    if _STYLE_DANGEROUS.search(value):
        return None  # drop the whole style rather than try to repair it
    if not allow_remote and _STYLE_URL.search(value):
        # Strip CSS url() so backgrounds can't smuggle in remote/tracking images.
        return _STYLE_URL.sub("none", value)
    return value


def _sanitize_style_elements(html: str, allow_remote: bool, state: dict[str, bool]) -> str:
    """Sanitize CSS inside ``<style>`` blocks (nh3 keeps the tag but not its text).

    Many mailers (e.g. Fastmail's calendar invites) put all layout in a ``<style>``
    block of class rules — dropping it loses tables/buttons entirely. We keep the
    block but neutralize script-like CSS and, unless remote images are approved,
    strip ``@import``/``url()`` so it can't fetch remote/tracking resources.
    """

    def repl(match: re.Match[str]) -> str:
        css = match.group(2)
        if _STYLE_BLOCK_DANGEROUS.search(css):
            return ""  # nuke the block entirely
        css = _STYLE_IMPORT.sub("", css)
        if not allow_remote and _STYLE_URL.search(css):
            state["has_remote"] = True
            css = _STYLE_URL.sub("none", css)
        return f"{match.group(1)}{css}{match.group(3)}"

    return _STYLE_ELEMENT.sub(repl, html)


def sanitize_html(
    html: str, message_id: uuid.UUID, *, allow_remote_images: bool = False
) -> SanitizeResult:
    """Sanitize ``html`` for the message ``message_id``.

    Returns the safe HTML, whether any remote images were blocked, and the
    sanitizer version. Pass ``allow_remote_images=True`` after the user approves
    a sender to render remote images and CSS backgrounds (SPEC §9.4).
    """
    state = {"has_remote": False}

    def attribute_filter(tag: str, attr: str, value: str) -> str | None:
        if attr == "style":
            return _sanitize_style(value, allow_remote_images)
        if tag == "a" and attr == "target":
            return "_blank"
        if tag == "img" and attr == "src":
            lowered = value.strip().lower()
            if lowered.startswith("cid:"):
                content_id = value.strip()[4:]
                return f"/api/messages/{message_id}/cid/{content_id}"
            if lowered.startswith(("http:", "https:")):
                state["has_remote"] = True
                return value if allow_remote_images else None  # block by default (FR-23)
            if lowered.startswith("data:image/"):
                return value
            return None
        return value

    cleaned = nh3.clean(
        html,
        tags=_ALLOWED_TAGS,
        attributes=_ALLOWED_ATTRS,
        url_schemes=_URL_SCHEMES,
        link_rel="noopener noreferrer nofollow",
        strip_comments=True,
        # nh3 strips <script>/<style> *content* by default; we allow <style> (above)
        # and sanitize its CSS ourselves, so only script content is stripped here.
        clean_content_tags={"script"},
        attribute_filter=attribute_filter,
    )
    cleaned = _sanitize_style_elements(cleaned, allow_remote_images, state)
    return SanitizeResult(
        html=cleaned,
        has_remote_images=state["has_remote"],
        sanitizer_v=SANITIZER_VERSION,
    )