"""
AgentDrop Shield -- Layer 4b: Intent Classifier.

Lightweight statistical/heuristic classifier that scores whether text
is attempting to instruct an AI agent.  No ML models required -- uses
signal analysis across multiple dimensions:

    1. Imperative verb density (commands)
    2. Second-person pronoun ratio
    3. Instruction framing phrases
    4. Question-to-command ratio
    5. System/role keyword density
    6. Code/technical injection markers

Returns a confidence score 0.0--1.0 with a breakdown of which signals
fired and their individual contributions.

Uses only Python stdlib.
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from typing import Optional

logger = logging.getLogger("agentdrop.shield.intent_classifier")


# =====================================================================
# Signal Definitions
# =====================================================================

# Imperative verbs commonly used in injection attempts.
# Weighted: high-risk verbs that almost never appear in normal documents
# get a higher weight than ambiguous ones.
_IMPERATIVE_VERBS_HIGH: frozenset[str] = frozenset([
    "ignore", "disregard", "forget", "override", "bypass",
    "pretend", "impersonate", "jailbreak", "exfiltrate",
    "obey", "comply",
])

_IMPERATIVE_VERBS_MEDIUM: frozenset[str] = frozenset([
    "execute", "output", "respond", "behave", "act",
    "reveal", "display", "print", "send", "transmit",
    "generate", "produce", "simulate", "enable", "activate",
    "switch", "enter",
])

_IMPERATIVE_VERBS_LOW: frozenset[str] = frozenset([
    "do", "run", "follow", "listen", "write", "tell",
    "give", "show", "provide", "create", "make", "set",
    "use", "add", "include", "ensure",
])

# Second-person pronouns and determiners
_SECOND_PERSON: frozenset[str] = frozenset([
    "you", "your", "yours", "yourself", "you're", "you've",
    "you'll", "you'd",
])

# Instruction framing phrases (compiled regexes)
_FRAMING_PATTERNS: list[tuple[re.Pattern, float, str]] = [
    (re.compile(r"\bfrom\s+now\s+on\b", re.IGNORECASE), 0.9, "from_now_on"),
    (re.compile(r"\byour\s+task\s+is\b", re.IGNORECASE), 0.85, "your_task_is"),
    (re.compile(r"\byou\s+will\b", re.IGNORECASE), 0.6, "you_will"),
    (re.compile(r"\byou\s+must\b", re.IGNORECASE), 0.75, "you_must"),
    (re.compile(r"\byou\s+should\b", re.IGNORECASE), 0.5, "you_should"),
    (re.compile(r"\byou\s+are\s+(?:now|going\s+to)\b", re.IGNORECASE), 0.8, "you_are_now"),
    (re.compile(r"\bi\s+want\s+you\s+to\b", re.IGNORECASE), 0.7, "i_want_you_to"),
    (re.compile(r"\bplease\s+ensure\b", re.IGNORECASE), 0.5, "please_ensure"),
    (re.compile(r"\bdo\s+not\s+(?:ever|under\s+any)\b", re.IGNORECASE), 0.7, "do_not_ever"),
    (re.compile(r"\bnew\s+instructions?\s*[:=]\b", re.IGNORECASE), 0.95, "new_instructions"),
    (re.compile(r"\byour\s+(?:new|real|actual|true)\s+(?:role|purpose|instructions?)\b", re.IGNORECASE), 0.9, "your_new_role"),
    (re.compile(r"\bact\s+as\s+(?:a|an|if)\b", re.IGNORECASE), 0.7, "act_as"),
    (re.compile(r"\bpretend\s+(?:you\s+are|to\s+be|that)\b", re.IGNORECASE), 0.8, "pretend_you_are"),
    (re.compile(r"\bwhen\s+(?:the\s+)?user\b", re.IGNORECASE), 0.65, "when_user"),
    (re.compile(r"\brespond\s+(?:only\s+)?(?:with|in|as)\b", re.IGNORECASE), 0.6, "respond_with"),
    (re.compile(r"\bhenceforth\b", re.IGNORECASE), 0.85, "henceforth"),
    (re.compile(r"\bgoing\s+forward\b", re.IGNORECASE), 0.6, "going_forward"),
    (re.compile(r"\bfor\s+(?:the\s+)?(?:rest\s+of|remainder)\b", re.IGNORECASE), 0.6, "for_the_rest"),
]

# System/role keywords that indicate AI-instruction context
_SYSTEM_KEYWORDS: frozenset[str] = frozenset([
    "system", "assistant", "prompt", "instructions", "context",
    "llm", "chatgpt", "gpt", "claude", "model", "ai",
    "jailbreak", "dan", "developer mode", "unrestricted",
])

# Regex to detect XML-like tags, backtick blocks, bracket notation
# in non-code files -- signals of delimiter injection
_INJECTION_MARKER_PATTERNS: list[tuple[re.Pattern, float, str]] = [
    (re.compile(r"<\|(?:im_start|im_end|system|user|assistant|endoftext)\|>"), 1.0, "chatml_delimiters"),
    (re.compile(r"\[INST\]|\[/INST\]|<<SYS>>|<</SYS>>"), 0.95, "llama_delimiters"),
    (re.compile(r"```\s*system\b"), 0.9, "code_block_system"),
    (re.compile(r"<\s*(?:system|instruction|prompt)\s*>", re.IGNORECASE), 0.8, "xml_system_tags"),
    (re.compile(r"###\s*(?:System|Instruction|Human|Assistant)\s*:", re.IGNORECASE), 0.75, "markdown_role_markers"),
]


# =====================================================================
# Result Model
# =====================================================================

@dataclass
class IntentSignal:
    """A single signal that contributed to the intent score."""
    signal_name: str
    score: float
    detail: str


@dataclass
class IntentClassification:
    """Result of the intent classification analysis."""

    intent_score: float = 0.0
    """Overall probability 0.0--1.0 that the text is trying to instruct an AI."""

    signals: list[IntentSignal] = field(default_factory=list)
    """Breakdown of individual signals that fired."""

    is_instructive: bool = False
    """True if intent_score exceeds the configured threshold."""


# =====================================================================
# Classifier
# =====================================================================

class IntentClassifier:
    """Layer 4b -- heuristic intent classifier for AI-instructive text.

    Analyzes text across multiple signal dimensions and returns a
    weighted composite score indicating the probability that the text
    is attempting to instruct/manipulate an AI agent.

    Thread-safe: holds no mutable state between calls.
    """

    # Signal weights for the composite score
    WEIGHT_IMPERATIVE_VERBS = 0.15
    WEIGHT_SECOND_PERSON = 0.15
    WEIGHT_FRAMING_PHRASES = 0.25
    WEIGHT_QUESTION_RATIO = 0.10
    WEIGHT_SYSTEM_KEYWORDS = 0.15
    WEIGHT_INJECTION_MARKERS = 0.20

    def __init__(self, threshold: float = 0.45) -> None:
        """
        Args:
            threshold: Score above which ``is_instructive`` is True.
        """
        self._threshold = threshold

    # Sliding window parameters for large-document dilution defense.
    # Window size in characters and step size (overlap = window - step).
    WINDOW_SIZE = 500
    WINDOW_STEP = 300
    # Minimum document length (chars) to trigger windowed analysis.
    # Below this threshold, global scoring is sufficient.
    WINDOW_MIN_DOC_LENGTH = 1500

    def classify(self, texts: list[tuple[str, str]]) -> IntentClassification:
        """Classify whether extracted text is AI-instructive.

        Uses a two-pass approach for large documents:
        1. Global pass: score the full document (original behavior).
        2. Windowed pass: score overlapping windows and take the max.

        The final score is max(global, best_window) to defeat dilution
        attacks where malicious lines are buried in large clean documents.

        Args:
            texts: List of (source_label, normalized_text) from Layer 3.

        Returns:
            IntentClassification with composite score and signal breakdown.
        """
        if not texts:
            return IntentClassification()

        # Combine all text for analysis
        combined = "\n".join(text for _, text in texts if text.strip())
        if not combined.strip():
            return IntentClassification()

        # --- Pass 1: Global scoring (original behavior) ---
        global_result = self._classify_text(combined)

        # --- Pass 2: Sliding window scoring for large documents ---
        if len(combined) >= self.WINDOW_MIN_DOC_LENGTH:
            window_result = self._classify_windowed(combined)
            if window_result.intent_score > global_result.intent_score:
                # Window found a hotspot -- use its score but merge signals
                merged_signals = list(global_result.signals)
                merged_signals.append(IntentSignal(
                    "sliding_window_hotspot",
                    window_result.intent_score,
                    f"Windowed analysis detected concentrated injection "
                    f"region (window_score={window_result.intent_score:.2f} "
                    f"vs global={global_result.intent_score:.2f})",
                ))
                return IntentClassification(
                    intent_score=window_result.intent_score,
                    signals=merged_signals,
                    is_instructive=window_result.intent_score >= self._threshold,
                )

        return global_result

    def _classify_text(self, text: str) -> IntentClassification:
        """Score a single block of text (the core classification logic).

        This is the original classify() logic, extracted so it can be
        reused for both global and per-window analysis.
        """
        signals: list[IntentSignal] = []

        # --- Signal 1: Imperative verb density ---
        verb_score = self._score_imperative_verbs(text)
        if verb_score > 0.0:
            signals.append(IntentSignal(
                "imperative_verb_density", verb_score,
                "Elevated density of command-style verbs",
            ))

        # --- Signal 2: Second-person pronoun ratio ---
        pronoun_score = self._score_second_person(text)
        if pronoun_score > 0.0:
            signals.append(IntentSignal(
                "second_person_ratio", pronoun_score,
                "High ratio of second-person pronouns (you/your)",
            ))

        # --- Signal 3: Instruction framing phrases ---
        framing_score, framing_matches = self._score_framing_phrases(text)
        if framing_score > 0.0:
            detail = f"Framing phrases detected: {', '.join(framing_matches)}"
            signals.append(IntentSignal(
                "instruction_framing", framing_score, detail,
            ))

        # --- Signal 4: Question-to-command ratio ---
        question_score = self._score_question_ratio(text)
        if question_score > 0.0:
            signals.append(IntentSignal(
                "command_heavy_ratio", question_score,
                "Text is predominantly commands, not questions",
            ))

        # --- Signal 5: System/role keyword density ---
        keyword_score = self._score_system_keywords(text)
        if keyword_score > 0.0:
            signals.append(IntentSignal(
                "system_keyword_density", keyword_score,
                "AI/system role keywords detected",
            ))

        # --- Signal 6: Code/technical injection markers ---
        marker_score, marker_matches = self._score_injection_markers(text)
        if marker_score > 0.0:
            detail = f"Injection markers: {', '.join(marker_matches)}"
            signals.append(IntentSignal(
                "injection_markers", marker_score, detail,
            ))

        # --- Composite score (weighted average) ---
        composite = (
            verb_score * self.WEIGHT_IMPERATIVE_VERBS
            + pronoun_score * self.WEIGHT_SECOND_PERSON
            + framing_score * self.WEIGHT_FRAMING_PHRASES
            + question_score * self.WEIGHT_QUESTION_RATIO
            + keyword_score * self.WEIGHT_SYSTEM_KEYWORDS
            + marker_score * self.WEIGHT_INJECTION_MARKERS
        )

        # Boost: if multiple strong signals co-occur, the text is
        # almost certainly instructive.  Apply a multiplicative boost.
        strong_signals = sum(1 for s in signals if s.score >= 0.6)
        if strong_signals >= 3:
            composite = min(composite * 1.3, 1.0)
        elif strong_signals >= 2:
            composite = min(composite * 1.15, 1.0)

        composite = round(min(composite, 1.0), 4)

        return IntentClassification(
            intent_score=composite,
            signals=signals,
            is_instructive=composite >= self._threshold,
        )

    def _classify_windowed(self, text: str) -> IntentClassification:
        """Run sliding window analysis and return the highest-scoring window.

        Slides a window of WINDOW_SIZE characters across the text with
        WINDOW_STEP overlap.  Returns the IntentClassification from the
        window with the highest composite score.

        Performance: for a 50,000-char document with 500-char windows
        and 250-char steps, this produces ~200 windows.  Each window
        analysis is O(n) where n = window size, so total overhead is
        O(doc_size) -- linear and bounded.
        """
        best_result = IntentClassification()

        pos = 0
        text_len = len(text)
        while pos < text_len:
            end = min(pos + self.WINDOW_SIZE, text_len)
            window = text[pos:end]

            # Only score windows that have meaningful content
            if len(window.strip()) > 50:
                result = self._classify_text(window)
                if result.intent_score > best_result.intent_score:
                    best_result = result

            pos += self.WINDOW_STEP

            # Early exit: if we already found a very high score,
            # no need to keep scanning.
            if best_result.intent_score >= 0.85:
                break

        return best_result

    # ------------------------------------------------------------------
    # Signal scorers
    # ------------------------------------------------------------------

    @staticmethod
    def _tokenize_lower(text: str) -> list[str]:
        """Split text into lowercase word tokens."""
        return re.findall(r"[a-z']+", text.lower())

    def _score_imperative_verbs(self, text: str) -> float:
        """Score 0.0--1.0 based on density of imperative verbs."""
        words = self._tokenize_lower(text)
        if not words:
            return 0.0

        high = sum(1 for w in words if w in _IMPERATIVE_VERBS_HIGH)
        medium = sum(1 for w in words if w in _IMPERATIVE_VERBS_MEDIUM)
        low = sum(1 for w in words if w in _IMPERATIVE_VERBS_LOW)

        # Weighted count: high-risk verbs count more
        weighted = high * 3.0 + medium * 1.5 + low * 0.5
        density = weighted / len(words)

        # Normalize: 0.05 density -> 0.5 score, 0.15+ -> 1.0
        score = min(density / 0.15, 1.0)
        return round(score, 4)

    def _score_second_person(self, text: str) -> float:
        """Score 0.0--1.0 based on second-person pronoun density."""
        words = self._tokenize_lower(text)
        if not words:
            return 0.0

        count = sum(1 for w in words if w in _SECOND_PERSON)
        ratio = count / len(words)

        # Normal documents rarely exceed 2% second-person.
        # Instructive text often hits 5-15%.
        # 0.03 -> 0.3, 0.08+ -> 1.0
        score = min(ratio / 0.08, 1.0)
        return round(score, 4)

    @staticmethod
    def _score_framing_phrases(text: str) -> tuple[float, list[str]]:
        """Score 0.0--1.0 based on instruction framing phrases found."""
        matches: list[str] = []
        max_score = 0.0

        for pattern, weight, name in _FRAMING_PATTERNS:
            if pattern.search(text):
                matches.append(name)
                max_score = max(max_score, weight)

        if not matches:
            return 0.0, []

        # Combine: use max score, boosted by count of distinct phrases
        count_boost = min(len(matches) * 0.08, 0.3)
        score = min(max_score + count_boost, 1.0)
        return round(score, 4), matches

    @staticmethod
    def _score_question_ratio(text: str) -> float:
        """Score 0.0--1.0 based on question-to-statement ratio.

        High score = mostly commands (few questions). Low score = mostly
        questions (less suspicious). A purely Q&A document would score low.
        """
        sentences = re.split(r"[.!?\n]+", text)
        sentences = [s.strip() for s in sentences if len(s.strip()) > 3]
        if not sentences:
            return 0.0

        questions = sum(1 for s in sentences if s.rstrip().endswith("?")
                        or s.lower().startswith(("who ", "what ", "when ",
                                                  "where ", "why ", "how ",
                                                  "is ", "are ", "do ", "does ",
                                                  "can ", "could ", "will ")))
        question_ratio = questions / len(sentences)

        # Invert: fewer questions = higher score (more commands)
        # If >50% questions, score is 0. If 0% questions, score depends
        # on sentence count (need at least a few sentences to be meaningful).
        if question_ratio > 0.5:
            return 0.0
        if len(sentences) < 3:
            return 0.0  # Too short to be meaningful

        score = 1.0 - (question_ratio * 2.0)
        # Scale down -- this signal alone shouldn't dominate
        score = score * 0.7
        return round(max(score, 0.0), 4)

    def _score_system_keywords(self, text: str) -> float:
        """Score 0.0--1.0 based on AI/system keyword density."""
        text_lower = text.lower()
        words = self._tokenize_lower(text)
        if not words:
            return 0.0

        count = 0
        for kw in _SYSTEM_KEYWORDS:
            count += text_lower.count(kw)

        density = count / len(words)
        # 0.02 -> 0.4, 0.05+ -> 1.0
        score = min(density / 0.05, 1.0)
        return round(score, 4)

    @staticmethod
    def _score_injection_markers(text: str) -> tuple[float, list[str]]:
        """Score 0.0--1.0 based on presence of injection delimiter markers."""
        matches: list[str] = []
        max_score = 0.0

        for pattern, weight, name in _INJECTION_MARKER_PATTERNS:
            if pattern.search(text):
                matches.append(name)
                max_score = max(max_score, weight)

        if not matches:
            return 0.0, []

        # Any injection marker is a strong signal
        count_boost = min((len(matches) - 1) * 0.1, 0.2)
        score = min(max_score + count_boost, 1.0)
        return round(score, 4), matches
