"""
TrustGuard integration for LangChain.

Provides protected document loaders, retrievers, and callbacks
that scan content for threats before processing.
"""

from __future__ import annotations

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

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.documents import Document
from langchain_core.document_loaders import BaseLoader
from langchain_core.retrievers import BaseRetriever
from langchain_core.callbacks import CallbackManagerForRetrieverRun

from agent_trust import TrustGuard, GuardResult, Verdict, ThreatLevel, ContentSource


logger = logging.getLogger(__name__)


class ThreatInDocumentError(Exception):
    """Raised when a threat is detected in a document."""
    
    def __init__(
        self,
        message: str,
        document_source: str,
        guard_result: GuardResult,
    ):
        super().__init__(message)
        self.document_source = document_source
        self.guard_result = guard_result


# ============================================================================
# TrustGuardLoader - Protected Document Loader
# ============================================================================

class TrustGuardLoader(BaseLoader):
    """
    A document loader wrapper that scans documents for threats before returning.
    
    Wraps any LangChain document loader and scans each document using TrustGuard
    before yielding it. Can be configured to block, warn, or filter documents.
    
    Example:
        from langchain_community.document_loaders import WebBaseLoader
        from agent_trust_langchain import TrustGuardLoader
        
        # Wrap a web loader with threat protection
        base_loader = WebBaseLoader(["https://example.com"])
        loader = TrustGuardLoader(
            base_loader,
            api_key="ta_xxx...",
            on_threat="filter",  # Skip documents with threats
        )
        
        docs = loader.load()  # Only returns safe documents
    """
    
    def __init__(
        self,
        loader: BaseLoader,
        api_key: Optional[str] = None,
        api_url: Optional[str] = None,
        on_threat: str = "warn",  # "block", "warn", "filter", "tag"
        min_block_level: ThreatLevel = ThreatLevel.HIGH,
        content_type: ContentSource = ContentSource.DOCUMENT,
        timeout: float = 30.0,
    ):
        """
        Initialize the protected loader.
        
        Args:
            loader: The base document loader to wrap
            api_key: TrustGuard API key
            api_url: Custom API URL (optional)
            on_threat: Action when threat detected:
                - "block": Raise ThreatInDocumentError
                - "warn": Log warning and continue
                - "filter": Skip the document
                - "tag": Add threat info to metadata and continue
            min_block_level: Minimum level to trigger blocking (for "block" mode)
            content_type: Type of content being loaded (document, web, etc.)
            timeout: Request timeout in seconds
        """
        self.loader = loader
        self.on_threat = on_threat
        self.min_block_level = min_block_level
        self.content_type = content_type
        
        # Initialize TrustGuard
        guard_kwargs = {"timeout": timeout}
        if api_key:
            guard_kwargs["api_key"] = api_key
        if api_url:
            guard_kwargs["api_url"] = api_url
        self.guard = TrustGuard(**guard_kwargs)
        
        # Stats
        self.scanned_count = 0
        self.threat_count = 0
        self.filtered_count = 0
    
    def _get_threat_level_value(self, level: ThreatLevel) -> int:
        """Convert threat level to numeric value."""
        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 threat level exceeds blocking threshold."""
        return self._get_threat_level_value(threat_level) >= self._get_threat_level_value(self.min_block_level)
    
    def _scan_document(self, doc: Document) -> GuardResult:
        """Scan a document for threats."""
        self.scanned_count += 1
        
        source = doc.metadata.get("source", "unknown")
        
        # Choose scan method based on content type
        if self.content_type == ContentSource.WEB:
            result = self.guard.scan_web(doc.page_content, source_url=source)
        elif self.content_type == ContentSource.RAG:
            result = self.guard.scan_rag(doc.page_content, source=source, metadata=doc.metadata)
        else:
            filename = doc.metadata.get("filename", doc.metadata.get("source"))
            result = self.guard.scan_document(doc.page_content, filename=filename)
        
        if not result.safe:
            self.threat_count += 1
        
        return result
    
    def _process_document(self, doc: Document) -> Optional[Document]:
        """Process a document, handling threats according to configuration."""
        try:
            result = self._scan_document(doc)
        except Exception as e:
            logger.warning(f"Failed to scan document: {e}")
            return doc  # Return unscanned on error
        
        if result.safe:
            return doc
        
        source = doc.metadata.get("source", "unknown")
        threat_names = [t.pattern_name for t in result.threats]
        
        if self.on_threat == "block" and self._should_block(result.threat_level):
            raise ThreatInDocumentError(
                f"Threat detected in document '{source}': {result.reasoning}",
                document_source=source,
                guard_result=result,
            )
        
        elif self.on_threat == "filter":
            self.filtered_count += 1
            logger.warning(
                f"Filtering document '{source}' due to threats: {threat_names}"
            )
            return None
        
        elif self.on_threat == "tag":
            # Add threat info to metadata
            doc.metadata["trust_guard"] = {
                "safe": False,
                "verdict": result.verdict.value,
                "threat_level": result.threat_level.value,
                "threats": [t.pattern_name for t in result.threats],
                "reasoning": result.reasoning,
            }
            return doc
        
        else:  # "warn" or unknown
            logger.warning(
                f"Threat detected in document '{source}': "
                f"{result.threat_level.value} - {threat_names}"
            )
            return doc
    
    def lazy_load(self) -> Iterator[Document]:
        """Lazily load and scan documents."""
        for doc in self.loader.lazy_load():
            processed = self._process_document(doc)
            if processed is not None:
                yield processed
    
    def load(self) -> List[Document]:
        """Load and scan all documents."""
        return list(self.lazy_load())
    
    def get_stats(self) -> Dict[str, Any]:
        """Get scanning statistics."""
        return {
            "scanned_count": self.scanned_count,
            "threat_count": self.threat_count,
            "filtered_count": self.filtered_count,
        }


# ============================================================================
# TrustGuardRetriever - Protected Retriever
# ============================================================================

class TrustGuardRetriever(BaseRetriever):
    """
    A retriever wrapper that scans retrieved documents for threats.
    
    Wraps any LangChain retriever and scans each retrieved document
    before returning. Useful for protecting RAG pipelines.
    
    Example:
        from langchain_community.vectorstores import Chroma
        from agent_trust_langchain import TrustGuardRetriever
        
        # Create a protected retriever
        base_retriever = vectorstore.as_retriever()
        retriever = TrustGuardRetriever(
            retriever=base_retriever,
            api_key="ta_xxx...",
            on_threat="filter",  # Skip documents with threats
        )
        
        docs = retriever.invoke("my query")  # Only returns safe documents
    """
    
    retriever: BaseRetriever
    """The base retriever to wrap."""
    
    guard: TrustGuard = None
    """TrustGuard instance."""
    
    on_threat: str = "warn"
    """Action on threat: 'block', 'warn', 'filter', 'tag'."""
    
    min_block_level: ThreatLevel = ThreatLevel.HIGH
    """Minimum level for blocking."""
    
    class Config:
        arbitrary_types_allowed = True
    
    def __init__(
        self,
        retriever: BaseRetriever,
        api_key: Optional[str] = None,
        api_url: Optional[str] = None,
        on_threat: str = "warn",
        min_block_level: ThreatLevel = ThreatLevel.HIGH,
        timeout: float = 30.0,
        **kwargs,
    ):
        """
        Initialize the protected retriever.
        
        Args:
            retriever: The base retriever to wrap
            api_key: TrustGuard API key
            api_url: Custom API URL (optional)
            on_threat: Action when threat detected
            min_block_level: Minimum level for blocking
            timeout: Request timeout
        """
        guard_kwargs = {"timeout": timeout}
        if api_key:
            guard_kwargs["api_key"] = api_key
        if api_url:
            guard_kwargs["api_url"] = api_url
        
        super().__init__(
            retriever=retriever,
            guard=TrustGuard(**guard_kwargs),
            on_threat=on_threat,
            min_block_level=min_block_level,
            **kwargs,
        )
    
    def _get_threat_level_value(self, level: ThreatLevel) -> int:
        """Convert threat level to numeric value."""
        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 threat level exceeds blocking threshold."""
        return self._get_threat_level_value(threat_level) >= self._get_threat_level_value(self.min_block_level)
    
    def _scan_document(self, doc: Document) -> GuardResult:
        """Scan a document for threats."""
        source = doc.metadata.get("source", "unknown")
        return self.guard.scan_rag(
            doc.page_content,
            source=source,
            metadata=doc.metadata,
        )
    
    def _process_documents(self, docs: List[Document]) -> List[Document]:
        """Process retrieved documents, handling threats."""
        safe_docs = []
        
        for doc in docs:
            try:
                result = self._scan_document(doc)
            except Exception as e:
                logger.warning(f"Failed to scan retrieved document: {e}")
                safe_docs.append(doc)
                continue
            
            if result.safe:
                safe_docs.append(doc)
                continue
            
            source = doc.metadata.get("source", "unknown")
            threat_names = [t.pattern_name for t in result.threats]
            
            if self.on_threat == "block" and self._should_block(result.threat_level):
                raise ThreatInDocumentError(
                    f"Threat detected in retrieved document '{source}': {result.reasoning}",
                    document_source=source,
                    guard_result=result,
                )
            
            elif self.on_threat == "filter":
                logger.warning(
                    f"Filtering retrieved document '{source}' due to threats: {threat_names}"
                )
                continue  # Skip this document
            
            elif self.on_threat == "tag":
                doc.metadata["trust_guard"] = {
                    "safe": False,
                    "verdict": result.verdict.value,
                    "threat_level": result.threat_level.value,
                    "threats": threat_names,
                    "reasoning": result.reasoning,
                }
                safe_docs.append(doc)
            
            else:  # "warn"
                logger.warning(
                    f"Threat in retrieved document '{source}': "
                    f"{result.threat_level.value} - {threat_names}"
                )
                safe_docs.append(doc)
        
        return safe_docs
    
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: Optional[CallbackManagerForRetrieverRun] = None,
    ) -> List[Document]:
        """Retrieve and scan documents."""
        # Get documents from base retriever
        docs = self.retriever.invoke(query)
        
        # Scan and filter
        return self._process_documents(docs)


