"""
ChatFeatureExtractor — zero-cost feature extraction from (query, response) pairs.

Computes 12 lexical and structural features that characterise response quality
and uncertainty without any LLM calls.  All features are derived from simple
regex patterns and word-level statistics, making extraction essentially free
at inference time.

Typical use::

    extractor = ChatFeatureExtractor()
    vec = extractor.extract(query, response)
    arr = vec.to_array()   # 12-element normalised list for downstream models
Purpose:    Computes 12 lexical and structural features that characterise response quality and uncertainty without any LLM calls. Feeds into ChatGuard's native scoring mode.
Exports:    ChatFeatureExtractor, ChatFeatureVector
Added:      v0.26.0
Status:     Stable
"""

from __future__ import annotations

import math
import re
from dataclasses import dataclass
from typing import Dict, List

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_HEDGE_WORDS: frozenset = frozenset({
    "think", "probably", "perhaps", "maybe", "might", "possibly", "could",
    "seem", "seems", "apparently", "arguably", "presumably", "likely",
    "unlikely", "generally", "often", "sometimes", "usually", "typically",
    "roughly", "approximately", "around", "about", "somewhat", "fairly",
    "relatively", "tend", "tends",
})
# Note: 'might' and 'could' appear in both _HEDGE_WORDS and _CONDITIONAL_RE.
# This means hedge_density and conditional_density are not fully independent.
# Known collinearity — acceptable for LogReg but should be noted.

_UNCERTAINTY_RE = re.compile(
    r"\b(i don'?t|i cannot|i can'?t|i won'?t|i'm not sure|i am not sure"
    r"|i don'?t know|i cannot confirm|i'?m unsure|not certain|unclear to me)\b",
    re.IGNORECASE,
)

_SENTENCE_SPLIT_RE = re.compile(r"[.!?]+")
_CONDITIONAL_RE = re.compile(r"\b(if|might|could|would|may)\b", re.IGNORECASE)
_NUMERIC_RE = re.compile(r"\b\d[\d,.%]*\b")


# ---------------------------------------------------------------------------
# Cosine helper (lazy MiniLM, falls back gracefully)
# ---------------------------------------------------------------------------

def _cosine(text_a: str, text_b: str) -> float:
    try:
        from llm_guard.scoring.mini_judge import _SentenceModel
        return _SentenceModel.get().similarity(text_a, text_b)
    except Exception:
        return 0.5


# ---------------------------------------------------------------------------
# Feature vector dataclass
# ---------------------------------------------------------------------------

@dataclass
class ChatFeatureVector:
    """Raw feature values extracted from a (query, response) pair.

    All fields store unscaled values.  Call ``to_array()`` to get a
    normalised 12-element list suitable for feeding into a linear model.
    """

    hedge_density: float
    length_ratio: float
    first_person_uncertainty: float
    query_response_cosine: float
    entity_overlap_ratio: float
    conditional_density: float
    lexical_diversity: float
    response_word_count: int
    query_word_count: int
    uppercase_ratio: float
    numeric_density: float
    sentence_count: int

    # ------------------------------------------------------------------
    # Serialisation helpers
    # ------------------------------------------------------------------

    def to_array(self) -> List[float]:
        """Return a 12-element normalised list for downstream models.

        Normalisation per element:
        0  hedge_density              — raw
        1  length_ratio               — min(x / 20, 1)
        2  first_person_uncertainty   — min(x / 3, 1)
        3  query_response_cosine      — raw
        4  entity_overlap_ratio       — raw
        5  conditional_density        — raw
        6  lexical_diversity          — raw
        7  response_word_count        — log1p(x) / 7
        8  query_word_count           — log1p(x) / 5
        9  uppercase_ratio            — raw
        10 numeric_density            — raw
        11 sentence_count             — log1p(x) / 4
        """
        return [
            self.hedge_density,
            min(self.length_ratio / 20.0, 1.0),
            min(self.first_person_uncertainty / 3.0, 1.0),
            self.query_response_cosine,
            self.entity_overlap_ratio,
            self.conditional_density,
            self.lexical_diversity,
            math.log1p(self.response_word_count) / 7.0,
            math.log1p(self.query_word_count) / 5.0,
            self.uppercase_ratio,
            self.numeric_density,
            math.log1p(self.sentence_count) / 4.0,
        ]

    def to_dict(self) -> Dict[str, object]:
        # Explicit dict for API stability — equivalent to dataclasses.asdict(self)
        # Kept for consistency with other result types in llm_guard (ChainTrustResult, etc.)
        """Return all 12 raw field values as a plain dict."""
        return {
            "hedge_density": self.hedge_density,
            "length_ratio": self.length_ratio,
            "first_person_uncertainty": self.first_person_uncertainty,
            "query_response_cosine": self.query_response_cosine,
            "entity_overlap_ratio": self.entity_overlap_ratio,
            "conditional_density": self.conditional_density,
            "lexical_diversity": self.lexical_diversity,
            "response_word_count": self.response_word_count,
            "query_word_count": self.query_word_count,
            "uppercase_ratio": self.uppercase_ratio,
            "numeric_density": self.numeric_density,
            "sentence_count": self.sentence_count,
        }


