"""Search result highlighting with fragment scoring.

Highlights matching terms in search result snippets.  Splits text into
fragments of approximately ``fragment_size`` characters, scores each
fragment by query term density and position, and returns the top
fragments with matching terms wrapped in configurable tags.

Fragment scoring algorithm:
    1. Split text at sentence/word boundaries into ~fragment_size chunks.
    2. For each query term match in a fragment, add 1/(position+1) where
       position is the ordinal index of the match within the fragment.
    3. Multiply by a boost equal to the number of unique query terms
       matched in the fragment.
    4. Return the top ``num_fragments`` fragments joined by ``separator``.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import List, Set


# Pre-compiled patterns
_WORD_RE = re.compile(r"\b\w+\b", re.UNICODE)
_SENTENCE_END_RE = re.compile(r"[.!?]\s+", re.UNICODE)


@dataclass
class HighlightConfig:
    """Configuration for the highlighter.

    Attributes:
        pre_tag:       Opening tag wrapped around matching terms.
        post_tag:      Closing tag wrapped around matching terms.
        fragment_size: Approximate character length per fragment.
        num_fragments: Maximum number of fragments to return.
        separator:     String used to join selected fragments.
    """

    pre_tag: str = "<b>"
    post_tag: str = "</b>"
    fragment_size: int = 150
    num_fragments: int = 3
    separator: str = " ... "


@dataclass
class Fragment:
    """A scored text fragment.

    Attributes:
        text:  The raw (un-highlighted) fragment text.
        score: Relevance score based on query term density.
        start: Character offset of the fragment start in the source text.
        end:   Character offset of the fragment end in the source text.
    """

    text: str
    score: float
    start: int
    end: int


class Highlighter:
    """Highlights matching query terms in text fragments.

    Args:
        config: Optional :class:`HighlightConfig`.  Uses defaults when *None*.
    """

    def __init__(self, config: HighlightConfig | None = None) -> None:
        self._config = config or HighlightConfig()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def highlight(self, text: str, query_terms: Set[str]) -> str:
        """Return highlighted text with best fragments joined by separator.

        Splits *text* into scored fragments, selects the top
        ``num_fragments``, and returns them in document order with
        matching terms wrapped in the configured tags.

        Args:
            text:        Source text to highlight.
            query_terms: Set of lower-cased query terms to match.

        Returns:
            Highlighted string with fragments joined by *separator*.
        """
        if not text or not query_terms:
            return text

        fragments = self.get_fragments(text, query_terms)
        if not fragments:
            return text

        cfg = self._config
        # Take top num_fragments by score, then sort by document position.
        top = sorted(fragments, key=lambda f: f.score, reverse=True)[
            : cfg.num_fragments
        ]
        top.sort(key=lambda f: f.start)

        parts: List[str] = []
        lower_terms = _lower_set(query_terms)
        for frag in top:
            parts.append(self._tag_matches(frag.text, lower_terms))

        return cfg.separator.join(parts)

    def get_fragments(
        self, text: str, query_terms: Set[str]
    ) -> List[Fragment]:
        """Get scored fragments from *text*.

        Args:
            text:        Source text.
            query_terms: Set of lower-cased query terms.

        Returns:
            List of :class:`Fragment` objects (unsorted).
        """
        if not text or not query_terms:
            return []

        lower_terms = _lower_set(query_terms)
        raw_fragments = self._split_fragments(text)
        scored: List[Fragment] = []

        for start, end in raw_fragments:
            chunk = text[start:end]
            score = self._score_fragment(chunk, lower_terms)
            scored.append(Fragment(text=chunk, score=score, start=start, end=end))

        return scored

    def highlight_field(self, text: str, query_terms: Set[str]) -> str:
        """Highlight the entire text without fragmentation.

        Every occurrence of a query term is wrapped in the configured
        tags.  Useful for short fields (titles, tags) where fragmentation
        is unnecessary.

        Args:
            text:        Source text.
            query_terms: Set of lower-cased query terms.

        Returns:
            Full text with matching terms highlighted.
        """
        if not text or not query_terms:
            return text
        return self._tag_matches(text, _lower_set(query_terms))

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _split_fragments(self, text: str) -> List[tuple]:
        """Split *text* into (start, end) fragment boundaries.

        Tries to break at sentence boundaries first, then falls back to
        word boundaries, producing chunks of approximately
        ``fragment_size`` characters.
        """
        size = self._config.fragment_size
        fragments: List[tuple] = []
        text_len = len(text)

        if text_len <= size:
            return [(0, text_len)]

        pos = 0
        while pos < text_len:
            end = min(pos + size, text_len)

            if end < text_len:
                # Try sentence boundary first
                best = self._find_break(text, pos + size // 2, end + size // 4)
                if best != -1:
                    end = best
                else:
                    # Fall back to word boundary
                    space = text.rfind(" ", pos, end + 20)
                    if space > pos:
                        end = space + 1

            fragments.append((pos, end))
            pos = end

        return fragments

    @staticmethod
    def _find_break(text: str, start: int, end: int) -> int:
        """Find a sentence-ending break between *start* and *end*.

        Returns the character position right after the sentence-ending
        whitespace, or -1 if none found.
        """
        region = text[start:end]
        match = _SENTENCE_END_RE.search(region)
        if match:
            return start + match.end()
        return -1

    @staticmethod
    def _score_fragment(chunk: str, lower_terms: Set[str]) -> float:
        """Score a fragment by query term density and coverage.

        For each word in the fragment that matches a query term, add
        ``1 / (word_position + 1)`` where *word_position* is the
        zero-based index of the word in the fragment — so earlier
        matches score higher.  The sum is then multiplied by the number
        of unique matched terms to boost fragments that cover more of
        the query.
        """
        words = _WORD_RE.findall(chunk)
        if not words:
            return 0.0

        unique_matched: Set[str] = set()
        score = 0.0

        for position, word in enumerate(words):
            low = word.lower()
            if low in lower_terms:
                score += 1.0 / (position + 1)
                unique_matched.add(low)

        if not unique_matched:
            return 0.0

        # Boost by number of unique query terms covered
        return score * len(unique_matched)

    def _tag_matches(self, text: str, lower_terms: Set[str]) -> str:
        """Wrap every word matching *lower_terms* in pre/post tags."""
        cfg = self._config

        def _replace(m: re.Match) -> str:
            word = m.group(0)
            if word.lower() in lower_terms:
                return f"{cfg.pre_tag}{word}{cfg.post_tag}"
            return word

        return _WORD_RE.sub(_replace, text)


# ------------------------------------------------------------------
# Module-level helpers
# ------------------------------------------------------------------

def _lower_set(terms: Set[str]) -> Set[str]:
    """Return a new set with all terms lower-cased."""
    return {t.lower() for t in terms}