# ============================================================================
# TrustGuardCallback - Chain Callback with Guard
# ============================================================================

class TrustGuardCallback(BaseCallbackHandler):
    """
    Callback handler that uses TrustGuard for comprehensive content scanning.
    
    Scans inputs and outputs using TrustGuard's specialized detection
    for web content, documents, and more.
    
    Example:
        from langchain_openai import ChatOpenAI
        from agent_trust_langchain import TrustGuardCallback
        
        callback = TrustGuardCallback(
            api_key="ta_xxx...",
            block_on_threat=True,
            scan_type="web",  # Optimized for web content
        )
        
        llm = ChatOpenAI(callbacks=[callback])
    """
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        api_url: Optional[str] = None,
        block_on_threat: bool = False,
        min_block_level: ThreatLevel = ThreatLevel.HIGH,
        scan_inputs: bool = True,
        scan_outputs: bool = False,
        scan_type: str = "text",  # "text", "web", "document", "memory"
        on_threat_detected: Optional[callable] = None,
        timeout: float = 30.0,
    ):
        """
        Initialize the callback.
        
        Args:
            api_key: TrustGuard API key
            api_url: Custom API URL
            block_on_threat: Raise exception on threat
            min_block_level: Minimum level for blocking
            scan_inputs: Scan input messages
            scan_outputs: Scan LLM outputs
            scan_type: Type of content scanning to use
            on_threat_detected: Custom callback for threats
            timeout: Request timeout
        """
        super().__init__()
        
        self.block_on_threat = block_on_threat
        self.min_block_level = min_block_level
        self.scan_inputs = scan_inputs
        self.scan_outputs = scan_outputs
        self.scan_type = scan_type
        self.on_threat_detected = on_threat_detected
        
        guard_kwargs = {"timeout": timeout}
        if api_key:
            guard_kwargs["api_key"] = api_key
        if api_url:
            guard_kwargs["api_url"] = api_url
        self.guard = TrustGuard(**guard_kwargs)
        
        # Stats
        self.scanned_count = 0
        self.threat_count = 0
        self.blocked_count = 0
        self.detected_threats: List[Dict[str, Any]] = []
    
    def _get_threat_level_value(self, level: ThreatLevel) -> int:
        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:
        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_content(self, content: str, source: str = "input") -> GuardResult:
        """Scan content using appropriate TrustGuard method."""
        self.scanned_count += 1
        
        if self.scan_type == "web":
            return self.guard.scan_web(content)
        elif self.scan_type == "document":
            return self.guard.scan_document(content)
        elif self.scan_type == "memory":
            return self.guard.scan_memory(content)
        else:
            # Default: use document scanning (good general-purpose detection)
            return self.guard.scan_document(content)
    
    def _handle_scan_result(self, result: GuardResult, source: str, content: str) -> None:
        """Handle scan result, potentially blocking or logging."""
        if result.safe:
            return
        
        self.threat_count += 1
        
        threat_record = {
            "source": source,
            "verdict": result.verdict.value,
            "threat_level": result.threat_level.value,
            "threats": [t.pattern_name for t in result.threats],
            "reasoning": result.reasoning,
            "content_preview": content[:100] + "..." if len(content) > 100 else content,
        }
        self.detected_threats.append(threat_record)
        
        logger.warning(
            f"TrustGuard threat in {source}: {result.threat_level.value} - "
            f"{[t.pattern_name for t in result.threats]}"
        )
        
        if self.on_threat_detected:
            self.on_threat_detected(threat_record)
        
        if self._should_block(result.threat_level):
            self.blocked_count += 1
            raise ThreatInDocumentError(
                f"Content blocked by TrustGuard: {result.reasoning}",
                document_source=source,
                guard_result=result,
            )
    
    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Scan prompts before sending to LLM."""
        if not self.scan_inputs:
            return
        
        for prompt in prompts:
            if prompt and prompt.strip():
                try:
                    result = self._scan_content(prompt, "prompt")
                    self._handle_scan_result(result, "prompt", prompt)
                except ThreatInDocumentError:
                    raise
                except Exception as e:
                    logger.warning(f"Failed to scan prompt: {e}")
    
    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[Any]],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Scan chat messages before sending."""
        if not self.scan_inputs:
            return
        
        for message_list in messages:
            for message in message_list:
                content = getattr(message, "content", None)
                if content and isinstance(content, str) and content.strip():
                    try:
                        result = self._scan_content(content, "message")
                        self._handle_scan_result(result, "message", content)
                    except ThreatInDocumentError:
                        raise
                    except Exception as e:
                        logger.warning(f"Failed to scan message: {e}")
    
    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Optionally scan LLM outputs."""
        if not self.scan_outputs:
            return
        
        try:
            if hasattr(response, 'generations'):
                for gen_list in response.generations:
                    for gen in gen_list:
                        content = getattr(gen, 'text', None) or getattr(getattr(gen, 'message', None), 'content', None)
                        if content and isinstance(content, str):
                            result = self._scan_content(content, "output")
                            self._handle_scan_result(result, "output", content)
        except ThreatInDocumentError:
            raise
        except Exception as e:
            logger.debug(f"Could not scan output: {e}")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get scanning statistics."""
        return {
            "scanned_count": self.scanned_count,
            "threat_count": self.threat_count,
            "blocked_count": self.blocked_count,
            "recent_threats": self.detected_threats[-10:],
        }
