import html
import re
from typing import List, Set

import nh3

from ._Sanitizer import Sanitizer, T

# Matches a complete HTML entity (named, decimal, or hex) OR a bare `&`.
# The complete-entity alternative is tried first, so `&gt;` is captured whole
# and only standalone `&` (or incomplete fragments like `&ltda`) fall through
# to the final `|&`. Protecting bare `&` too stops nh3 from normalising
# incomplete entities (e.g. `&ltda` -> `&lt;da`) during clean().
_ENTITY_PATTERN = re.compile(r"&(?:[a-zA-Z][a-zA-Z0-9]*|#\d+|#[xX][0-9a-fA-F]+);|&")

# Placeholder format written during entity protection: two PUA sentinels
# (U+E000, U+E001) around the entity's index. PUA chars are never produced
# by nh3 or html.unescape, so they round-trip safely.
_PLACEHOLDER_PATTERN = re.compile("\ue000(\\d+)\ue001")

# Hoisted to module scope so each _sanitize call doesn't allocate a new set
# just to pass `tags=` into nh3.clean.
_EMPTY_TAGS: Set[str] = set()


class Nh3Sanitizer(Sanitizer):
    def _types_to_sanitize(self) -> List[type]:
        return [str]

    def _sanitize(self, value: T) -> T:
        if not isinstance(value, str):
            return value

        # Tiered fast paths. nh3.clean is the dominant cost (Rust FFI + full
        # HTML parse), so skipping it whenever it would be a no-op matters:
        #   - no `<`, `>`, `&` => nothing to strip or escape; return as-is.
        #   - `<`/`>` but no `&` => no entities to preserve, so the placeholder
        #     round-trip is unnecessary. nh3 escapes bare brackets, then
        #     html.unescape reverses that — result matches the full pipeline.
        #   - `&` present => full pipeline (entity protection is required so
        #     existing entities survive nh3's normalization).
        if "&" not in value:
            if "<" not in value and ">" not in value:
                return value
            return html.unescape(nh3.clean(value, tags=_EMPTY_TAGS))

        # Full pipeline: protect every `&` (complete entity or bare) with a
        # placeholder, run nh3 to strip tags and escape bare `<`/`>`,
        # html.unescape to turn those nh3-introduced entities back into raw
        # chars, then restore the original `&...` sequences from placeholders.
        entities: List[str] = []

        def _store(match: re.Match) -> str:
            entities.append(match.group(0))
            return f"\ue000{len(entities) - 1}\ue001"

        protected = _ENTITY_PATTERN.sub(_store, value)
        cleaned = nh3.clean(protected, tags=_EMPTY_TAGS)
        unescaped = html.unescape(cleaned)

        return _PLACEHOLDER_PATTERN.sub(
            lambda m: entities[int(m.group(1))],
            unescaped,
        )
