"""
LangChain Callback Handler for automatic threat scanning.

Automatically scans incoming messages for threats and can block
suspicious content before it reaches the LLM.
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Union
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage

from agent_trust import AgentTrustClient, Verdict, ThreatLevel


logger = logging.getLogger(__name__)


class ThreatDetectedError(Exception):
    """Raised when a threat is detected and blocking is enabled."""
    
    def __init__(
        self,
        message: str,
        verdict: Verdict,
        threat_level: ThreatLevel,
        threats: list,
        reasoning: str = ""
    ):
        super().__init__(message)
        self.verdict = verdict
        self.threat_level = threat_level
        self.threats = threats
        self.reasoning = reasoning


class TrustVerificationCallback(BaseCallbackHandler):
    """
    Callback handler that automatically scans messages for threats.
    
    Can be configured to:
    - Log detected threats
    - Block messages that exceed a threat threshold
    - Track suspicious agents
    
    Example:
        from langchain_openai import ChatOpenAI
        from agent_trust_langchain import TrustVerificationCallback
        
        callback = TrustVerificationCallback(
            block_on_threat=True,
            min_block_level=ThreatLevel.HIGH
        )
        
        llm = ChatOpenAI(callbacks=[callback])
        
        try:
            response = llm.invoke("Some message")
        except ThreatDetectedError as e:
            print(f"Blocked: {e.reasoning}")
    """
    
    def __init__(
        self,
        api_url: Optional[str] = None,
        api_key: Optional[str] = None,
        block_on_threat: bool = False,
        min_block_level: ThreatLevel = ThreatLevel.HIGH,
        log_threats: bool = True,
        scan_human_messages: bool = True,
        scan_ai_messages: bool = False,
        on_threat_detected: Optional[callable] = None,
    ):
        """
        Initialize the callback handler.
        
        Args:
            api_url: Custom API URL (optional)
            api_key: API key for authentication (optional)
            block_on_threat: If True, raise exception when threat detected
            min_block_level: Minimum threat level to trigger blocking
            log_threats: If True, log detected threats
            scan_human_messages: Scan incoming human messages
            scan_ai_messages: Scan AI responses (usually not needed)
            on_threat_detected: Custom callback when threat detected
        """
        super().__init__()
        self.block_on_threat = block_on_threat
        self.min_block_level = min_block_level
        self.log_threats = log_threats
        self.scan_human_messages = scan_human_messages
        self.scan_ai_messages = scan_ai_messages
        self.on_threat_detected = on_threat_detected
        
        # Initialize client
        kwargs = {}
        if api_url:
            kwargs["api_url"] = api_url
        if api_key:
            kwargs["api_key"] = api_key
        self.client = AgentTrustClient(**kwargs)
        
        # Threat tracking
        self.detected_threats: List[Dict[str, Any]] = []
        self.scanned_count = 0
        self.blocked_count = 0
    
    def _get_threat_level_value(self, level: ThreatLevel) -> int:
        """Convert threat level to numeric value for comparison."""
        levels = {
            ThreatLevel.SAFE: 0,
            ThreatLevel.LOW: 1,
            ThreatLevel.MEDIUM: 2,
            ThreatLevel.HIGH: 3,
            ThreatLevel.CRITICAL: 4,
        }
        return levels.get(level, 0)
    
    def _should_block(self, threat_level: ThreatLevel) -> bool:
        """Check if the threat level should trigger blocking."""
        if not self.block_on_threat:
            return False
        return self._get_threat_level_value(threat_level) >= self._get_threat_level_value(self.min_block_level)
    
    def _scan_text(self, text: str, source: str = "unknown") -> None:
        """Scan text for threats and handle accordingly."""
        if not text or not text.strip():
            return
        
        self.scanned_count += 1
        
        try:
            result = self.client.scan_text(text)
        except Exception as e:
            logger.warning(f"Failed to scan message: {e}")
            return
        
        if result.is_safe:
            return
        
        # Record the threat
        threat_record = {
            "source": source,
            "verdict": result.verdict.value,
            "threat_level": result.threat_level.value,
            "threats": [
                {
                    "name": t.pattern_name,
                    "severity": t.severity.value,
                    "matched_text": t.matched_text[:100] if t.matched_text else None,
                }
                for t in result.threats
            ],
            "reasoning": result.reasoning,
            "text_preview": text[:100] + "..." if len(text) > 100 else text,
        }
        self.detected_threats.append(threat_record)
        
        # Log if enabled
        if self.log_threats:
            threat_names = [t.pattern_name for t in result.threats]
            logger.warning(
                f"Threat detected in {source}: {result.threat_level.value} - "
                f"{threat_names} - {result.reasoning}"
            )
        
        # Custom callback
        if self.on_threat_detected:
            self.on_threat_detected(threat_record)
        
        # Block if configured
        if self._should_block(result.threat_level):
            self.blocked_count += 1
            raise ThreatDetectedError(
                f"Message blocked due to {result.threat_level.value} threat: {result.reasoning}",
                verdict=result.verdict,
                threat_level=result.threat_level,
                threats=result.threats,
                reasoning=result.reasoning,
            )
    
    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Scan messages before they're sent to the chat model."""
        if not self.scan_human_messages:
            return
        
        for message_list in messages:
            for message in message_list:
                if isinstance(message, HumanMessage):
                    content = message.content
                    if isinstance(content, str):
                        self._scan_text(content, source="human_message")
                    elif isinstance(content, list):
                        # Handle multi-modal messages
                        for part in content:
                            if isinstance(part, dict) and part.get("type") == "text":
                                self._scan_text(part.get("text", ""), source="human_message")
                            elif isinstance(part, str):
                                self._scan_text(part, source="human_message")
    
    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Scan prompts before they're sent to the LLM."""
        if not self.scan_human_messages:
            return
        
        for prompt in prompts:
            self._scan_text(prompt, source="prompt")
    
    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Optionally scan AI responses."""
        if not self.scan_ai_messages:
            return
        
        try:
            # Handle different response formats
            if hasattr(response, 'generations'):
                for generation_list in response.generations:
                    for generation in generation_list:
                        if hasattr(generation, 'text'):
                            self._scan_text(generation.text, source="ai_response")
                        elif hasattr(generation, 'message'):
                            content = generation.message.content
                            if isinstance(content, str):
                                self._scan_text(content, source="ai_response")
        except Exception as e:
            logger.debug(f"Could not scan AI response: {e}")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get scanning statistics."""
        return {
            "scanned_count": self.scanned_count,
            "blocked_count": self.blocked_count,
            "threats_detected": len(self.detected_threats),
            "recent_threats": self.detected_threats[-10:],
        }
    
    def clear_history(self) -> None:
        """Clear threat history and reset counters."""
        self.detected_threats = []
        self.scanned_count = 0
        self.blocked_count = 0
