"""
Emergency Controls - Agent Freeze/Resume Capabilities.

Provides:
- Immediate agent freeze
- Controlled resume
- Emergency killswitch
- Circuit breaker patterns
"""
import asyncio
import time
import logging
from dataclasses import dataclass, field
from typing import Optional, Any, Callable
from datetime import datetime, timezone
from enum import Enum
from functools import wraps

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


class FreezeReason(str, Enum):
    """Reasons for freezing an agent."""
    MANUAL = "manual"
    POLICY_VIOLATION = "policy_violation"
    ANOMALY_DETECTED = "anomaly_detected"
    RATE_LIMIT = "rate_limit"
    SECURITY_THREAT = "security_threat"
    SYSTEM_OVERLOAD = "system_overload"
    CIRCUIT_BREAKER = "circuit_breaker"
    MAINTENANCE = "maintenance"


class FreezeScope(str, Enum):
    """Scope of freeze."""
    ALL_ACTIONS = "all_actions"
    WRITE_ACTIONS = "write_actions"
    SPECIFIC_RESOURCE = "specific_resource"
    SPECIFIC_ACTION_TYPE = "specific_action_type"


@dataclass
class FreezeState:
    """State of an agent freeze."""
    is_frozen: bool
    reason: Optional[FreezeReason] = None
    scope: FreezeScope = FreezeScope.ALL_ACTIONS
    frozen_at: Optional[datetime] = None
    frozen_by: Optional[str] = None
    freeze_until: Optional[datetime] = None
    affected_resources: list[str] = field(default_factory=list)
    affected_action_types: list[str] = field(default_factory=list)
    message: Optional[str] = None
    metadata: dict = field(default_factory=dict)
    
    def is_expired(self) -> bool:
        """Check if freeze has expired."""
        if not self.freeze_until:
            return False
        return datetime.now(timezone.utc) > self.freeze_until
    
    def to_dict(self) -> dict:
        return {
            "is_frozen": self.is_frozen,
            "reason": self.reason.value if self.reason else None,
            "scope": self.scope.value,
            "frozen_at": self.frozen_at.isoformat() if self.frozen_at else None,
            "frozen_by": self.frozen_by,
            "freeze_until": self.freeze_until.isoformat() if self.freeze_until else None,
            "affected_resources": self.affected_resources,
            "affected_action_types": self.affected_action_types,
            "message": self.message,
            "metadata": self.metadata,
        }


@dataclass
class CircuitBreakerConfig:
    """Configuration for circuit breaker."""
    failure_threshold: int = 5  # Failures before opening
    success_threshold: int = 3  # Successes before closing
    timeout_seconds: float = 60.0  # Time before attempting reset
    half_open_max_calls: int = 3  # Max calls in half-open state


class CircuitState(str, Enum):
    """Circuit breaker states."""
    CLOSED = "closed"  # Normal operation
    OPEN = "open"  # Blocking all calls
    HALF_OPEN = "half_open"  # Testing if service recovered


@dataclass
class CircuitBreaker:
    """Circuit breaker for resilient operations."""
    name: str
    config: CircuitBreakerConfig
    state: CircuitState = CircuitState.CLOSED
    failure_count: int = 0
    success_count: int = 0
    last_failure_time: Optional[float] = None
    half_open_calls: int = 0
    
    def is_open(self) -> bool:
        """Check if circuit is open (blocking)."""
        if self.state == CircuitState.OPEN:
            # Check if timeout has passed
            if self.last_failure_time:
                elapsed = time.time() - self.last_failure_time
                if elapsed >= self.config.timeout_seconds:
                    self.state = CircuitState.HALF_OPEN
                    self.half_open_calls = 0
                    return False
            return True
        return False
    
    def record_success(self):
        """Record a successful call."""
        if self.state == CircuitState.HALF_OPEN:
            self.success_count += 1
            if self.success_count >= self.config.success_threshold:
                self.state = CircuitState.CLOSED
                self.failure_count = 0
                self.success_count = 0
        elif self.state == CircuitState.CLOSED:
            self.failure_count = max(0, self.failure_count - 1)
    
    def record_failure(self):
        """Record a failed call."""
        self.failure_count += 1
        self.last_failure_time = time.time()
        
        if self.state == CircuitState.HALF_OPEN:
            self.state = CircuitState.OPEN
            self.success_count = 0
        elif self.failure_count >= self.config.failure_threshold:
            self.state = CircuitState.OPEN
            self.success_count = 0
    
    def allow_request(self) -> bool:
        """Check if request should be allowed."""
        if self.is_open():
            return False
        
        if self.state == CircuitState.HALF_OPEN:
            if self.half_open_calls >= self.config.half_open_max_calls:
                return False
            self.half_open_calls += 1
        
        return True


