"""Firewall wrapper for easy integration"""

from typing import Dict, Any, List, Optional, Callable, Awaitable
from .client import AegisClient, AegisAsyncClient
from .types import (
    ToolCall,
    ToolCallResult,
    ToolCallDeclaration,
    Span,
    SpanSource,
    TrustLevel,
    SensitivityLevel,
)
from .errors import CapabilityBlockedError, ApprovalRequiredError


class AegisFirewall:
    """Synchronous firewall wrapper"""
    
    def __init__(self, client: AegisClient):
        self.client = client
        self.session_spans: Dict[str, Span] = {}
    
    def add_content(
        self,
        content: str,
        source: SpanSource,
        trust_level: TrustLevel,
        sensitivity: SensitivityLevel = SensitivityLevel.PUBLIC,
    ) -> Span:
        """Add content to session"""
        span = self.client.create_span({
            "content": content,
            "source": source.value,
            "trust_level": trust_level.value,
            "sensitivity": sensitivity.value,
        })
        self.session_spans[span.id] = span
        return span
    
    def add_user_message(self, content: str) -> Span:
        """Add user message"""
        return self.add_content(
            content,
            SpanSource.USER,
            TrustLevel.UNTRUSTED_USER,
            SensitivityLevel.PUBLIC,
        )
    
    def add_document(
        self,
        content: str,
        sensitivity: SensitivityLevel = SensitivityLevel.INTERNAL,
    ) -> Span:
        """Add document from RAG"""
        return self.add_content(
            content,
            SpanSource.DOCUMENT,
            TrustLevel.UNTRUSTED_DOCUMENT,
            sensitivity,
        )
    
    def add_system_prompt(self, content: str) -> Span:
        """Add system prompt"""
        return self.add_content(
            content,
            SpanSource.SYSTEM,
            TrustLevel.TRUSTED_SYSTEM,
            SensitivityLevel.PUBLIC,
        )
    
    def execute_tool_call(
        self,
        capability: str,
        args: Dict[str, Any],
        source_span_ids: List[str],
    ) -> ToolCallResult:
        """Validate and execute tool call"""
        sources = []
        for span_id in source_span_ids:
            span = self.session_spans.get(span_id)
            if not span:
                raise ValueError(f"Span {span_id} not found in session")
            sources.append({
                "spanId": span_id,
                "trust": span.trust_level.value,
                "sensitivity": span.sensitivity.value,
            })
        
        tool_call = ToolCall(
            capability=capability,
            arguments=args,
            declaration=ToolCallDeclaration(sources=sources),
        )
        
        return self.client.validate_tool_call(tool_call)
    
    def safe_execute_tool_call(
        self,
        capability: str,
        args: Dict[str, Any],
        source_span_ids: List[str],
        on_approval_required: Optional[Callable[[str], None]] = None,
    ) -> Dict[str, Any]:
        """Safe tool call execution with error handling"""
        try:
            result = self.execute_tool_call(capability, args, source_span_ids)
            
            if result.requires_approval and result.approval_id and on_approval_required:
                on_approval_required(result.approval_id)
            
            return {
                "success": result.allowed,
                "result": result,
            }
        except (CapabilityBlockedError, ApprovalRequiredError) as e:
            if isinstance(e, ApprovalRequiredError) and on_approval_required:
                on_approval_required(e.approval_id)
            
            return {
                "success": False,
                "error": e,
            }
    
    def get_session_spans(self) -> List[Span]:
        """Get all spans in session"""
        return list(self.session_spans.values())
    
    def clear_session(self) -> None:
        """Clear session"""
        self.session_spans.clear()


class AegisAsyncFirewall:
    """Asynchronous firewall wrapper"""
    
    def __init__(self, client: AegisAsyncClient):
        self.client = client
        self.session_spans: Dict[str, Span] = {}
    
    async def add_content(
        self,
        content: str,
        source: SpanSource,
        trust_level: TrustLevel,
        sensitivity: SensitivityLevel = SensitivityLevel.PUBLIC,
    ) -> Span:
        """Add content to session"""
        span = await self.client.create_span({
            "content": content,
            "source": source.value,
            "trust_level": trust_level.value,
            "sensitivity": sensitivity.value,
        })
        self.session_spans[span.id] = span
        return span
    
    async def add_user_message(self, content: str) -> Span:
        """Add user message"""
        return await self.add_content(
            content,
            SpanSource.USER,
            TrustLevel.UNTRUSTED_USER,
            SensitivityLevel.PUBLIC,
        )
    
    async def add_document(
        self,
        content: str,
        sensitivity: SensitivityLevel = SensitivityLevel.INTERNAL,
    ) -> Span:
        """Add document from RAG"""
        return await self.add_content(
            content,
            SpanSource.DOCUMENT,
            TrustLevel.UNTRUSTED_DOCUMENT,
            sensitivity,
        )
    
    async def add_system_prompt(self, content: str) -> Span:
        """Add system prompt"""
        return await self.add_content(
            content,
            SpanSource.SYSTEM,
            TrustLevel.TRUSTED_SYSTEM,
            SensitivityLevel.PUBLIC,
        )
    
    async def execute_tool_call(
        self,
        capability: str,
        args: Dict[str, Any],
        source_span_ids: List[str],
    ) -> ToolCallResult:
        """Validate and execute tool call"""
        sources = []
        for span_id in source_span_ids:
            span = self.session_spans.get(span_id)
            if not span:
                raise ValueError(f"Span {span_id} not found in session")
            sources.append({
                "spanId": span_id,
                "trust": span.trust_level.value,
                "sensitivity": span.sensitivity.value,
            })
        
        tool_call = ToolCall(
            capability=capability,
            arguments=args,
            declaration=ToolCallDeclaration(sources=sources),
        )
        
        return await self.client.validate_tool_call(tool_call)
    
    async def safe_execute_tool_call(
        self,
        capability: str,
        args: Dict[str, Any],
        source_span_ids: List[str],
        on_approval_required: Optional[Callable[[str], Awaitable[None]]] = None,
    ) -> Dict[str, Any]:
        """Safe tool call execution with error handling"""
        try:
            result = await self.execute_tool_call(capability, args, source_span_ids)
            
            if result.requires_approval and result.approval_id and on_approval_required:
                await on_approval_required(result.approval_id)
            
            return {
                "success": result.allowed,
                "result": result,
            }
        except (CapabilityBlockedError, ApprovalRequiredError) as e:
            if isinstance(e, ApprovalRequiredError) and on_approval_required:
                await on_approval_required(e.approval_id)
            
            return {
                "success": False,
                "error": e,
            }
    
    def get_session_spans(self) -> List[Span]:
        """Get all spans in session"""
        return list(self.session_spans.values())
    
    def clear_session(self) -> None:
        """Clear session"""
        self.session_spans.clear()
