"""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.
# v6: links keep only http(s) and mailto (no data: pages).
# v7: tracking pixels are removed (and counted), even when images are allowed.
# v8: CSS sizes in other units (in, em, %, pt) no longer read as pixels — v7
#     took 0.31in signature logos for trackers.
SANITIZER_VERSION = 8

# 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
    trackers_blocked: int = 0


# --- Tracking pixels ---------------------------------------------------------------
# A remote image that is invisible (1×1, or hidden by CSS) has no use but to tell
# the sender the mail was opened, and by whom. A few well-known open-tracking URL
# shapes are caught whatever their size. Matching runs on the raw HTML before nh3:
# a miss only leaves an image in, and nh3 still cleans whatever remains.
_IMG_TAG = re.compile(r"<img\b[^>]*>", re.IGNORECASE)
_ATTR = re.compile(r"""([a-zA-Z_:][-a-zA-Z0-9_:.]*)\s*=\s*("[^"]*"|'[^']*'|[^\s"'>]+)""")
_CSS_HIDDEN = re.compile(
    r"display\s*:\s*none|visibility\s*:\s*hidden|opacity\s*:\s*0(?![.\d]*[1-9])",
    re.IGNORECASE,
)
_TRACKER_URL = re.compile(
    r"//[^/]*mailtrack\.io/"
    r"|/track/open\b"
    r"|/wf/open\b"
    r"|/open\.(?:php|aspx|gif)\b"
    r"|[?&]tracking[_-]?pixel=",
    re.IGNORECASE,
)


_PX = re.compile(r"\s*(\d+(?:\.\d*)?|\.\d+)\s*(px)?\s*$", re.IGNORECASE)


def _px(value: str | None) -> float | None:
    """A size in pixels ("1", "1px"), or None for anything else: other units
    (0.31in is a 30-pixel signature logo), auto, %, or malformed (".", "1.2.3")."""
    if value is None:
        return None
    match = _PX.match(value)
    return float(match.group(1)) if match else None


def _tiny(value: str | None) -> bool:
    px = _px(value)
    return px is not None and px <= 2


def _css(style: str) -> dict[str, str]:
    """An inline style's declarations, lower-cased, without !important."""
    out: dict[str, str] = {}
    for decl in style.split(";"):
        name, _, value = decl.partition(":")
        value = value.lower().replace("!important", "").strip()
        if name.strip() and value:
            out[name.strip().lower()] = value
    return out


def is_tracking_pixel(tag: str) -> bool:
    """Whether an ``<img …>`` tag is a remote open-tracking pixel."""
    attrs = {name.lower(): value.strip("\"'") for name, value in _ATTR.findall(tag)}
    src = attrs.get("src", "").strip()
    if not src.lower().startswith(("http:", "https:", "//")):
        return False  # inline (cid:/data:) images can't report anything
    if _TRACKER_URL.search(src):
        return True
    style = attrs.get("style", "")
    if _CSS_HIDDEN.search(style):
        return True
    # The size the browser draws: a style width/height overrides the attribute
    # (and "auto", % or another unit is no tiny size); a max-width/max-height
    # of a pixel or two caps it; height "auto" follows the width.
    css = _css(style)
    wide = _tiny(css["width"]) if "width" in css else _tiny(attrs.get("width"))
    tall = _tiny(css["height"]) if "height" in css else _tiny(attrs.get("height"))
    wide = wide or _tiny(css.get("max-width"))
    tall = tall or _tiny(css.get("max-height"))
    if css.get("height") == "auto":
        tall = wide
    return wide and tall


def strip_tracking_pixels(html: str) -> tuple[str, int]:
    """Remove tracking pixels from raw HTML; returns the HTML and how many."""
    count = 0

    def repl(match: re.Match[str]) -> str:
        nonlocal count
        if is_tracking_pixel(match.group(0)):
            count += 1
            return ""
        return match.group(0)

    return _IMG_TAG.sub(repl, html), count


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,
    block_trackers: bool = True,
) -> 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).
    ``block_trackers=False`` (Settings → Privacy) leaves tracking pixels in; they
    are remote images like any other then.
    """
    state = {"has_remote": False}
    trackers = 0
    if block_trackers:
        html, trackers = strip_tracking_pixels(html)

    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 attr == "href":
            # A link goes to the web or to a mail address: data: (a page made
            # up by the sender) and cid: are kept for inline images only.
            scheme = value.strip().lower().split(":", 1)[0] if ":" in value else ""
            return value if scheme in ("", "http", "https", "mailto") else None
        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,
        trackers_blocked=trackers,
    )