"""Message and response inspector for scandar-guard — enhanced with decoding and multi-turn."""
from __future__ import annotations

import time
from typing import List, Optional

from .decoder import ContentDecoder
from .findings import FindingSource, InspectionResult, RuntimeFinding
from .multi_turn import MultiTurnTracker
from .patterns import (
    COMPLETION_ATTACK_PATTERNS,
    PII_PATTERNS,
    PROMPT_INJECTION_PATTERNS,
    SHELL_INJECTION_PATTERNS,
    PatternDef,
)

_decoder = ContentDecoder()


class MessageInspector:
    def __init__(self) -> None:
        self._multi_turn: Optional[MultiTurnTracker] = None

    def set_multi_turn_tracker(self, tracker: MultiTurnTracker) -> None:
        self._multi_turn = tracker

    def inspect_message(
        self, role: str, content: str, session_id: str
    ) -> InspectionResult:
        """Inspect an inbound message with decoding and multi-turn tracking."""
        start = time.perf_counter()
        findings: List[RuntimeFinding] = []

        role_lower = role.lower()

        # ── Decode obfuscated content first ───────────────────────────
        decoded_variants, evasion_findings = _decoder.check_evasion_attempt(content, session_id)
        findings.extend(evasion_findings)

        # Build list of texts to scan (original + decoded variants)
        texts_to_scan = [content] + decoded_variants

        if role_lower == "user":
            for text in texts_to_scan:
                findings.extend(self._run_patterns(text, PROMPT_INJECTION_PATTERNS, "message", session_id))
        elif role_lower == "assistant":
            for text in texts_to_scan:
                findings.extend(self._run_patterns(text, COMPLETION_ATTACK_PATTERNS, "response", session_id))
                findings.extend(self._run_patterns(text, PII_PATTERNS, "response", session_id))
        elif role_lower in ("tool", "function"):
            for text in texts_to_scan:
                findings.extend(self._run_patterns(text, PROMPT_INJECTION_PATTERNS, "tool_result", session_id))
                findings.extend(self._run_patterns(text, SHELL_INJECTION_PATTERNS, "tool_result", session_id))

        # ── Multi-turn tracking ───────────────────────────────────────
        if self._multi_turn and role_lower in ("user", "tool"):
            findings.extend(self._multi_turn.record_message(content, role_lower))

        # Deduplicate by category (same category found in original + decoded = keep one)
        findings = self._deduplicate(findings)

        return InspectionResult(
            findings=findings,
            session_id=session_id,
            event_type=f"message_inspected:{role_lower}",
            duration_ms=(time.perf_counter() - start) * 1000,
        )

    def inspect_response(self, content: str, session_id: str) -> InspectionResult:
        """Inspect a model response with decoding."""
        start = time.perf_counter()
        findings: List[RuntimeFinding] = []

        decoded_variants, evasion_findings = _decoder.check_evasion_attempt(content, session_id)
        findings.extend(evasion_findings)

        texts_to_scan = [content] + decoded_variants
        for text in texts_to_scan:
            findings.extend(self._run_patterns(text, COMPLETION_ATTACK_PATTERNS, "response", session_id))
            findings.extend(self._run_patterns(text, PII_PATTERNS, "response", session_id))

        findings = self._deduplicate(findings)

        return InspectionResult(
            findings=findings,
            session_id=session_id,
            event_type="response_inspected",
            duration_ms=(time.perf_counter() - start) * 1000,
        )

    def _run_patterns(
        self,
        text: str,
        patterns: list[PatternDef],
        source: FindingSource,
        session_id: str,
    ) -> List[RuntimeFinding]:
        findings: List[RuntimeFinding] = []
        inspect_text = text[:20_000] if len(text) > 20_000 else text
        for p in patterns:
            match = p.pattern.search(inspect_text)
            if match:
                findings.append(
                    RuntimeFinding(
                        session_id=session_id,
                        severity=p.severity,
                        category=p.category,
                        title=p.title,
                        description=f"Detected in {source}: {p.title.lower()}",
                        source=source,
                        confidence=p.confidence,
                        context={
                            "matched_excerpt": self._truncate_context(
                                inspect_text, match.start(), match.end()
                            )
                        },
                    )
                )
        return findings

    def _truncate_context(self, text: str, start: int, end: int, window: int = 80) -> str:
        ctx_start = max(0, start - window)
        ctx_end = min(len(text), end + window)
        excerpt = text[ctx_start:ctx_end]
        if ctx_start > 0:
            excerpt = "..." + excerpt
        if ctx_end < len(text):
            excerpt = excerpt + "..."
        return excerpt

    def _deduplicate(self, findings: List[RuntimeFinding]) -> List[RuntimeFinding]:
        """Keep only the highest-confidence finding per category."""
        best: dict[str, RuntimeFinding] = {}
        severity_rank = {"critical": 0, "high": 1, "medium": 2, "low": 3, "info": 4}
        for f in findings:
            key = f.category
            if key not in best:
                best[key] = f
            else:
                existing = best[key]
                existing_rank = severity_rank.get(existing.severity, 5)
                new_rank = severity_rank.get(f.severity, 5)
                if new_rank < existing_rank or (new_rank == existing_rank and f.confidence > existing.confidence):
                    best[key] = f
        return list(best.values())