class EmergencyController:
    """
    Emergency control system for AI agents.
    
    Provides:
    - Immediate freeze capabilities
    - Controlled resume with verification
    - Circuit breaker for external services
    - Rate limiting
    """
    
    def __init__(self, agent_id: str):
        self.agent_id = agent_id
        self._freeze_state = FreezeState(is_frozen=False)
        self._circuit_breakers: dict[str, CircuitBreaker] = {}
        self._freeze_callbacks: list[Callable] = []
        self._resume_callbacks: list[Callable] = []
        self._action_history: list[dict] = []
        self._rate_limits: dict[str, list[float]] = {}
    
    @property
    def is_frozen(self) -> bool:
        """Check if agent is currently frozen."""
        if self._freeze_state.is_expired():
            self._auto_unfreeze()
        return self._freeze_state.is_frozen
    
    @property
    def freeze_state(self) -> FreezeState:
        """Get current freeze state."""
        return self._freeze_state
    
    def freeze(
        self,
        reason: FreezeReason,
        scope: FreezeScope = FreezeScope.ALL_ACTIONS,
        duration_seconds: Optional[float] = None,
        frozen_by: str = "system",
        message: Optional[str] = None,
        affected_resources: Optional[list[str]] = None,
        affected_action_types: Optional[list[str]] = None,
        metadata: Optional[dict] = None,
    ) -> FreezeState:
        """
        Freeze the agent immediately.
        
        Args:
            reason: Why the agent is being frozen
            scope: What actions are affected
            duration_seconds: Auto-unfreeze after this time
            frozen_by: Who/what initiated the freeze
            message: Human-readable message
            affected_resources: Resources affected (for targeted freeze)
            affected_action_types: Action types affected
            metadata: Additional context
            
        Returns:
            FreezeState
        """
        now = datetime.now(timezone.utc)
        freeze_until = None
        if duration_seconds:
            from datetime import timedelta
            freeze_until = now + timedelta(seconds=duration_seconds)
        
        self._freeze_state = FreezeState(
            is_frozen=True,
            reason=reason,
            scope=scope,
            frozen_at=now,
            frozen_by=frozen_by,
            freeze_until=freeze_until,
            affected_resources=affected_resources or [],
            affected_action_types=affected_action_types or [],
            message=message,
            metadata=metadata or {},
        )
        
        logger.warning(
            f"Agent {self.agent_id} FROZEN: reason={reason.value}, scope={scope.value}, by={frozen_by}"
        )
        
        # Execute callbacks
        for callback in self._freeze_callbacks:
            try:
                callback(self._freeze_state)
            except Exception as e:
                logger.error(f"Freeze callback failed: {e}")
        
        return self._freeze_state
    
    def resume(
        self,
        resumed_by: str = "system",
        verification_passed: bool = True,
        notes: Optional[str] = None,
    ) -> FreezeState:
        """
        Resume agent operations.
        
        Args:
            resumed_by: Who/what is resuming
            verification_passed: Whether verification checks passed
            notes: Notes about the resume
            
        Returns:
            Updated FreezeState (unfrozen)
        """
        if not self._freeze_state.is_frozen:
            return self._freeze_state
        
        if not verification_passed:
            logger.warning(f"Resume denied for agent {self.agent_id}: verification failed")
            return self._freeze_state
        
        old_state = self._freeze_state
        self._freeze_state = FreezeState(is_frozen=False)
        
        logger.info(
            f"Agent {self.agent_id} RESUMED: by={resumed_by}, "
            f"was_frozen_for={old_state.reason.value if old_state.reason else 'unknown'}"
        )
        
        # Execute callbacks
        for callback in self._resume_callbacks:
            try:
                callback(old_state, notes)
            except Exception as e:
                logger.error(f"Resume callback failed: {e}")
        
        return self._freeze_state
    
    def _auto_unfreeze(self):
        """Auto-unfreeze when time expires."""
        logger.info(f"Agent {self.agent_id} auto-unfreezing (time expired)")
        self._freeze_state = FreezeState(is_frozen=False)
    
    def is_action_blocked(
        self,
        action_type: str,
        resource: Optional[str] = None,
    ) -> tuple[bool, Optional[str]]:
        """
        Check if a specific action is blocked.
        
        Args:
            action_type: Type of action
            resource: Target resource
            
        Returns:
            Tuple of (is_blocked, reason_message)
        """
        if not self.is_frozen:
            return False, None
        
        state = self._freeze_state
        
        # Check scope
        if state.scope == FreezeScope.ALL_ACTIONS:
            return True, state.message or f"Agent frozen: {state.reason.value}"
        
        if state.scope == FreezeScope.WRITE_ACTIONS:
            write_actions = {"create", "update", "delete", "write", "modify", "insert"}
            if any(w in action_type.lower() for w in write_actions):
                return True, state.message or "Write actions frozen"
        
        if state.scope == FreezeScope.SPECIFIC_RESOURCE:
            if resource and resource in state.affected_resources:
                return True, state.message or f"Resource {resource} is frozen"
        
        if state.scope == FreezeScope.SPECIFIC_ACTION_TYPE:
            if action_type in state.affected_action_types:
                return True, state.message or f"Action type {action_type} is frozen"
        
        return False, None
    
    def on_freeze(self, callback: Callable[[FreezeState], None]):
        """Register a callback for freeze events."""
        self._freeze_callbacks.append(callback)
    
    def on_resume(self, callback: Callable[[FreezeState, Optional[str]], None]):
        """Register a callback for resume events."""
        self._resume_callbacks.append(callback)
    
    # Circuit Breaker Methods
    
    def get_circuit_breaker(
        self,
        name: str,
        config: Optional[CircuitBreakerConfig] = None,
    ) -> CircuitBreaker:
        """Get or create a circuit breaker."""
        if name not in self._circuit_breakers:
            self._circuit_breakers[name] = CircuitBreaker(
                name=name,
                config=config or CircuitBreakerConfig(),
            )
        return self._circuit_breakers[name]
    
    def circuit_breaker(self, name: str):
        """
        Decorator for circuit breaker protection.
        
        Usage:
            @controller.circuit_breaker("external_api")
            async def call_external_api():
                ...
        """
        def decorator(func: Callable) -> Callable:
            cb = self.get_circuit_breaker(name)
            
            @wraps(func)
            async def async_wrapper(*args, **kwargs):
                if not cb.allow_request():
                    raise CircuitBreakerOpenError(f"Circuit breaker '{name}' is open")
                
                try:
                    result = await func(*args, **kwargs)
                    cb.record_success()
                    return result
                except Exception as e:
                    cb.record_failure()
                    raise
            
            @wraps(func)
            def sync_wrapper(*args, **kwargs):
                if not cb.allow_request():
                    raise CircuitBreakerOpenError(f"Circuit breaker '{name}' is open")
                
                try:
                    result = func(*args, **kwargs)
                    cb.record_success()
                    return result
                except Exception as e:
                    cb.record_failure()
                    raise
            
            if asyncio.iscoroutinefunction(func):
                return async_wrapper
            return sync_wrapper
        
        return decorator
    
    # Rate Limiting
    
    def check_rate_limit(
        self,
        key: str,
        max_requests: int,
        window_seconds: float,
    ) -> tuple[bool, int]:
        """
        Check if rate limit is exceeded.
        
        Args:
            key: Rate limit key
            max_requests: Max requests in window
            window_seconds: Time window
            
        Returns:
            Tuple of (allowed, remaining_requests)
        """
        now = time.time()
        
        if key not in self._rate_limits:
            self._rate_limits[key] = []
        
        # Clean old entries
        cutoff = now - window_seconds
        self._rate_limits[key] = [t for t in self._rate_limits[key] if t > cutoff]
        
        current_count = len(self._rate_limits[key])
        
        if current_count >= max_requests:
            return False, 0
        
        # Record this request
        self._rate_limits[key].append(now)
        return True, max_requests - current_count - 1


class CircuitBreakerOpenError(Exception):
    """Raised when circuit breaker is open."""
    pass


class AgentFrozenError(Exception):
    """Raised when attempting an action while agent is frozen."""
    def __init__(self, message: str, freeze_state: FreezeState):
        super().__init__(message)
        self.freeze_state = freeze_state


def frozen_guard(controller: EmergencyController, action_type: str = "", resource: str = ""):
    """
    Decorator to guard functions against frozen state.
    
    Usage:
        @frozen_guard(controller, action_type="database_write")
        async def write_to_db():
            ...
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            blocked, message = controller.is_action_blocked(action_type, resource)
            if blocked:
                raise AgentFrozenError(message or "Agent is frozen", controller.freeze_state)
            return await func(*args, **kwargs)
        
        @wraps(func)
        def sync_wrapper(*args, **kwargs):
            blocked, message = controller.is_action_blocked(action_type, resource)
            if blocked:
                raise AgentFrozenError(message or "Agent is frozen", controller.freeze_state)
            return func(*args, **kwargs)
        
        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        return sync_wrapper
    
    return decorator