# ---------------------------------------------------------------------------
# Extractor
# ---------------------------------------------------------------------------

class ChatFeatureExtractor:
    """Compute 12 zero-cost features from a (query, response) text pair.

    The extractor holds no mutable state after construction and is therefore
    safe to share across threads.

    Example::

        extractor = ChatFeatureExtractor()
        vec = extractor.extract("What is the capital of France?", "Paris.")
        print(vec.hedge_density)   # 0.0
    """

    # The instance is stateless; __init__ is a no-op kept for API clarity.
    def __init__(self) -> None:
        pass

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

    def extract(self, query: str, response: str) -> ChatFeatureVector:
        """Extract features from a query/response pair.

        Parameters
        ----------
        query:
            The user's input text.
        response:
            The model's reply text.

        Returns
        -------
        ChatFeatureVector
            All 12 raw features.
        """
        q_words = self._tokenize(query)
        r_words = self._tokenize(response)

        q_word_count = len(q_words)
        r_word_count = len(r_words)

        # Sentences (non-empty segments after splitting on .!?)
        sentences = [s.strip() for s in _SENTENCE_SPLIT_RE.split(response) if s.strip()]
        sentence_count = max(len(sentences), 1)

        # 1. Hedge density
        r_lower = [w.lower() for w in r_words]
        hedge_count = sum(1 for w in r_lower if w in _HEDGE_WORDS)
        hedge_density = hedge_count / max(r_word_count, 1)

        # 2. Length ratio
        length_ratio = r_word_count / max(q_word_count, 1)

        # 3. First-person uncertainty (matches per sentence)
        uncertainty_matches = len(_UNCERTAINTY_RE.findall(response))
        first_person_uncertainty = uncertainty_matches / sentence_count

        # 4. Query-response cosine (lazy MiniLM, fallback 0.5)
        query_response_cosine = _cosine(query, response)

        # 5. Entity overlap ratio
        q_caps = {w for w in q_words if len(w) > 1 and w[0].isupper()}
        r_caps = {w for w in r_words if len(w) > 1 and w[0].isupper()}
        if r_caps:
            entity_overlap_ratio = len(q_caps & r_caps) / len(r_caps)
        else:
            entity_overlap_ratio = 0.0

        # 6. Conditional density (fraction of sentences with conditional words)
        cond_sentences = sum(
            1 for s in sentences if _CONDITIONAL_RE.search(s)
        )
        conditional_density = cond_sentences / sentence_count

        # 7. Lexical diversity
        lexical_diversity = len(set(r_lower)) / max(r_word_count, 1)

        # 8. Uppercase ratio (words that are all-caps and length >= 2)
        upper_count = sum(1 for w in r_words if len(w) >= 2 and w.isupper())
        uppercase_ratio = upper_count / max(r_word_count, 1)

        # 9. Numeric density
        numeric_hits = len(_NUMERIC_RE.findall(response))
        numeric_density = numeric_hits / max(r_word_count, 1)

        return ChatFeatureVector(
            hedge_density=hedge_density,
            length_ratio=length_ratio,
            first_person_uncertainty=first_person_uncertainty,
            query_response_cosine=query_response_cosine,
            entity_overlap_ratio=entity_overlap_ratio,
            conditional_density=conditional_density,
            lexical_diversity=lexical_diversity,
            response_word_count=r_word_count,
            query_word_count=q_word_count,
            uppercase_ratio=uppercase_ratio,
            numeric_density=numeric_density,
            sentence_count=sentence_count,
        )

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

    @staticmethod
    def _tokenize(text: str) -> List[str]:
        """Split text into word tokens (alphanumeric runs only)."""
        return re.findall(r"[A-Za-z0-9']+", text)
