"""
RBSP Query Rewriter.

Query rewriting and normalization pipeline that transforms user queries into
optimized forms before execution.  Each rewrite rule is a composable unit
that handles one transformation (lowercasing, stop-word removal, stemming,
etc.).  The ``QueryRewriter`` chains rules in order, producing a
``RewrittenQuery`` that downstream components can consume directly.

All parsing uses ``re`` from stdlib.  Zero external dependencies.
"""

from __future__ import annotations

import re
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Dict, List, Optional


# ---------------------------------------------------------------------------
# Stop words — a compact, high-recall English stop list
# ---------------------------------------------------------------------------

_STOP_WORDS: frozenset[str] = frozenset({
    "a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for",
    "of", "with", "by", "from", "is", "it", "as", "was", "are", "be",
    "been", "being", "have", "has", "had", "do", "does", "did", "will",
    "would", "could", "should", "may", "might", "shall", "can", "this",
    "that", "these", "those", "i", "you", "he", "she", "we", "they", "me",
    "him", "her", "us", "them", "my", "your", "his", "its", "our", "their",
    "what", "which", "who", "whom", "where", "when", "how", "not", "no",
    "if", "then", "than", "so", "up", "out", "about", "into", "over",
    "after", "before", "between", "under", "again", "further", "once",
})


# ---------------------------------------------------------------------------
# RewrittenQuery — the output of the rewriting pipeline
# ---------------------------------------------------------------------------

@dataclass(slots=True)
class RewrittenQuery:
    """Result of applying the rewriting pipeline to a raw query.

    Attributes:
        text:          The rewritten query text.
        phrases:       Extracted quoted-phrase queries (in appearance order).
        field_filters: Extracted ``field:value`` pairs.
        original:      The original, unmodified query string.
    """

    text: str
    phrases: List[str] = field(default_factory=list)
    field_filters: Dict[str, str] = field(default_factory=dict)
    original: str = ""


# ---------------------------------------------------------------------------
# RewriteRule base class
# ---------------------------------------------------------------------------

class RewriteRule(ABC):
    """Abstract base for a single query-rewrite transformation.

    Each rule receives a ``RewrittenQuery`` (which may already have been
    mutated by earlier rules) and returns a new or mutated
    ``RewrittenQuery``.
    """

    @abstractmethod
    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        """Apply this transformation to *query* and return the result."""


# ---------------------------------------------------------------------------
# Concrete rules
# ---------------------------------------------------------------------------

class LowercaseRule(RewriteRule):
    """Lowercase the query text."""

    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        query.text = query.text.lower()
        return query


class StripPunctuationRule(RewriteRule):
    """Remove non-essential punctuation.

    Preserves colons (used in field:value syntax), double-quotes (phrase
    delimiters), tildes and equals (query operators), and hyphens within
    words.
    """

    _STRIP_RE = re.compile(r'[^\w\s:~="\'\\-]', re.UNICODE)

    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        query.text = self._STRIP_RE.sub("", query.text)
        return query


class CollapseWhitespaceRule(RewriteRule):
    """Normalize runs of whitespace to a single space and strip edges."""

    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        query.text = re.sub(r"\s+", " ", query.text).strip()
        return query


class StopWordRemovalRule(RewriteRule):
    """Remove common English stop words from the query text.

    Words that appear inside quoted phrases or as part of ``field:value``
    tokens are left untouched — the rule only strips bare stop-word tokens.
    Boolean operators (AND / OR / NOT) are preserved.
    """

    _BOOL_OPS = frozenset({"AND", "OR", "NOT"})

    def __init__(self, stop_words: Optional[frozenset[str]] = None) -> None:
        self._stop_words = stop_words if stop_words is not None else _STOP_WORDS

    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        tokens: List[str] = []
        in_quote = False
        for token in query.text.split():
            # Crude quoted-region tracking — don't strip inside quotes
            if token.startswith('"') and not token.endswith('"'):
                in_quote = True
                tokens.append(token)
                continue
            if in_quote:
                tokens.append(token)
                if token.endswith('"'):
                    in_quote = False
                continue

            # Preserve boolean operators
            if token in self._BOOL_OPS:
                tokens.append(token)
                continue

            # Preserve field:value tokens
            if ":" in token and not token.startswith(":"):
                tokens.append(token)
                continue

            if token.lower() not in self._stop_words:
                tokens.append(token)

        query.text = " ".join(tokens)
        return query


