"""
Starvex NLI Engine - Natural Language Inference for policy compliance

Uses NLI models to check if agent responses contradict business policies.
"""

import logging
from dataclasses import dataclass, field
from typing import Dict, Any, List, Optional, Tuple

logger = logging.getLogger(__name__)


@dataclass
class NLIResult:
    """Result of NLI check"""
    entails: bool
    contradicts: bool
    neutral: bool
    scores: Dict[str, float] = field(default_factory=dict)


class ConversationContext:
    """Tracks multi-turn conversation history for context-aware policy checking"""

    def __init__(self, max_turns: int = 10):
        self.max_turns = max_turns
        self._history: List[Dict[str, str]] = []

    def add_turn(self, role: str, content: str) -> None:
        """Add a conversation turn"""
        self._history.append({"role": role, "content": content})
        if len(self._history) > self.max_turns:
            self._history = self._history[-self.max_turns:]

    def get_context_summary(self) -> str:
        """Get a summary of the conversation context"""
        if not self._history:
            return ""
        parts = []
        for turn in self._history[-5:]:  # Last 5 turns
            parts.append(f"{turn['role']}: {turn['content'][:200]}")
        return "\n".join(parts)

    def get_user_intent(self) -> Optional[str]:
        """Extract the likely user intent from conversation history"""
        user_messages = [
            t["content"] for t in self._history if t["role"] == "user"
        ]
        if not user_messages:
            return None
        return user_messages[-1]

    def clear(self) -> None:
        self._history = []

    @property
    def turn_count(self) -> int:
        return len(self._history)


class NLIEngine:
    """
    Natural Language Inference engine for policy compliance checking.

    Uses transformer models to determine if agent responses
    contradict specified business policies. Supports conversation
    context for multi-turn dialogue tracking.
    """

    def __init__(
        self,
        model_name: str = "facebook/bart-large-mnli",
        enable_context: bool = False,
        max_context_turns: int = 10,
    ):
        """
        Initialize NLI engine.

        Args:
            model_name: HuggingFace model name for NLI
            enable_context: Enable conversation context tracking
            max_context_turns: Maximum conversation turns to track
        """
        self.model_name = model_name
        self._model = None
        self._tokenizer = None
        self.rules: List[str] = []
        self.enable_context = enable_context
        self._context = ConversationContext(max_turns=max_context_turns) if enable_context else None

    def _ensure_model(self):
        """Load model if not already loaded"""
        if self._model is not None:
            return

        try:
            from transformers import AutoModelForSequenceClassification, AutoTokenizer
            import torch

            self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self._model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
            self._torch = torch
            logger.debug(f"NLI model loaded: {self.model_name}")

        except ImportError:
            logger.warning("transformers not installed, NLI disabled")
            raise ImportError("transformers and torch required for NLI")

    def add_rule(self, rule: str) -> None:
        """Add a policy rule to check against"""
        self.rules.append(rule)

    def clear_rules(self) -> None:
        """Clear all policy rules"""
        self.rules = []

    def check(self, premise: str, hypothesis: str) -> NLIResult:
        """
        Check if hypothesis is entailed by, contradicts, or is neutral to premise.

        Args:
            premise: The base statement (e.g., policy)
            hypothesis: The statement to check (e.g., agent response)

        Returns:
            NLIResult with entailment status
        """
        self._ensure_model()

        try:
            # Tokenize
            inputs = self._tokenizer(
                premise,
                hypothesis,
                return_tensors="pt",
                truncation=True,
                max_length=512,
            )

            # Get predictions
            with self._torch.no_grad():
                outputs = self._model(**inputs)
                logits = outputs.logits
                probs = self._torch.softmax(logits, dim=1)[0]

            # BART-MNLI labels: contradiction, neutral, entailment
            scores = {
                "contradiction": float(probs[0]),
                "neutral": float(probs[1]),
                "entailment": float(probs[2]),
            }

            return NLIResult(
                entails=scores["entailment"] > 0.5,
                contradicts=scores["contradiction"] > 0.5,
                neutral=scores["neutral"] > 0.5,
                scores=scores,
            )

        except Exception as e:
            logger.error(f"NLI check error: {e}")
            return NLIResult(
                entails=False,
                contradicts=False,
                neutral=True,
                scores={"error": str(e)},
            )

    def check_against_rules(self, text: str) -> Tuple[bool, List[NLIResult]]:
        """
        Check text against all policy rules.

        Args:
            text: Text to check (typically agent response)

        Returns:
            Tuple of (all_passed, list of results)
        """
        if not self.rules:
            return True, []

        results = []
        all_passed = True

        for rule in self.rules:
            result = self.check(rule, text)
            results.append(result)

            if result.contradicts:
                all_passed = False

        return all_passed, results

    def check_with_context(
        self,
        input_text: str,
        output_text: str,
    ) -> Tuple[bool, List[NLIResult]]:
        """
        Check output against rules with conversation context for better accuracy.

        Enhances policy checking by considering the full conversation history,
        enabling detection of multi-turn policy violations.

        Args:
            input_text: User's input message
            output_text: Agent's response to check

        Returns:
            Tuple of (all_passed, list of results)
        """
        if not self.rules:
            return True, []

        # Add turn to context
        if self._context:
            self._context.add_turn("user", input_text)
            self._context.add_turn("assistant", output_text)

        results = []
        all_passed = True

        for rule in self.rules:
            # Build context-enriched premise
            if self._context and self._context.turn_count > 2:
                context_summary = self._context.get_context_summary()
                enriched_premise = (
                    f"Policy: {rule}\n"
                    f"Conversation context:\n{context_summary}"
                )
            else:
                enriched_premise = rule

            result = self.check(enriched_premise, output_text)
            results.append(result)

            if result.contradicts:
                all_passed = False

        return all_passed, results

    def add_conversation_turn(self, role: str, content: str) -> None:
        """Add a turn to the conversation context"""
        if self._context:
            self._context.add_turn(role, content)

    def clear_context(self) -> None:
        """Clear conversation context"""
        if self._context:
            self._context.clear()


# Alias
PolicyChecker = NLIEngine
