"""
Vaikora Interceptor - Function wrapper for AI agent security.

This module provides the @vaikora.interceptor decorator that wraps function
calls with comprehensive security checks:
1. Data validation (garbage data handling)
2. Policy evaluation
3. Human-in-the-loop approval for sensitive actions

Usage:
    from vaikora import VaikoraClient, interceptor
    
    client = VaikoraClient(api_key="...", agent_id="...")
    
    @interceptor(
        client=client,
        action_type="database.write",
        resource="production",
        require_approval=True,
        validate_inputs=True,
        validate_outputs=True,
    )
    async def write_to_database(data: dict) -> bool:
        # Your code here
        return True
"""
import asyncio
import functools
import hashlib
import json
import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Awaitable, Callable, Optional, ParamSpec, TypeVar, Union

P = ParamSpec("P")
T = TypeVar("T")

logger = logging.getLogger("vaikora.interceptor")


class InterceptorAction(str, Enum):
    """Action to take when policy denies or issues are found."""
    BLOCK = "block"  # Raise exception
    WARN = "warn"  # Log warning but allow
    QUEUE = "queue"  # Queue for approval
    FALLBACK = "fallback"  # Use fallback function


class ApprovalStatus(str, Enum):
    """Status of approval request."""
    PENDING = "pending"
    APPROVED = "approved"
    DENIED = "denied"
    EXPIRED = "expired"
    CANCELLED = "cancelled"


@dataclass
class InterceptorResult:
    """Result from an intercepted function call."""
    success: bool
    action_id: Optional[str] = None
    result: Any = None
    error: Optional[str] = None
    validation_issues: list = field(default_factory=list)
    policy_violation: Optional[str] = None
    approval_required: bool = False
    approval_id: Optional[str] = None
    execution_time_ms: int = 0


@dataclass
class InterceptorConfig:
    """Configuration for the interceptor."""
    action_type: str
    resource: Optional[str] = None
    agent_id: Optional[str] = None
    
    # Validation settings
    validate_inputs: bool = True
    validate_outputs: bool = False
    validation_strict: bool = False
    check_pii: bool = True
    check_anomalies: bool = True
    check_toxicity: bool = True
    auto_clean: bool = False
    
    # Policy settings
    evaluate_policy: bool = True
    on_policy_deny: InterceptorAction = InterceptorAction.BLOCK
    
    # Approval settings
    require_approval: bool = False
    approval_timeout_seconds: int = 3600  # 1 hour
    approval_notify_channels: list = field(default_factory=lambda: ["email"])
    approval_message: Optional[str] = None
    
    # Fallback settings
    fallback_function: Optional[Callable] = None
    
    # Metadata
    metadata: dict = field(default_factory=dict)
    tags: list = field(default_factory=list)


class InterceptorException(Exception):
    """Base exception for interceptor errors."""
    pass


class PolicyDeniedException(InterceptorException):
    """Raised when policy denies the action."""
    def __init__(
        self,
        message: str,
        policy_id: Optional[str] = None,
        action_id: Optional[str] = None,
    ):
        super().__init__(message)
        self.policy_id = policy_id
        self.action_id = action_id


class ValidationException(InterceptorException):
    """Raised when input/output validation fails."""
    def __init__(self, message: str, issues: list):
        super().__init__(message)
        self.issues = issues


class ApprovalRequiredException(InterceptorException):
    """Raised when action requires human approval."""
    def __init__(
        self,
        message: str,
        approval_id: str,
        action_id: Optional[str] = None,
    ):
        super().__init__(message)
        self.approval_id = approval_id
        self.action_id = action_id


class ApprovalDeniedException(InterceptorException):
    """Raised when human approval is denied."""
    def __init__(
        self,
        message: str,
        approval_id: str,
        denied_by: Optional[str] = None,
        reason: Optional[str] = None,
    ):
        super().__init__(message)
        self.approval_id = approval_id
        self.denied_by = denied_by
        self.reason = reason