class StemRule(RewriteRule):
    """Stem each query term using the Porter stemmer.

    Preserves quoted phrases, boolean operators, and ``field:value`` tokens
    in their original form.
    """

    _BOOL_OPS = frozenset({"AND", "OR", "NOT"})

    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        from rbsp.core.stemmer import stem

        tokens: List[str] = []
        in_quote = False
        for token in query.text.split():
            if token.startswith('"') and not token.endswith('"'):
                in_quote = True
                tokens.append(token)
                continue
            if in_quote:
                tokens.append(token)
                if token.endswith('"'):
                    in_quote = False
                continue

            if token in self._BOOL_OPS:
                tokens.append(token)
                continue

            if ":" in token and not token.startswith(":"):
                tokens.append(token)
                continue

            tokens.append(stem(token))

        query.text = " ".join(tokens)
        return query


class FieldExtractionRule(RewriteRule):
    """Extract ``field:value`` syntax from the query.

    Tokens matching ``field:value`` are removed from the query text and
    stored in ``RewrittenQuery.field_filters``.  If a field appears more
    than once the last value wins.

    Example::

        "type:python error message" -> field_filters={"type": "python"},
                                       text="error message"
    """

    _FIELD_RE = re.compile(r"\b([a-zA-Z_]\w*):(\S+)")

    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        filters: Dict[str, str] = dict(query.field_filters)
        text = query.text

        for match in self._FIELD_RE.finditer(text):
            field_name = match.group(1)
            field_value = match.group(2)
            filters[field_name] = field_value

        # Remove matched field:value tokens from the text
        text = self._FIELD_RE.sub("", text)
        text = re.sub(r"\s+", " ", text).strip()

        query.text = text
        query.field_filters = filters
        return query


class QuotedPhraseRule(RewriteRule):
    """Detect ``"quoted phrases"`` and store them for phrase query handling.

    Quoted content is extracted into ``RewrittenQuery.phrases`` and removed
    from the main text so downstream rules don't process it further.
    """

    _PHRASE_RE = re.compile(r'"([^"]*)"')

    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        phrases: List[str] = list(query.phrases)

        for match in self._PHRASE_RE.finditer(query.text):
            phrase = match.group(1).strip()
            if phrase:
                phrases.append(phrase)

        # Remove quoted sections from the text
        text = self._PHRASE_RE.sub("", query.text)
        text = re.sub(r"\s+", " ", text).strip()

        query.text = text
        query.phrases = phrases
        return query


class OperatorNormalizationRule(RewriteRule):
    """Normalize boolean operator variants to canonical uppercase form.

    Maps common variants:
        ``and``, ``&&`` -> ``AND``
        ``or``, ``||``  -> ``OR``
        ``not``, ``!``  -> ``NOT``
    """

    _OP_MAP = {
        "and": "AND", "&&": "AND",
        "or": "OR", "||": "OR",
        "not": "NOT", "!": "NOT",
    }

    def rewrite(self, query: RewrittenQuery) -> RewrittenQuery:
        tokens: List[str] = []
        for token in query.text.split():
            normalized = self._OP_MAP.get(token, None)
            if normalized is not None:
                tokens.append(normalized)
            else:
                tokens.append(token)
        query.text = " ".join(tokens)
        return query


# ---------------------------------------------------------------------------
# QueryRewriter — the pipeline orchestrator
# ---------------------------------------------------------------------------

_DEFAULT_RULES: tuple[type[RewriteRule], ...] = (
    QuotedPhraseRule,
    FieldExtractionRule,
    OperatorNormalizationRule,
    LowercaseRule,
    StripPunctuationRule,
    CollapseWhitespaceRule,
    StopWordRemovalRule,
    StemRule,
)


class QueryRewriter:
    """Composable query-rewrite pipeline.

    Applies an ordered sequence of ``RewriteRule`` instances to transform a
    raw query string into a ``RewrittenQuery``.

    Args:
        rules: Explicit rule list.  When *None*, a sensible default chain is
               used (phrase extraction -> field extraction -> operator
               normalization -> lowercase -> strip punctuation -> collapse
               whitespace -> stop-word removal -> stemming).

    Usage::

        rewriter = QueryRewriter()
        result = rewriter.rewrite('type:python "exact match" running tests')
    """

    def __init__(self, rules: Optional[List[RewriteRule]] = None) -> None:
        if rules is not None:
            self._rules: List[RewriteRule] = list(rules)
        else:
            self._rules = [cls() for cls in _DEFAULT_RULES]

    def rewrite(self, query: str) -> RewrittenQuery:
        """Run the full rewrite pipeline on *query*.

        Args:
            query: The raw user query string.

        Returns:
            A ``RewrittenQuery`` with the transformed text and any
            extracted metadata (phrases, field filters).
        """
        rq = RewrittenQuery(text=query, original=query)
        for rule in self._rules:
            rq = rule.rewrite(rq)
        return rq

    def add_rule(self, rule: RewriteRule) -> None:
        """Append a rewrite rule to the end of the pipeline.

        Args:
            rule: The rule instance to add.
        """
        self._rules.append(rule)

    @property
    def rules(self) -> List[RewriteRule]:
        """Return a copy of the current rule list."""
        return list(self._rules)
