"""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).
# v5: this server's own inline-image URLs quoted in a message load from here.
SANITIZER_VERSION = 5

# This server's inline-image URL, however it ended up written: "/api/…",
# "https://host/api/…", or mangled to "http:///api/…" by an earlier copy.
INTERNAL_CID = re.compile(r"/api/messages/([0-9a-fA-F-]{36})/cid/([^\"'?#\s<>]+)")

_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}"
            # This server's own inline-image URL, quoted from another message
            # (by an earlier version, or another of the user's own mails) —
            # possibly mangled to "http:///api/…", which a browser takes for a
            # host named "api": load it from here, never look that host up.
            internal = INTERNAL_CID.search(value)
            if internal is not None and not lowered.startswith("data:"):
                return f"/api/messages/{internal.group(1)}/cid/{internal.group(2)}"
            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 its content stays. <title> is dropped
        # with its text: as a disallowed tag nh3 would unwrap it, and the document
        # title ("mailer", the subject) would render as a line above the body.
        clean_content_tags={"script", "title"},
        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,
    )