# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Guardrails for Familiar

Provides runtime safety and cost controls:
- Token budget limits per request
- Cost caps (USD) with circuit breaker
- Iteration limits for agent loops
- Content safety filters
- Rate limiting per user/channel

Usage:
    from familiar.core.guardrails import (
        Guardrails, GuardrailConfig, GuardrailViolation,
        check_guardrails, get_default_guardrails
    )

    guardrails = Guardrails(GuardrailConfig(
        max_tokens_per_request=100_000,
        max_cost_per_request=1.00,
        max_iterations=15
    ))

    # In agent loop:
    for i in range(max_iterations):
        guardrails.check_iteration(i, trace)  # Raises if exceeded
        response = provider.chat(...)
        guardrails.track_usage(trace)
"""

import logging
import re
import threading
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set

if TYPE_CHECKING:
    from familiar.core.tool_registry import ToolRegistry
from collections import defaultdict
from datetime import datetime, timedelta
from enum import Enum

# Import pattern safety for ReDoS-resistant matching
from .pattern_safety import safe_search

logger = logging.getLogger(__name__)


# ============================================================
# CONFIGURATION
# ============================================================


class ViolationType(str, Enum):
    """Types of guardrail violations."""

    TOKEN_LIMIT = "token_limit"
    COST_LIMIT = "cost_limit"
    ITERATION_LIMIT = "iteration_limit"
    RATE_LIMIT = "rate_limit"
    CONTENT_SAFETY = "content_safety"
    TOOL_BLOCKED = "tool_blocked"
    TIMEOUT = "timeout"
    PHI_CLOUD_BLOCK = "phi_cloud_block"


class ViolationAction(str, Enum):
    """What to do when a guardrail is violated."""

    BLOCK = "block"  # Raise exception, halt execution
    WARN = "warn"  # Log warning, continue
    DEGRADE = "degrade"  # Switch to cheaper model/mode
    THROTTLE = "throttle"  # Slow down execution


@dataclass
class GuardrailConfig:
    """Configuration for guardrail limits."""

    # Token limits
    max_tokens_per_request: int = 100_000  # Total tokens in one chat()
    max_tokens_per_turn: int = 16_000  # Tokens in single LLM call
    warn_tokens_threshold: float = 0.8  # Warn at 80% of limit

    # Cost limits
    max_cost_per_request: float = 1.00  # USD per chat()
    max_cost_per_hour: float = 10.00  # USD per hour per user
    max_cost_per_day: float = 50.00  # USD per day per user

    # Iteration limits
    max_iterations: int = 15  # Max tool-use loops
    max_tool_calls_per_iteration: int = 5  # Max parallel tool calls

    # Rate limits
    max_requests_per_minute: int = 20  # Per user
    max_requests_per_hour: int = 200  # Per user

    # Timeout
    request_timeout_seconds: float = 300  # 5 minutes
    tool_timeout_seconds: float = 60  # Per tool call

    # Tool restrictions
    blocked_tools: List[str] = field(default_factory=list)
    require_approval_tools: List[str] = field(default_factory=list)

    # Content safety
    enable_content_filter: bool = True

    # PII Detection (v1.4.0)
    enable_pii_detection: bool = False
    pii_config: Optional["PIIConfig"] = None  # Uses default if None

    # Actions
    token_violation_action: ViolationAction = ViolationAction.BLOCK
    cost_violation_action: ViolationAction = ViolationAction.BLOCK
    rate_violation_action: ViolationAction = ViolationAction.THROTTLE


class GuardrailViolation(Exception):
    """Raised when a guardrail is violated."""

    def __init__(
        self,
        violation_type: ViolationType,
        message: str,
        current_value: Any = None,
        limit_value: Any = None,
        action: ViolationAction = ViolationAction.BLOCK,
        details: Optional[Dict] = None,
    ):
        super().__init__(message)
        self.violation_type = violation_type
        self.current_value = current_value
        self.limit_value = limit_value
        self.action = action
        self.details = details or {}

    def __str__(self):
        return f"[{self.violation_type.value}] {super().__str__()}"


# ============================================================
# USAGE TRACKING
# ============================================================


@dataclass
class UsageTracker:
    """Tracks usage within a single request."""

    total_tokens: int = 0
    total_cost: float = 0.0
    iterations: int = 0
    tool_calls: int = 0
    start_time: float = field(default_factory=time.time)

    # Per-iteration tracking
    tokens_by_iteration: List[int] = field(default_factory=list)
    costs_by_iteration: List[float] = field(default_factory=list)

    def add_usage(self, input_tokens: int, output_tokens: int, cost: float):
        """Record token usage from an LLM call."""
        tokens = input_tokens + output_tokens
        self.total_tokens += tokens
        self.total_cost += cost

        if self.tokens_by_iteration:
            self.tokens_by_iteration[-1] += tokens
            self.costs_by_iteration[-1] += cost

    def start_iteration(self):
        """Start tracking a new iteration."""
        self.iterations += 1
        self.tokens_by_iteration.append(0)
        self.costs_by_iteration.append(0.0)

    def add_tool_call(self):
        """Record a tool call."""
        self.tool_calls += 1

    @property
    def elapsed_seconds(self) -> float:
        return time.time() - self.start_time


@dataclass
class UserUsageHistory:
    """Tracks usage history for rate limiting."""

    # Request timestamps
    request_times: List[float] = field(default_factory=list)

    # Hourly/daily cost accumulation
    hourly_costs: Dict[str, float] = field(default_factory=dict)  # hour_key -> cost
    daily_costs: Dict[str, float] = field(default_factory=dict)  # date_key -> cost

    def record_request(self, cost: float = 0.0):
        """Record a new request."""
        now = time.time()
        self.request_times.append(now)

        # Clean old entries
        cutoff = now - 3600  # 1 hour
        self.request_times = [t for t in self.request_times if t > cutoff]

        # Track costs
        hour_key = datetime.now().strftime("%Y%m%d%H")
        date_key = datetime.now().strftime("%Y%m%d")

        self.hourly_costs[hour_key] = self.hourly_costs.get(hour_key, 0) + cost
        self.daily_costs[date_key] = self.daily_costs.get(date_key, 0) + cost

        # Clean old cost entries
        self._cleanup_old_costs()

    def _cleanup_old_costs(self):
        """Remove cost entries older than 48 hours."""
        cutoff_hour = (datetime.now() - timedelta(hours=48)).strftime("%Y%m%d%H")
        cutoff_date = (datetime.now() - timedelta(days=7)).strftime("%Y%m%d")

        self.hourly_costs = {k: v for k, v in self.hourly_costs.items() if k >= cutoff_hour}
        self.daily_costs = {k: v for k, v in self.daily_costs.items() if k >= cutoff_date}

    def requests_last_minute(self) -> int:
        """Count requests in last 60 seconds."""
        cutoff = time.time() - 60
        return sum(1 for t in self.request_times if t > cutoff)

    def requests_last_hour(self) -> int:
        """Count requests in last hour."""
        return len(self.request_times)

    def cost_last_hour(self) -> float:
        """Get cost in current hour."""
        hour_key = datetime.now().strftime("%Y%m%d%H")
        return self.hourly_costs.get(hour_key, 0)

    def cost_today(self) -> float:
        """Get cost today."""
        date_key = datetime.now().strftime("%Y%m%d")
        return self.daily_costs.get(date_key, 0)


# ============================================================
# CONTENT SAFETY
# ============================================================


class ContentFilter:
    """
    Filters content for safety concerns.

    This is a basic implementation - production systems should use
    dedicated moderation APIs (OpenAI Moderation, Perspective API, etc.)
    """

    # Patterns that should be blocked
    BLOCK_PATTERNS = [
        # Violence
        r"\b(kill|murder|assassinate)\s+(him|her|them|yourself|people)\b",
        # Weapons/explosives
        r"\b(how to|instructions for|recipe for)\s+(make|build|create)\s+(bomb|explosive|weapon|gun)\b",
        # Self-harm (these should trigger support, not block)
        # r'\b(how to|ways to)\s+(kill|hurt)\s+myself\b',
    ]

    # Patterns that should trigger warning
    WARN_PATTERNS = [
        r"\b(hack|crack|bypass)\s+(password|security|authentication)\b",
        r"\b(pirate|torrent|crack)\s+(software|game|movie)\b",
    ]

    def __init__(self):
        # Patterns are used directly with safe_search (no pre-compilation needed)
        pass

    def check(self, text: str) -> tuple:
        """
        Check text for safety concerns.

        Uses ReDoS-resistant pattern matching via safe_search.

        Returns:
            (is_safe: bool, reason: Optional[str], severity: str)
        """
        # Check block patterns using safe_search
        for pattern in self.BLOCK_PATTERNS:
            if safe_search(pattern, text):
                return (False, "Content blocked: matched safety filter", "block")

        # Check warn patterns using safe_search
        for pattern in self.WARN_PATTERNS:
            if safe_search(pattern, text):
                return (True, "Content warning: potentially sensitive", "warn")

        return (True, None, "ok")


# ============================================================
# PII DETECTION
# ============================================================


class PIIAction(str, Enum):
    """Actions to take when PII is detected."""

    BLOCK = "block"  # Block the request entirely
    REDACT = "redact"  # Replace PII with placeholders
    WARN = "warn"  # Log warning but allow
    HASH = "hash"  # Replace with hashed values (reversible with key)


class PIIEntityType(str, Enum):
    """Types of PII that can be detected."""

    # Identifiers
    SSN = "SSN"  # US Social Security Number
    CREDIT_CARD = "CREDIT_CARD"  # Credit card numbers
    IBAN = "IBAN"  # International Bank Account Number

    # Personal info
    PERSON = "PERSON"  # Person names
    EMAIL = "EMAIL"  # Email addresses
    PHONE = "PHONE"  # Phone numbers

    # Location
    ADDRESS = "ADDRESS"  # Physical addresses
    IP_ADDRESS = "IP_ADDRESS"  # IP addresses

    # Healthcare
    MEDICAL_LICENSE = "MEDICAL_LICENSE"

    # Government IDs
    PASSPORT = "PASSPORT"
    DRIVERS_LICENSE = "DRIVERS_LICENSE"

    # Financial
    BANK_ACCOUNT = "BANK_ACCOUNT"

    # Dates
    DATE_OF_BIRTH = "DATE_OF_BIRTH"

    # Custom
    CUSTOM = "CUSTOM"


@dataclass
class PIIFinding:
    """A detected PII entity."""

    entity_type: str
    text: str
    start: int
    end: int
    confidence: float

    def to_dict(self) -> dict:
        return {
            "entity_type": self.entity_type,
            "text": self.text,
            "start": self.start,
            "end": self.end,
            "confidence": self.confidence,
        }


@dataclass
class PIIConfig:
    """Configuration for PII detection."""

    # Which entity types to detect
    enabled_entities: Set[str] = field(
        default_factory=lambda: {
            PIIEntityType.SSN.value,
            PIIEntityType.CREDIT_CARD.value,
            PIIEntityType.EMAIL.value,
            PIIEntityType.PHONE.value,
            PIIEntityType.IBAN.value,
            PIIEntityType.IP_ADDRESS.value,
        }
    )

    # Action to take when PII is detected
    action: PIIAction = PIIAction.REDACT

    # Minimum confidence score to consider a detection valid
    min_confidence: float = 0.7

    # Allow list - patterns/values that should not be flagged
    allow_list: Set[str] = field(default_factory=set)

    # Deny list - additional patterns to always flag
    deny_list: Set[str] = field(default_factory=set)

    # Custom redaction format (default: <ENTITY_TYPE>)
    redaction_format: str = "<{entity_type}>"

    # Whether to log detected PII (for audit)
    log_detections: bool = True

    # Hash key for HASH action (should be set securely in production)
    hash_key: Optional[str] = None


class PIIDetector:
    """
    Detects and handles PII (Personally Identifiable Information).

    Uses Microsoft Presidio when available, falls back to regex patterns.

    Usage:
        detector = PIIDetector(PIIConfig(action=PIIAction.REDACT))

        # Check for PII
        has_pii, findings = detector.detect("My SSN is 123-45-6789")

        # Redact PII
        clean_text = detector.redact("My email is test@example.com")
        # Returns: "My email is <EMAIL>"

        # Check and act based on config
        result = detector.process("Contact me at 555-1234")
    """

    # Regex patterns for fallback detection
    PATTERNS = {
        PIIEntityType.SSN.value: [
            r"\b\d{3}-\d{2}-\d{4}\b",  # 123-45-6789
            r"\b\d{9}\b",  # 123456789 (less confident)
        ],
        PIIEntityType.CREDIT_CARD.value: [
            r"\b4\d{3}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b",  # Visa
            r"\b5[1-5]\d{2}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b",  # Mastercard
            r"\b3[47]\d{2}[\s-]?\d{6}[\s-]?\d{5}\b",  # Amex
            r"\b6(?:011|5\d{2})[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b",  # Discover
        ],
        PIIEntityType.EMAIL.value: [
            r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
        ],
        PIIEntityType.PHONE.value: [
            r"\b\d{3}[-.\s]?\d{3}[-.\s]?\d{4}\b",  # 555-123-4567, 555.123.4567, 555 123 4567
            r"\(\d{3}\)\s*\d{3}[-.\s]?\d{4}",  # (555) 123-4567, (555)123-4567
            r"\b\+1[-.\s]?\d{3}[-.\s]?\d{3}[-.\s]?\d{4}\b",  # +1-555-123-4567
            r"\b\+1[-.\s]?\(\d{3}\)\s*\d{3}[-.\s]?\d{4}",  # +1 (555) 123-4567
        ],
        PIIEntityType.IP_ADDRESS.value: [
            r"\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\b",
        ],
        PIIEntityType.IBAN.value: [
            r"\b[A-Z]{2}\d{2}[A-Z0-9]{4}\d{7}([A-Z0-9]?){0,16}\b",
        ],
    }

    # Confidence scores for regex patterns (lower than Presidio NER)
    PATTERN_CONFIDENCE = {
        PIIEntityType.SSN.value: 0.85,
        PIIEntityType.CREDIT_CARD.value: 0.90,
        PIIEntityType.EMAIL.value: 0.95,
        PIIEntityType.PHONE.value: 0.75,
        PIIEntityType.IP_ADDRESS.value: 0.90,
        PIIEntityType.IBAN.value: 0.85,
    }

    def __init__(self, config: Optional[PIIConfig] = None):
        self.config = config or PIIConfig()
        self._presidio_available = False
        self._analyzer = None
        self._anonymizer = None
        self._init_presidio()
        self._compile_patterns()

    def _init_presidio(self):
        """Try to initialize Presidio."""
        try:
            from presidio_analyzer import AnalyzerEngine
            from presidio_anonymizer import AnonymizerEngine

            self._analyzer = AnalyzerEngine()
            self._anonymizer = AnonymizerEngine()
            self._presidio_available = True
            logger.info("PIIDetector: Using Microsoft Presidio for PII detection")
        except ImportError:
            logger.warning("PIIDetector: Presidio not installed. Using regex fallback.")
        except Exception as e:
            logger.warning(f"PIIDetector: Failed to initialize Presidio: {e}. Using fallback.")

    def _compile_patterns(self):
        """
        Prepare patterns for detection.

        Note: We no longer pre-compile patterns since safe_search handles
        compilation with caching internally.
        """
        # Store patterns for use with safe_search
        self._pattern_list = self.PATTERNS

    def detect(self, text: str) -> tuple:
        """
        Detect PII in text.

        Args:
            text: Text to scan

        Returns:
            (has_pii: bool, findings: List[PIIFinding])
        """
        if not text:
            return (False, [])

        # Check allow list
        for allowed in self.config.allow_list:
            text = text.replace(allowed, " " * len(allowed))

        findings = []

        if self._presidio_available:
            findings = self._detect_presidio(text)
        else:
            findings = self._detect_regex(text)

        # Add deny list matches
        findings.extend(self._check_deny_list(text))

        # Filter by confidence and enabled entities
        findings = [
            f
            for f in findings
            if f.confidence >= self.config.min_confidence
            and f.entity_type in self.config.enabled_entities
        ]

        # Log if configured
        if self.config.log_detections and findings:
            logger.warning(
                f"PII detected: {len(findings)} entities - "
                f"{', '.join(set(f.entity_type for f in findings))}"
            )

        return (len(findings) > 0, findings)

    def _detect_presidio(self, text: str) -> List[PIIFinding]:
        """Detect PII using Presidio."""
        try:
            # Map our entity types to Presidio's
            presidio_entities = list(self.config.enabled_entities)

            results = self._analyzer.analyze(text=text, entities=presidio_entities, language="en")

            return [
                PIIFinding(
                    entity_type=r.entity_type,
                    text=text[r.start : r.end],
                    start=r.start,
                    end=r.end,
                    confidence=r.score,
                )
                for r in results
            ]
        except Exception as e:
            logger.warning(f"Presidio detection failed: {e}. Falling back to regex.")
            return self._detect_regex(text)

    def _detect_regex(self, text: str) -> List[PIIFinding]:
        """
        Detect PII using regex patterns with ReDoS protection.

        Uses safe_search from pattern_safety module.
        """
        findings = []

        for entity_type, patterns in self.PATTERNS.items():
            if entity_type not in self.config.enabled_entities:
                continue

            confidence = self.PATTERN_CONFIDENCE.get(entity_type, 0.7)

            for pattern in patterns:
                # Use re.finditer with compiled pattern for multiple matches
                # safe_search handles single matches, so we compile carefully here
                try:
                    compiled = re.compile(pattern, re.IGNORECASE)
                    # Limit input length for safety
                    search_text = text[:100_000] if len(text) > 100_000 else text
                    for match in compiled.finditer(search_text):
                        findings.append(
                            PIIFinding(
                                entity_type=entity_type,
                                text=match.group(),
                                start=match.start(),
                                end=match.end(),
                                confidence=confidence,
                            )
                        )
                except re.error as e:
                    logger.warning(f"Invalid PII pattern {pattern}: {e}")
                    continue

        # Deduplicate overlapping findings (keep highest confidence)
        findings = self._deduplicate_findings(findings)

        return findings

    def _check_deny_list(self, text: str) -> List[PIIFinding]:
        """Check text against deny list."""
        findings = []
        text_lower = text.lower()

        for denied in self.config.deny_list:
            denied_lower = denied.lower()
            start = 0
            while True:
                pos = text_lower.find(denied_lower, start)
                if pos == -1:
                    break
                findings.append(
                    PIIFinding(
                        entity_type=PIIEntityType.CUSTOM.value,
                        text=text[pos : pos + len(denied)],
                        start=pos,
                        end=pos + len(denied),
                        confidence=1.0,
                    )
                )
                start = pos + 1

        return findings

    def _deduplicate_findings(self, findings: List[PIIFinding]) -> List[PIIFinding]:
        """Remove overlapping findings, keeping highest confidence."""
        if not findings:
            return []

        # Sort by start position, then by confidence (descending)
        sorted_findings = sorted(findings, key=lambda f: (f.start, -f.confidence))

        result = []
        last_end = -1

        for finding in sorted_findings:
            if finding.start >= last_end:
                result.append(finding)
                last_end = finding.end
            elif finding.confidence > result[-1].confidence:
                # Higher confidence for overlapping region
                result[-1] = finding
                last_end = finding.end

        return result

    def redact(self, text: str, findings: List[PIIFinding] = None) -> str:
        """
        Redact PII from text.

        Args:
            text: Original text
            findings: Pre-computed findings (will detect if not provided)

        Returns:
            Text with PII replaced by placeholders
        """
        if findings is None:
            _, findings = self.detect(text)

        if not findings:
            return text

        # Sort by position (reverse order to preserve indices)
        sorted_findings = sorted(findings, key=lambda f: f.start, reverse=True)

        result = text
        for finding in sorted_findings:
            replacement = self.config.redaction_format.format(entity_type=finding.entity_type)
            result = result[: finding.start] + replacement + result[finding.end :]

        return result

    def hash_pii(self, text: str, findings: List[PIIFinding] = None) -> tuple:
        """
        Replace PII with hashed values (reversible with mapping).

        Args:
            text: Original text
            findings: Pre-computed findings

        Returns:
            (hashed_text: str, mapping: Dict[str, str])
        """
        import hashlib

        if findings is None:
            _, findings = self.detect(text)

        if not findings:
            return (text, {})

        mapping = {}
        sorted_findings = sorted(findings, key=lambda f: f.start, reverse=True)

        result = text
        key = (self.config.hash_key or "familiar-pii").encode()

        for finding in sorted_findings:
            # Create deterministic hash
            hash_input = f"{finding.entity_type}:{finding.text}".encode()
            hashed = hashlib.blake2b(hash_input, key=key, digest_size=8).hexdigest()

            placeholder = f"<{finding.entity_type}:{hashed}>"
            mapping[placeholder] = finding.text

            result = result[: finding.start] + placeholder + result[finding.end :]

        return (result, mapping)

    def process(self, text: str) -> tuple:
        """
        Process text according to configured action.

        Args:
            text: Text to process

        Returns:
            (processed_text: str, findings: List[PIIFinding], blocked: bool)

        Raises:
            GuardrailViolation: If action is BLOCK and PII is found
        """
        has_pii, findings = self.detect(text)

        if not has_pii:
            return (text, [], False)

        action = self.config.action

        if action == PIIAction.BLOCK:
            entity_types = list(set(f.entity_type for f in findings))
            raise GuardrailViolation(
                ViolationType.CONTENT_SAFETY,
                f"PII detected: {', '.join(entity_types)}",
                action=ViolationAction.BLOCK,
                details={"pii_findings": [f.to_dict() for f in findings]},
            )

        elif action == PIIAction.REDACT:
            processed = self.redact(text, findings)
            return (processed, findings, False)

        elif action == PIIAction.HASH:
            processed, mapping = self.hash_pii(text, findings)
            for f in findings:
                f.text = "<hashed>"  # Don't expose original in findings
            return (processed, findings, False)

        elif action == PIIAction.WARN:
            return (text, findings, False)

        return (text, findings, False)

    @property
    def is_presidio_available(self) -> bool:
        """Check if Presidio is available."""
        return self._presidio_available


# ============================================================
# MAIN GUARDRAILS CLASS
# ============================================================


class Guardrails:
    """
    Runtime guardrails for agent execution.

    Tracks usage across requests and enforces limits.
    """

    def __init__(self, config: Optional[GuardrailConfig] = None):
        self.config = config or GuardrailConfig()

        # Per-user tracking
        self._user_history: Dict[str, UserUsageHistory] = defaultdict(UserUsageHistory)
        self._lock = threading.RLock()

        # Content filter
        self._content_filter = ContentFilter() if self.config.enable_content_filter else None

        # PII detector (v1.4.0)
        self._pii_detector = None
        if self.config.enable_pii_detection:
            pii_config = self.config.pii_config or PIIConfig()
            self._pii_detector = PIIDetector(pii_config)

        # Current request tracker (set per request)
        self._current_tracker: Optional[UsageTracker] = None

        # PII config reference for compliance mode overrides
        self.pii_config = self.config.pii_config or (
            PIIConfig() if self.config.enable_pii_detection else None
        )

    def set_compliance_mode(self, mode) -> None:
        """Apply compliance-mode overrides to PII handling.

        When mode is HIPAA and PII detection is already enabled by the user,
        upgrades the action to BLOCK. Does NOT force-enable PII detection —
        that is always a manual toggle via the Settings panel.
        """
        if mode is None:
            return
        mode_str = mode.value if hasattr(mode, "value") else str(mode)
        if mode_str == "hipaa":
            if self._pii_detector:
                # User already enabled PII detection — upgrade action to BLOCK
                self._pii_detector.config.action = PIIAction.BLOCK
                if self.pii_config:
                    self.pii_config.action = PIIAction.BLOCK
                logger.info("Guardrails: HIPAA mode — PII action upgraded to BLOCK")
            else:
                logger.info(
                    "Guardrails: HIPAA mode active but PII detection not enabled. "
                    "Enable it manually in Settings to activate PHI blocking."
                )

    # --------------------------------------------------------
    # REQUEST LIFECYCLE
    # --------------------------------------------------------

    def start_request(self, user_id: str = "default") -> UsageTracker:
        """
        Start tracking a new request.

        Call this at the beginning of Agent.chat()
        """
        with self._lock:
            # Check rate limits
            history = self._user_history[user_id]

            if history.requests_last_minute() >= self.config.max_requests_per_minute:
                if self.config.rate_violation_action == ViolationAction.BLOCK:
                    raise GuardrailViolation(
                        ViolationType.RATE_LIMIT,
                        f"Rate limit exceeded: {history.requests_last_minute()} requests/minute",
                        current_value=history.requests_last_minute(),
                        limit_value=self.config.max_requests_per_minute,
                    )
                elif self.config.rate_violation_action == ViolationAction.THROTTLE:
                    # Release lock before sleeping to avoid thread starvation
                    self._lock.release()
                    try:
                        time.sleep(1)
                    finally:
                        self._lock.acquire()

            # Check hourly/daily cost limits
            if history.cost_last_hour() >= self.config.max_cost_per_hour:
                raise GuardrailViolation(
                    ViolationType.COST_LIMIT,
                    f"Hourly cost limit exceeded: ${history.cost_last_hour():.2f}",
                    current_value=history.cost_last_hour(),
                    limit_value=self.config.max_cost_per_hour,
                )

            if history.cost_today() >= self.config.max_cost_per_day:
                raise GuardrailViolation(
                    ViolationType.COST_LIMIT,
                    f"Daily cost limit exceeded: ${history.cost_today():.2f}",
                    current_value=history.cost_today(),
                    limit_value=self.config.max_cost_per_day,
                )

            # Create tracker
            tracker = UsageTracker()
            self._current_tracker = tracker

            return tracker

    def end_request(self, user_id: str = "default", tracker: Optional[UsageTracker] = None):
        """
        Finalize request tracking.

        Call this at the end of Agent.chat()
        """
        tracker = tracker or self._current_tracker
        if not tracker:
            return

        with self._lock:
            self._user_history[user_id].record_request(tracker.total_cost)
            self._current_tracker = None

    # --------------------------------------------------------
    # PER-ITERATION CHECKS
    # --------------------------------------------------------

    def check_iteration(self, iteration: int, tracker: Optional[UsageTracker] = None):
        """
        Check guardrails at the start of each agent loop iteration.

        Raises GuardrailViolation if limits exceeded.
        """
        tracker = tracker or self._current_tracker
        if tracker:
            tracker.start_iteration()

        # Iteration limit
        if iteration >= self.config.max_iterations:
            raise GuardrailViolation(
                ViolationType.ITERATION_LIMIT,
                f"Maximum iterations ({self.config.max_iterations}) reached",
                current_value=iteration,
                limit_value=self.config.max_iterations,
            )

        # Timeout check
        if tracker and tracker.elapsed_seconds > self.config.request_timeout_seconds:
            raise GuardrailViolation(
                ViolationType.TIMEOUT,
                f"Request timeout ({self.config.request_timeout_seconds}s) exceeded",
                current_value=tracker.elapsed_seconds,
                limit_value=self.config.request_timeout_seconds,
            )

    def check_pre_llm_call(self, estimated_tokens: int, tracker: Optional[UsageTracker] = None):
        """
        Check before making an LLM call.

        Can be used to estimate if we'll exceed limits.
        """
        tracker = tracker or self._current_tracker
        if not tracker:
            return

        projected_total = tracker.total_tokens + estimated_tokens

        # Token limit check
        if projected_total > self.config.max_tokens_per_request:
            raise GuardrailViolation(
                ViolationType.TOKEN_LIMIT,
                f"Would exceed token limit: {projected_total:,} > {self.config.max_tokens_per_request:,}",
                current_value=projected_total,
                limit_value=self.config.max_tokens_per_request,
                action=self.config.token_violation_action,
            )

        # Warning threshold
        if projected_total > self.config.max_tokens_per_request * self.config.warn_tokens_threshold:
            logger.warning(
                f"Approaching token limit: {projected_total:,} / {self.config.max_tokens_per_request:,}"
            )

    def track_llm_usage(
        self,
        input_tokens: int,
        output_tokens: int,
        cost: float,
        tracker: Optional[UsageTracker] = None,
    ):
        """
        Record usage from an LLM call.

        Call after each provider.chat() response.
        """
        tracker = tracker or self._current_tracker
        if not tracker:
            return

        tracker.add_usage(input_tokens, output_tokens, cost)

        # Check limits after recording
        if tracker.total_tokens > self.config.max_tokens_per_request:
            if self.config.token_violation_action == ViolationAction.BLOCK:
                raise GuardrailViolation(
                    ViolationType.TOKEN_LIMIT,
                    f"Token limit exceeded: {tracker.total_tokens:,}",
                    current_value=tracker.total_tokens,
                    limit_value=self.config.max_tokens_per_request,
                )

        if tracker.total_cost > self.config.max_cost_per_request:
            if self.config.cost_violation_action == ViolationAction.BLOCK:
                raise GuardrailViolation(
                    ViolationType.COST_LIMIT,
                    f"Cost limit exceeded: ${tracker.total_cost:.4f}",
                    current_value=tracker.total_cost,
                    limit_value=self.config.max_cost_per_request,
                )

    # --------------------------------------------------------
    # TOOL CHECKS
    # --------------------------------------------------------

    def check_tool_call(
        self,
        tool_name: str,
        tool_input: dict,
        tool_registry: "ToolRegistry" = None,  # v1.4.0: Optional registry for safety checks
    ) -> tuple:
        """
        Check if a tool call is allowed.

        Args:
            tool_name: Name of the tool
            tool_input: Input parameters for the tool
            tool_registry: Optional ToolRegistry for safety level checks (v1.4.0)

        Returns:
            (allowed: bool, reason: Optional[str], requires_approval: bool)
        """
        # Blocked tools
        if tool_name in self.config.blocked_tools:
            return (False, f"Tool '{tool_name}' is blocked", False)

        # v1.4.0: Check tool registry for safety level
        if tool_registry is not None:
            tool_def = tool_registry.get(tool_name)
            if tool_def is not None:
                # Import here to avoid circular imports
                from .tool_registry import SafetyLevel, ToolStatus

                # Check if tool is disabled
                if tool_def.status == ToolStatus.DISABLED:
                    return (False, f"Tool '{tool_name}' is disabled", False)

                # Check safety level - dangerous/critical require approval
                if tool_def.safety_level in (SafetyLevel.DANGEROUS, SafetyLevel.CRITICAL):
                    return (
                        True,
                        f"Tool '{tool_name}' is {tool_def.safety_level.value} - requires approval",
                        True,
                    )

                # Validate input against tool schema
                is_valid, errors = tool_def.validate_input(tool_input)
                if not is_valid:
                    return (False, f"Invalid input for '{tool_name}': {'; '.join(errors)}", False)

        # Approval required (explicit config)
        if tool_name in self.config.require_approval_tools:
            return (True, f"Tool '{tool_name}' requires approval", True)

        # Track call count
        if self._current_tracker:
            self._current_tracker.add_tool_call()

        return (True, None, False)

    # --------------------------------------------------------
    # CONTENT SAFETY
    # --------------------------------------------------------

    def check_content(self, text: str) -> tuple:
        """
        Check content for safety (content filter + PII).

        Returns:
            (is_safe: bool, reason: Optional[str])
        """
        # Check content filter first
        if self._content_filter:
            is_safe, reason, severity = self._content_filter.check(text)

            if not is_safe:
                raise GuardrailViolation(
                    ViolationType.CONTENT_SAFETY, reason, action=ViolationAction.BLOCK
                )

        # Check PII (this may raise if action is BLOCK)
        if self._pii_detector:
            processed_text, findings, blocked = self._pii_detector.process(text)
            if findings:
                return (True, f"PII detected: {len(findings)} entities")

        return (True, None)

    def check_pii(self, text: str) -> tuple:
        """
        Check text for PII specifically.

        Returns:
            (has_pii: bool, findings: List[PIIFinding])
        """
        if not self._pii_detector:
            return (False, [])

        return self._pii_detector.detect(text)

    def redact_pii(self, text: str) -> str:
        """
        Redact PII from text.

        Returns:
            Text with PII replaced by placeholders
        """
        if not self._pii_detector:
            return text

        return self._pii_detector.redact(text)

    def process_with_pii(self, text: str) -> tuple:
        """
        Process text through all content checks including PII.

        Returns:
            (processed_text: str, findings: List[PIIFinding], blocked: bool)

        Raises:
            GuardrailViolation: If content is blocked
        """
        # Content filter (may raise)
        self.check_content(text)

        # PII detection (may redact or raise depending on config)
        if self._pii_detector:
            return self._pii_detector.process(text)

        return (text, [], False)

    @property
    def pii_detector(self) -> Optional[PIIDetector]:
        """Access the PII detector directly."""
        return self._pii_detector

    # --------------------------------------------------------
    # STATUS & REPORTING
    # --------------------------------------------------------

    def get_usage_summary(self, tracker: Optional[UsageTracker] = None) -> dict:
        """Get current usage summary."""
        tracker = tracker or self._current_tracker
        if not tracker:
            return {}

        return {
            "total_tokens": tracker.total_tokens,
            "total_cost_usd": tracker.total_cost,
            "iterations": tracker.iterations,
            "tool_calls": tracker.tool_calls,
            "elapsed_seconds": tracker.elapsed_seconds,
            "token_utilization": tracker.total_tokens / self.config.max_tokens_per_request,
            "cost_utilization": tracker.total_cost / self.config.max_cost_per_request,
        }

    def get_user_status(self, user_id: str) -> dict:
        """Get usage status for a user."""
        with self._lock:
            history = self._user_history.get(user_id)
            if not history:
                return {"status": "no_history"}

            return {
                "requests_last_minute": history.requests_last_minute(),
                "requests_last_hour": history.requests_last_hour(),
                "cost_last_hour": history.cost_last_hour(),
                "cost_today": history.cost_today(),
                "rate_limit_remaining": self.config.max_requests_per_minute
                - history.requests_last_minute(),
                "hourly_budget_remaining": self.config.max_cost_per_hour - history.cost_last_hour(),
                "daily_budget_remaining": self.config.max_cost_per_day - history.cost_today(),
            }


# ============================================================
# CONVENIENCE FUNCTIONS
# ============================================================

# Global default instance
_default_guardrails: Optional[Guardrails] = None


def get_default_guardrails() -> Guardrails:
    """Get or create the default guardrails instance."""
    global _default_guardrails
    if _default_guardrails is None:
        _default_guardrails = Guardrails()
    return _default_guardrails


def set_default_guardrails(guardrails: Guardrails):
    """Set the default guardrails instance."""
    global _default_guardrails
    _default_guardrails = guardrails


def check_guardrails(
    iteration: int = 0,
    tokens: int = 0,
    cost: float = 0.0,
    content: str = "",
    user_id: str = "default",
):
    """
    Convenience function to check all guardrails at once.

    Raises GuardrailViolation if any limit is exceeded.
    """
    guardrails = get_default_guardrails()

    guardrails.check_iteration(iteration)

    if tokens > 0:
        guardrails.check_pre_llm_call(tokens)

    if content:
        guardrails.check_content(content)


# ============================================================
# AUDIT LOGGING
# ============================================================


class AuditEventType(str, Enum):
    """Types of audit events."""

    # Request lifecycle
    REQUEST_START = "request_start"
    REQUEST_END = "request_end"
    REQUEST_ERROR = "request_error"

    # Guardrail events
    GUARDRAIL_CHECK = "guardrail_check"
    GUARDRAIL_VIOLATION = "guardrail_violation"
    GUARDRAIL_WARNING = "guardrail_warning"

    # PII events
    PII_DETECTED = "pii_detected"
    PII_REDACTED = "pii_redacted"
    PII_BLOCKED = "pii_blocked"

    # Content safety
    CONTENT_BLOCKED = "content_blocked"
    CONTENT_WARNING = "content_warning"

    # Usage events
    TOKEN_USAGE = "token_usage"
    COST_INCURRED = "cost_incurred"
    RATE_LIMIT_HIT = "rate_limit_hit"

    # Tool events
    TOOL_CALLED = "tool_called"
    TOOL_BLOCKED = "tool_blocked"
    TOOL_RESULT = "tool_result"

    # Custom
    CUSTOM = "custom"


@dataclass
class AuditEvent:
    """
    A structured audit event for compliance logging.

    All fields are designed to be JSON-serializable and
    compatible with log aggregation systems (ELK, Splunk, etc.)
    """

    # Event identity
    event_id: str = ""
    event_type: str = ""
    timestamp: str = ""

    # Request context
    request_id: str = ""
    user_id: str = ""
    session_id: str = ""

    # Event details
    action: str = ""  # What happened
    status: str = ""  # success, failure, warning
    message: str = ""  # Human-readable description

    # Metrics
    tokens_used: int = 0
    cost_usd: float = 0.0
    latency_ms: float = 0.0

    # Context
    skill_name: str = ""
    tool_name: str = ""
    model_name: str = ""

    # PII-specific (redacted)
    pii_entity_types: List[str] = field(default_factory=list)
    pii_count: int = 0

    # Error details
    error_type: str = ""
    error_message: str = ""

    # Additional metadata
    metadata: Dict[str, Any] = field(default_factory=dict)

    # Hashes for integrity
    input_hash: str = ""  # Hash of input (for correlation without storing content)
    output_hash: str = ""  # Hash of output

    def __post_init__(self):
        import secrets

        if not self.event_id:
            self.event_id = secrets.token_hex(12)
        if not self.timestamp:
            self.timestamp = datetime.now().isoformat()

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "event_id": self.event_id,
            "event_type": self.event_type,
            "timestamp": self.timestamp,
            "request_id": self.request_id,
            "user_id": self.user_id,
            "session_id": self.session_id,
            "action": self.action,
            "status": self.status,
            "message": self.message,
            "tokens_used": self.tokens_used,
            "cost_usd": self.cost_usd,
            "latency_ms": self.latency_ms,
            "skill_name": self.skill_name,
            "tool_name": self.tool_name,
            "model_name": self.model_name,
            "pii_entity_types": self.pii_entity_types,
            "pii_count": self.pii_count,
            "error_type": self.error_type,
            "error_message": self.error_message,
            "metadata": self.metadata,
            "input_hash": self.input_hash,
            "output_hash": self.output_hash,
        }

    def to_json(self, indent: int = None) -> str:
        """Convert to JSON string."""
        import json

        return json.dumps(self.to_dict(), indent=indent, default=str)

    @classmethod
    def from_dict(cls, data: dict) -> "AuditEvent":
        """Create from dictionary."""
        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})


@dataclass
class AuditConfig:
    """Configuration for audit logging."""

    # Output destinations
    log_to_file: bool = True
    log_file_path: str = "audit.jsonl"
    log_to_stdout: bool = False
    log_to_logger: bool = True  # Use Python logging

    # What to log
    log_request_lifecycle: bool = True
    log_guardrail_checks: bool = True
    log_pii_events: bool = True
    log_content_safety: bool = True
    log_usage_metrics: bool = True
    log_tool_calls: bool = True

    # Privacy settings
    hash_inputs: bool = True  # Store hash instead of raw input
    hash_outputs: bool = True  # Store hash instead of raw output
    redact_pii_in_logs: bool = True  # Redact any PII in log messages

    # Retention (for file rotation)
    max_file_size_mb: int = 100
    max_files: int = 10

    # Custom handler
    custom_handler: Optional[Callable[[AuditEvent], None]] = None

    # Filtering
    min_severity: str = "info"  # debug, info, warning, error
    exclude_event_types: Set[str] = field(default_factory=set)


class AuditLogger:
    """
    Structured audit logging for compliance and observability.

    Provides:
    - JSON-structured logging (JSONL format)
    - Multiple output destinations (file, stdout, Python logging, custom)
    - Privacy-preserving hashing of inputs/outputs
    - Automatic PII redaction in logs
    - Request correlation via request_id
    - Compliance-ready event structure

    Usage:
        # Basic usage
        audit = AuditLogger()
        audit.log_request_start(request_id="req-123", user_id="user-456")
        audit.log_pii_detected(request_id="req-123", entity_types=["SSN", "EMAIL"])
        audit.log_request_end(request_id="req-123", tokens=1500, cost=0.05)

        # With custom handler
        def my_handler(event: AuditEvent):
            send_to_splunk(event.to_dict())

        audit = AuditLogger(AuditConfig(custom_handler=my_handler))

        # Get events for analysis
        events = audit.get_events(request_id="req-123")
    """

    SEVERITY_LEVELS = {"debug": 0, "info": 1, "warning": 2, "error": 3}

    def __init__(self, config: Optional[AuditConfig] = None):
        self.config = config or AuditConfig()
        self._lock = threading.RLock()
        self._events: List[AuditEvent] = []  # In-memory buffer
        self._max_buffer_size = 1000
        self._file_handle = None
        self._pii_detector = None

        # Initialize file logging
        if self.config.log_to_file:
            self._init_file_logging()

        # Initialize PII detector for log redaction
        if self.config.redact_pii_in_logs:
            try:
                self._pii_detector = PIIDetector(PIIConfig(action=PIIAction.REDACT))
            except Exception:
                pass  # PII detection optional for audit logging

    def _init_file_logging(self):
        """Initialize file-based logging."""
        import os

        log_dir = os.path.dirname(self.config.log_file_path)
        if log_dir and not os.path.exists(log_dir):
            os.makedirs(log_dir, exist_ok=True)

    def _hash_content(self, content: str) -> str:
        """Create a hash of content for correlation without storing raw data."""
        if not content:
            return ""
        import hashlib

        return hashlib.sha256(content.encode()).hexdigest()[:16]

    def _redact_message(self, message: str) -> str:
        """Redact any PII from log messages."""
        if not message or not self._pii_detector:
            return message
        try:
            return self._pii_detector.redact(message)
        except Exception:
            return message

    def _should_log(self, event_type: str, severity: str = "info") -> bool:
        """Check if this event should be logged based on config."""
        # Check severity
        min_level = self.SEVERITY_LEVELS.get(self.config.min_severity, 1)
        event_level = self.SEVERITY_LEVELS.get(severity, 1)
        if event_level < min_level:
            return False

        # Check exclusions
        if event_type in self.config.exclude_event_types:
            return False

        # Check category
        if event_type.startswith("request_") and not self.config.log_request_lifecycle:
            return False
        if event_type.startswith("guardrail_") and not self.config.log_guardrail_checks:
            return False
        if event_type.startswith("pii_") and not self.config.log_pii_events:
            return False
        if event_type.startswith("content_") and not self.config.log_content_safety:
            return False
        if event_type in ("token_usage", "cost_incurred") and not self.config.log_usage_metrics:
            return False
        if event_type.startswith("tool_") and not self.config.log_tool_calls:
            return False

        return True

    def _write_event(self, event: AuditEvent, severity: str = "info"):
        """Write event to all configured destinations."""
        if not self._should_log(event.event_type, severity):
            return

        # Redact PII in message if configured
        if self.config.redact_pii_in_logs:
            event.message = self._redact_message(event.message)

        with self._lock:
            # Buffer in memory
            self._events.append(event)
            if len(self._events) > self._max_buffer_size:
                self._events.pop(0)

            # Write to file
            if self.config.log_to_file:
                self._write_to_file(event)

            # Write to stdout
            if self.config.log_to_stdout:
                print(event.to_json())

            # Write to Python logger
            if self.config.log_to_logger:
                log_level = getattr(logging, severity.upper(), logging.INFO)
                logger.log(
                    log_level,
                    f"[AUDIT] {event.event_type}: {event.message}",
                    extra={"audit_event": event.to_dict()},
                )

            # Custom handler
            if self.config.custom_handler:
                try:
                    self.config.custom_handler(event)
                except Exception as e:
                    logger.error(f"Audit custom handler error: {e}")

    def _write_to_file(self, event: AuditEvent):
        """Write event to JSONL file."""
        try:
            with open(self.config.log_file_path, "a") as f:
                f.write(event.to_json() + "\n")
        except Exception as e:
            logger.error(f"Failed to write audit log: {e}")

    # --------------------------------------------------------
    # REQUEST LIFECYCLE
    # --------------------------------------------------------

    def log_request_start(
        self,
        request_id: str,
        user_id: str = "",
        session_id: str = "",
        input_text: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log the start of a request."""
        event = AuditEvent(
            event_type=AuditEventType.REQUEST_START.value,
            request_id=request_id,
            user_id=user_id,
            session_id=session_id,
            action="request_started",
            status="started",
            message=f"Request started for user {user_id}",
            input_hash=self._hash_content(input_text) if self.config.hash_inputs else "",
            metadata=metadata or {},
        )
        self._write_event(event)
        return event

    def log_request_end(
        self,
        request_id: str,
        user_id: str = "",
        tokens_used: int = 0,
        cost_usd: float = 0.0,
        latency_ms: float = 0.0,
        output_text: str = "",
        status: str = "success",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log the end of a request."""
        event = AuditEvent(
            event_type=AuditEventType.REQUEST_END.value,
            request_id=request_id,
            user_id=user_id,
            action="request_completed",
            status=status,
            message=f"Request completed: {tokens_used} tokens, ${cost_usd:.4f}",
            tokens_used=tokens_used,
            cost_usd=cost_usd,
            latency_ms=latency_ms,
            output_hash=self._hash_content(output_text) if self.config.hash_outputs else "",
            metadata=metadata or {},
        )
        self._write_event(event)
        return event

    def log_request_error(
        self, request_id: str, error: Exception, user_id: str = "", metadata: Dict = None
    ) -> AuditEvent:
        """Log a request error."""
        event = AuditEvent(
            event_type=AuditEventType.REQUEST_ERROR.value,
            request_id=request_id,
            user_id=user_id,
            action="request_failed",
            status="error",
            message=f"Request failed: {str(error)}",
            error_type=type(error).__name__,
            error_message=str(error),
            metadata=metadata or {},
        )
        self._write_event(event, severity="error")
        return event

    # --------------------------------------------------------
    # GUARDRAIL EVENTS
    # --------------------------------------------------------

    def log_guardrail_check(
        self,
        request_id: str,
        check_type: str,
        passed: bool,
        current_value: Any = None,
        limit_value: Any = None,
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log a guardrail check."""
        event = AuditEvent(
            event_type=AuditEventType.GUARDRAIL_CHECK.value,
            request_id=request_id,
            action=f"guardrail_check_{check_type}",
            status="passed" if passed else "failed",
            message=f"Guardrail {check_type}: {'passed' if passed else 'failed'} ({current_value}/{limit_value})",
            metadata={
                "check_type": check_type,
                "current_value": current_value,
                "limit_value": limit_value,
                **(metadata or {}),
            },
        )
        self._write_event(event)
        return event

    def log_guardrail_violation(
        self,
        request_id: str,
        violation: "GuardrailViolation",
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log a guardrail violation."""
        event = AuditEvent(
            event_type=AuditEventType.GUARDRAIL_VIOLATION.value,
            request_id=request_id,
            user_id=user_id,
            action=f"violation_{violation.violation_type.value}",
            status="blocked",
            message=str(violation),
            error_type=violation.violation_type.value,
            error_message=str(violation),
            metadata={
                "violation_type": violation.violation_type.value,
                "current_value": violation.current_value,
                "limit_value": violation.limit_value,
                "action_taken": violation.action.value,
                **(metadata or {}),
            },
        )
        self._write_event(event, severity="warning")
        return event

    # --------------------------------------------------------
    # PII EVENTS
    # --------------------------------------------------------

    def log_pii_detected(
        self,
        request_id: str,
        entity_types: List[str],
        count: int,
        action_taken: str = "detected",
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log PII detection."""
        event = AuditEvent(
            event_type=AuditEventType.PII_DETECTED.value,
            request_id=request_id,
            user_id=user_id,
            action="pii_detected",
            status="detected",
            message=f"PII detected: {count} entities ({', '.join(entity_types)})",
            pii_entity_types=entity_types,
            pii_count=count,
            metadata={"action_taken": action_taken, **(metadata or {})},
        )
        self._write_event(event, severity="warning")
        return event

    def log_pii_redacted(
        self,
        request_id: str,
        entity_types: List[str],
        count: int,
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log PII redaction."""
        event = AuditEvent(
            event_type=AuditEventType.PII_REDACTED.value,
            request_id=request_id,
            user_id=user_id,
            action="pii_redacted",
            status="redacted",
            message=f"PII redacted: {count} entities",
            pii_entity_types=entity_types,
            pii_count=count,
            metadata=metadata or {},
        )
        self._write_event(event)
        return event

    def log_pii_blocked(
        self,
        request_id: str,
        entity_types: List[str],
        count: int,
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log PII block action."""
        event = AuditEvent(
            event_type=AuditEventType.PII_BLOCKED.value,
            request_id=request_id,
            user_id=user_id,
            action="pii_blocked",
            status="blocked",
            message=f"Request blocked due to PII: {', '.join(entity_types)}",
            pii_entity_types=entity_types,
            pii_count=count,
            metadata=metadata or {},
        )
        self._write_event(event, severity="warning")
        return event

    # --------------------------------------------------------
    # CONTENT SAFETY EVENTS
    # --------------------------------------------------------

    def log_content_blocked(
        self, request_id: str, reason: str, user_id: str = "", metadata: Dict = None
    ) -> AuditEvent:
        """Log content being blocked."""
        event = AuditEvent(
            event_type=AuditEventType.CONTENT_BLOCKED.value,
            request_id=request_id,
            user_id=user_id,
            action="content_blocked",
            status="blocked",
            message=f"Content blocked: {reason}",
            metadata=metadata or {},
        )
        self._write_event(event, severity="warning")
        return event

    def log_content_warning(
        self, request_id: str, reason: str, user_id: str = "", metadata: Dict = None
    ) -> AuditEvent:
        """Log a content warning."""
        event = AuditEvent(
            event_type=AuditEventType.CONTENT_WARNING.value,
            request_id=request_id,
            user_id=user_id,
            action="content_warning",
            status="warning",
            message=f"Content warning: {reason}",
            metadata=metadata or {},
        )
        self._write_event(event, severity="warning")
        return event

    # --------------------------------------------------------
    # USAGE EVENTS
    # --------------------------------------------------------

    def log_token_usage(
        self,
        request_id: str,
        tokens: int,
        model_name: str = "",
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log token usage."""
        event = AuditEvent(
            event_type=AuditEventType.TOKEN_USAGE.value,
            request_id=request_id,
            user_id=user_id,
            model_name=model_name,
            action="token_usage",
            status="recorded",
            message=f"Tokens used: {tokens}",
            tokens_used=tokens,
            metadata=metadata or {},
        )
        self._write_event(event, severity="debug")
        return event

    def log_cost_incurred(
        self,
        request_id: str,
        cost_usd: float,
        model_name: str = "",
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log cost incurred."""
        event = AuditEvent(
            event_type=AuditEventType.COST_INCURRED.value,
            request_id=request_id,
            user_id=user_id,
            model_name=model_name,
            action="cost_incurred",
            status="recorded",
            message=f"Cost: ${cost_usd:.4f}",
            cost_usd=cost_usd,
            metadata=metadata or {},
        )
        self._write_event(event, severity="debug")
        return event

    def log_rate_limit_hit(
        self,
        request_id: str,
        limit_type: str,
        current_value: int,
        limit_value: int,
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log rate limit being hit."""
        event = AuditEvent(
            event_type=AuditEventType.RATE_LIMIT_HIT.value,
            request_id=request_id,
            user_id=user_id,
            action="rate_limit_hit",
            status="throttled",
            message=f"Rate limit hit: {limit_type} ({current_value}/{limit_value})",
            metadata={
                "limit_type": limit_type,
                "current_value": current_value,
                "limit_value": limit_value,
                **(metadata or {}),
            },
        )
        self._write_event(event, severity="warning")
        return event

    # --------------------------------------------------------
    # TOOL EVENTS
    # --------------------------------------------------------

    def log_tool_called(
        self,
        request_id: str,
        tool_name: str,
        skill_name: str = "",
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log a tool being called."""
        event = AuditEvent(
            event_type=AuditEventType.TOOL_CALLED.value,
            request_id=request_id,
            user_id=user_id,
            tool_name=tool_name,
            skill_name=skill_name,
            action="tool_called",
            status="invoked",
            message=f"Tool called: {skill_name}.{tool_name}"
            if skill_name
            else f"Tool called: {tool_name}",
            metadata=metadata or {},
        )
        self._write_event(event, severity="debug")
        return event

    def log_tool_blocked(
        self,
        request_id: str,
        tool_name: str,
        reason: str,
        skill_name: str = "",
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log a tool being blocked."""
        event = AuditEvent(
            event_type=AuditEventType.TOOL_BLOCKED.value,
            request_id=request_id,
            user_id=user_id,
            tool_name=tool_name,
            skill_name=skill_name,
            action="tool_blocked",
            status="blocked",
            message=f"Tool blocked: {tool_name} - {reason}",
            metadata=metadata or {},
        )
        self._write_event(event, severity="warning")
        return event

    def log_tool_result(
        self,
        request_id: str,
        tool_name: str,
        success: bool,
        latency_ms: float = 0.0,
        skill_name: str = "",
        error: str = "",
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log a tool result."""
        event = AuditEvent(
            event_type=AuditEventType.TOOL_RESULT.value,
            request_id=request_id,
            user_id=user_id,
            tool_name=tool_name,
            skill_name=skill_name,
            action="tool_completed",
            status="success" if success else "error",
            message=f"Tool {tool_name}: {'success' if success else 'failed'} ({latency_ms:.0f}ms)",
            latency_ms=latency_ms,
            error_message=error if not success else "",
            metadata=metadata or {},
        )
        self._write_event(event, severity="debug" if success else "warning")
        return event

    # --------------------------------------------------------
    # CUSTOM EVENTS
    # --------------------------------------------------------

    def log_custom(
        self,
        request_id: str,
        action: str,
        message: str,
        status: str = "info",
        severity: str = "info",
        user_id: str = "",
        metadata: Dict = None,
    ) -> AuditEvent:
        """Log a custom audit event."""
        event = AuditEvent(
            event_type=AuditEventType.CUSTOM.value,
            request_id=request_id,
            user_id=user_id,
            action=action,
            status=status,
            message=message,
            metadata=metadata or {},
        )
        self._write_event(event, severity=severity)
        return event

    # --------------------------------------------------------
    # RETRIEVAL & ANALYSIS
    # --------------------------------------------------------

    def get_events(
        self, request_id: str = None, user_id: str = None, event_type: str = None, limit: int = 100
    ) -> List[AuditEvent]:
        """
        Get events from the in-memory buffer.

        Args:
            request_id: Filter by request ID
            user_id: Filter by user ID
            event_type: Filter by event type
            limit: Maximum events to return
        """
        with self._lock:
            events = self._events.copy()

        # Apply filters
        if request_id:
            events = [e for e in events if e.request_id == request_id]
        if user_id:
            events = [e for e in events if e.user_id == user_id]
        if event_type:
            events = [e for e in events if e.event_type == event_type]

        return events[-limit:]

    def get_request_summary(self, request_id: str) -> Dict[str, Any]:
        """Get a summary of all events for a request."""
        events = self.get_events(request_id=request_id)

        if not events:
            return {"request_id": request_id, "status": "not_found"}

        return {
            "request_id": request_id,
            "event_count": len(events),
            "event_types": list(set(e.event_type for e in events)),
            "total_tokens": sum(e.tokens_used for e in events),
            "total_cost_usd": sum(e.cost_usd for e in events),
            "pii_detected": any(e.pii_count > 0 for e in events),
            "errors": [e.error_message for e in events if e.error_message],
            "start_time": events[0].timestamp if events else None,
            "end_time": events[-1].timestamp if events else None,
        }

    def get_user_summary(self, user_id: str) -> Dict[str, Any]:
        """Get a summary of all events for a user."""
        events = self.get_events(user_id=user_id)

        return {
            "user_id": user_id,
            "event_count": len(events),
            "request_count": len(set(e.request_id for e in events)),
            "total_tokens": sum(e.tokens_used for e in events),
            "total_cost_usd": sum(e.cost_usd for e in events),
            "pii_events": len([e for e in events if e.event_type.startswith("pii_")]),
            "violations": len(
                [e for e in events if e.event_type == AuditEventType.GUARDRAIL_VIOLATION.value]
            ),
        }

    def export_jsonl(self, filepath: str, request_id: str = None) -> int:
        """
        Export events to a JSONL file.

        Returns number of events exported.
        """
        events = self.get_events(request_id=request_id, limit=10000)

        with open(filepath, "w") as f:
            for event in events:
                f.write(event.to_json() + "\n")

        return len(events)

    def clear_buffer(self):
        """Clear the in-memory event buffer."""
        with self._lock:
            self._events.clear()


# Global audit logger instance
_audit_logger: Optional[AuditLogger] = None


def get_audit_logger() -> AuditLogger:
    """Get the global audit logger instance."""
    global _audit_logger
    if _audit_logger is None:
        _audit_logger = AuditLogger()
    return _audit_logger


def set_audit_logger(logger: AuditLogger):
    """Set the global audit logger instance."""
    global _audit_logger
    _audit_logger = logger


# ============================================================
# DECORATORS
# ============================================================


def with_guardrails(config: Optional[GuardrailConfig] = None, user_id_arg: str = "user_id"):
    """
    Decorator to add guardrails to a function.

    Usage:
        @with_guardrails(GuardrailConfig(max_cost_per_request=0.50))
        def my_agent_function(message: str, user_id: str = "default"):
            ...
    """

    def decorator(func):
        import functools

        guardrails = Guardrails(config)

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Extract user_id from kwargs
            uid = kwargs.get(user_id_arg, "default")

            tracker = guardrails.start_request(uid)
            try:
                result = func(*args, **kwargs)
                return result
            finally:
                guardrails.end_request(uid, tracker)

        # Attach guardrails for introspection
        wrapper.guardrails = guardrails
        return wrapper

    return decorator