class ApprovalExpiredException(InterceptorException):
    """Raised when approval request expires."""
    def __init__(self, message: str, approval_id: str):
        super().__init__(message)
        self.approval_id = approval_id


class VaikoraInterceptor:
    """
    Interceptor that wraps function calls with security checks.
    
    Provides:
    - Input/output data validation (garbage data handling)
    - Policy evaluation
    - Human-in-the-loop approval workflow
    - Fallback execution
    - Comprehensive logging and audit trail
    """
    
    def __init__(self, client: "VaikoraClient"):  # type: ignore
        """
        Initialize the interceptor with a Vaikora client.
        
        Args:
            client: VaikoraClient instance for API communication
        """
        self._client = client
        self._approval_cache: dict[str, ApprovalStatus] = {}
    
    def __call__(
        self,
        action_type: str,
        resource: Optional[str] = None,
        agent_id: Optional[str] = None,
        validate_inputs: bool = True,
        validate_outputs: bool = False,
        validation_strict: bool = False,
        check_pii: bool = True,
        check_anomalies: bool = True,
        check_toxicity: bool = True,
        auto_clean: bool = False,
        evaluate_policy: bool = True,
        on_policy_deny: Union[InterceptorAction, str] = InterceptorAction.BLOCK,
        require_approval: bool = False,
        approval_timeout_seconds: int = 3600,
        approval_notify_channels: Optional[list[str]] = None,
        approval_message: Optional[str] = None,
        fallback_function: Optional[Callable] = None,
        metadata: Optional[dict] = None,
        tags: Optional[list[str]] = None,
    ) -> Callable[[Callable[P, T]], Callable[P, Awaitable[T]]]:
        """
        Decorator to intercept function calls with security checks.
        
        Args:
            action_type: Type of action (e.g., "database.write", "file.delete")
            resource: Resource being accessed (e.g., "production", "users")
            agent_id: Agent ID (overrides client default)
            validate_inputs: Enable input validation
            validate_outputs: Enable output validation
            validation_strict: Fail on any validation issue
            check_pii: Check for PII in data
            check_anomalies: Check for anomalies in data
            check_toxicity: Check for toxic content
            auto_clean: Auto-clean data issues
            evaluate_policy: Evaluate against policies
            on_policy_deny: Action when policy denies (block, warn, queue, fallback)
            require_approval: Require human approval
            approval_timeout_seconds: Approval expiration time
            approval_notify_channels: Notification channels for approval
            approval_message: Custom approval request message
            fallback_function: Function to call on denial/failure
            metadata: Additional metadata for logging
            tags: Tags for categorization
            
        Returns:
            Decorated function with security checks
        """
        if isinstance(on_policy_deny, str):
            on_policy_deny = InterceptorAction(on_policy_deny)
        
        config = InterceptorConfig(
            action_type=action_type,
            resource=resource,
            agent_id=agent_id,
            validate_inputs=validate_inputs,
            validate_outputs=validate_outputs,
            validation_strict=validation_strict,
            check_pii=check_pii,
            check_anomalies=check_anomalies,
            check_toxicity=check_toxicity,
            auto_clean=auto_clean,
            evaluate_policy=evaluate_policy,
            on_policy_deny=on_policy_deny,
            require_approval=require_approval,
            approval_timeout_seconds=approval_timeout_seconds,
            approval_notify_channels=approval_notify_channels or ["email"],
            approval_message=approval_message,
            fallback_function=fallback_function,
            metadata=metadata or {},
            tags=tags or [],
        )
        
        def decorator(func: Callable[P, T]) -> Callable[P, Awaitable[T]]:
            @functools.wraps(func)
            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
                return await self._execute_intercepted(func, config, *args, **kwargs)
            
            # Attach config for introspection
            wrapper._interceptor_config = config  # type: ignore
            wrapper._interceptor = self  # type: ignore
            
            return wrapper
        
        return decorator
    
    async def _execute_intercepted(
        self,
        func: Callable[P, T],
        config: InterceptorConfig,
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> T:
        """Execute the intercepted function with all security checks."""
        start_time = time.time()
        action_id: Optional[str] = None
        
        try:
            # 1. Validate inputs
            if config.validate_inputs:
                await self._validate_data(
                    data={"args": args, "kwargs": kwargs},
                    config=config,
                    context="input",
                )
            
            # 2. Determine agent ID
            target_agent_id = config.agent_id or self._client._default_agent_id
            if not target_agent_id:
                raise InterceptorException(
                    "agent_id must be provided either to interceptor or client"
                )
            
            # 3. Submit action for policy evaluation
            if config.evaluate_policy:
                payload = {
                    "args": self._serialize_args(args),
                    "kwargs": self._serialize_kwargs(kwargs),
                    "function": func.__name__,
                    "module": func.__module__,
                }
                
                result = await self._client.actions.submit(
                    agent_id=target_agent_id,
                    action_type=config.action_type,
                    resource=config.resource,
                    payload=payload,
                    metadata={
                        **config.metadata,
                        "tags": config.tags,
                        "require_approval": config.require_approval,
                    },
                )
                
                action_id = str(result.action_id)
                
                # Handle policy denial
                if not result.approved:
                    await self._handle_policy_denial(
                        config=config,
                        denial_reason=result.denial_reason,
                        policy_id=str(result.policy_id) if result.policy_id else None,
                        action_id=action_id,
                        func=func,
                        args=args,
                        kwargs=kwargs,
                    )
            
            # 4. Check if approval is required
            if config.require_approval:
                approval_id = await self._request_approval(
                    config=config,
                    action_id=action_id,
                    func=func,
                    args=args,
                    kwargs=kwargs,
                )
                
                # Wait for approval or handle async
                approval_status = await self._wait_for_approval(
                    approval_id=approval_id,
                    timeout_seconds=config.approval_timeout_seconds,
                )
                
                if approval_status == ApprovalStatus.DENIED:
                    raise ApprovalDeniedException(
                        message="Action was denied by human reviewer",
                        approval_id=approval_id,
                    )
                elif approval_status == ApprovalStatus.EXPIRED:
                    raise ApprovalExpiredException(
                        message="Approval request expired",
                        approval_id=approval_id,
                    )
                elif approval_status != ApprovalStatus.APPROVED:
                    raise ApprovalRequiredException(
                        message="Action requires human approval",
                        approval_id=approval_id,
                        action_id=action_id,
                    )
            
            # 5. Execute the function
            if asyncio.iscoroutinefunction(func):
                output = await func(*args, **kwargs)
            else:
                # Run sync function in executor
                loop = asyncio.get_event_loop()
                output = await loop.run_in_executor(
                    None, lambda: func(*args, **kwargs)
                )
            
            # 6. Validate outputs
            if config.validate_outputs and output is not None:
                await self._validate_data(
                    data=output,
                    config=config,
                    context="output",
                )
            
            # 7. Mark action as completed
            if action_id:
                execution_time_ms = int((time.time() - start_time) * 1000)
                await self._client.actions.complete(
                    action_id=action_id,
                    status="executed",
                    execution_time_ms=execution_time_ms,
                    result_metadata={
                        "output_type": type(output).__name__,
                        "output_hash": self._hash_data(output),
                    },
                )
            
            return output
            
        except (PolicyDeniedException, ApprovalDeniedException, ApprovalExpiredException):
            # Re-raise security exceptions
            raise
            
        except Exception as e:
            # Mark action as failed
            if action_id:
                execution_time_ms = int((time.time() - start_time) * 1000)
                try:
                    await self._client.actions.complete(
                        action_id=action_id,
                        status="failed",
                        execution_time_ms=execution_time_ms,
                        error_message=str(e),
                    )
                except Exception as complete_error:
                    logger.error(f"Failed to mark action as failed: {complete_error}")
            
            raise
    
    async def _validate_data(
        self,
        data: Any,
        config: InterceptorConfig,
        context: str,
    ) -> None:
        """Validate data using the validation API."""
        try:
            result = await self._client.validation.validate(
                data=data,
                check_pii=config.check_pii,
                check_anomalies=config.check_anomalies,
                check_toxicity=config.check_toxicity,
                auto_clean=config.auto_clean,
                strict_mode=config.validation_strict,
            )
            
            if not result.is_valid:
                issues = []
                if result.pii_count > 0:
                    issues.append(f"PII detected ({result.pii_count} instances)")
                if result.anomaly_count > 0:
                    issues.append(f"Anomalies detected ({result.anomaly_count} instances)")
                if result.toxicity_count > 0:
                    issues.append(f"Toxic content detected ({result.toxicity_count} instances)")
                
                if config.validation_strict:
                    raise ValidationException(
                        message=f"Validation failed for {context}: {', '.join(issues)}",
                        issues=issues,
                    )
                else:
                    logger.warning(f"Validation issues in {context}: {', '.join(issues)}")
                    
        except ValidationException:
            raise
        except Exception as e:
            logger.error(f"Validation error: {e}")
            if config.validation_strict:
                raise ValidationException(
                    message=f"Validation service error: {e}",
                    issues=[str(e)],
                )
    
    async def _handle_policy_denial(
        self,
        config: InterceptorConfig,
        denial_reason: Optional[str],
        policy_id: Optional[str],
        action_id: Optional[str],
        func: Callable,
        args: tuple,
        kwargs: dict,
    ) -> None:
        """Handle policy denial based on configured action."""
        message = denial_reason or "Action denied by policy"
        
        if config.on_policy_deny == InterceptorAction.BLOCK:
            raise PolicyDeniedException(
                message=message,
                policy_id=policy_id,
                action_id=action_id,
            )
            
        elif config.on_policy_deny == InterceptorAction.WARN:
            logger.warning(f"Policy warning (allowed): {message}")
            
        elif config.on_policy_deny == InterceptorAction.QUEUE:
            # Request approval for denied action
            approval_id = await self._request_approval(
                config=config,
                action_id=action_id,
                func=func,
                args=args,
                kwargs=kwargs,
                reason=f"Policy denied: {message}",
            )
            raise ApprovalRequiredException(
                message=f"Action queued for approval: {message}",
                approval_id=approval_id,
                action_id=action_id,
            )
            
        elif config.on_policy_deny == InterceptorAction.FALLBACK:
            if config.fallback_function:
                logger.info(f"Using fallback function due to policy denial: {message}")
                # Will be handled by caller
                raise PolicyDeniedException(
                    message=message,
                    policy_id=policy_id,
                    action_id=action_id,
                )
            else:
                raise PolicyDeniedException(
                    message=f"{message} (no fallback configured)",
                    policy_id=policy_id,
                    action_id=action_id,
                )
    
    async def _request_approval(
        self,
        config: InterceptorConfig,
        action_id: Optional[str],
        func: Callable,
        args: tuple,
        kwargs: dict,
        reason: Optional[str] = None,
    ) -> str:
        """Request human approval for an action."""
        # Build approval request message
        message = config.approval_message or (
            f"Action '{config.action_type}' on resource '{config.resource}' "
            f"requires approval"
        )
        if reason:
            message = f"{message}. Reason: {reason}"
        
        # Determine agent ID
        target_agent_id = config.agent_id or self._client._default_agent_id
        
        # Make approval request
        response = await self._client._request(
            "POST",
            "/approvals/request",
            json={
                "agent_id": target_agent_id,
                "action_type": config.action_type,
                "resource": config.resource,
                "action_log_id": action_id,
                "action_details": {
                    "function": func.__name__,
                    "module": func.__module__,
                    "args_summary": str(args)[:200],
                    "kwargs_keys": list(kwargs.keys()),
                },
                "notification_channels": config.approval_notify_channels,
                "message": message,
                "expires_in_hours": max(1, int(config.approval_timeout_seconds / 3600)),
            },
        )
        
        approval_id = response.get("approval_id") or response.get("id")
        logger.info(f"Approval requested: {approval_id}")
        
        return str(approval_id)
    
    async def _wait_for_approval(
        self,
        approval_id: str,
        timeout_seconds: int,
        poll_interval: float = 2.0,
    ) -> ApprovalStatus:
        """Wait for approval status to change from pending."""
        start_time = time.time()
        
        while time.time() - start_time < timeout_seconds:
            # Check cache first
            if approval_id in self._approval_cache:
                status = self._approval_cache[approval_id]
                if status != ApprovalStatus.PENDING:
                    return status
            
            # Poll API
            try:
                response = await self._client._request(
                    "GET",
                    f"/approvals/status/{approval_id}",
                )
                status_str = response.get("status", "pending")
                status = ApprovalStatus(status_str)
                
                self._approval_cache[approval_id] = status
                
                if status != ApprovalStatus.PENDING:
                    return status
                    
            except Exception as e:
                logger.warning(f"Error polling approval status: {e}")
            
            await asyncio.sleep(poll_interval)
        
        return ApprovalStatus.EXPIRED
    
    async def check_approval_status(self, approval_id: str) -> ApprovalStatus:
        """Check the current status of an approval request."""
        try:
            response = await self._client._request(
                "GET",
                f"/approvals/status/{approval_id}",
            )
            return ApprovalStatus(response.get("status", "pending"))
        except Exception as e:
            logger.error(f"Error checking approval status: {e}")
            raise
    
    async def cancel_approval(self, approval_id: str) -> bool:
        """Cancel a pending approval request."""
        try:
            await self._client._request(
                "POST",
                f"/approvals/resolve/{approval_id}",
                json={"resolution": "cancelled", "reason": "Cancelled by SDK"},
            )
            self._approval_cache[approval_id] = ApprovalStatus.CANCELLED
            return True
        except Exception as e:
            logger.error(f"Error cancelling approval: {e}")
            return False
    
    def _serialize_args(self, args: tuple) -> list:
        """Serialize function arguments to JSON-safe format."""
        result = []
        for arg in args:
            try:
                # Try JSON serialization
                json.dumps(arg)
                result.append(arg)
            except (TypeError, ValueError):
                # Fall back to string representation
                result.append(str(arg))
        return result
    
    def _serialize_kwargs(self, kwargs: dict) -> dict:
        """Serialize function keyword arguments to JSON-safe format."""
        result = {}
        for key, value in kwargs.items():
            try:
                # Try JSON serialization
                json.dumps(value)
                result[key] = value
            except (TypeError, ValueError):
                # Fall back to string representation
                result[key] = str(value)
        return result
    
    def _hash_data(self, data: Any) -> str:
        """Create a hash of data for audit purposes."""
        try:
            serialized = json.dumps(data, sort_keys=True, default=str)
            return hashlib.sha256(serialized.encode()).hexdigest()[:16]
        except Exception:
            return hashlib.sha256(str(data).encode()).hexdigest()[:16]


# Convenience function for creating interceptor
def create_interceptor(client: "VaikoraClient") -> VaikoraInterceptor:  # type: ignore
    """Create an interceptor instance from a Vaikora client."""
    return VaikoraInterceptor(client)


# Type alias for the decorator
interceptor = VaikoraInterceptor

__all__ = [
    "VaikoraInterceptor",
    "InterceptorConfig",
    "InterceptorResult",
    "InterceptorAction",
    "ApprovalStatus",
    "InterceptorException",
    "PolicyDeniedException",
    "ValidationException",
    "ApprovalRequiredException",
    "ApprovalDeniedException",
    "ApprovalExpiredException",
    "create_interceptor",
    "interceptor",
]
