# 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, forces PII action to BLOCK so PHI is stopped
        at the guardrail level before reaching any provider.  If PII
        detection was not enabled (default), it is force-created here.
        """
        if mode is None:
            return
        mode_str = mode.value if hasattr(mode, "value") else str(mode)
        if mode_str == "hipaa":
            if not self._pii_detector:
                # Force-enable PII detection in HIPAA mode even if default is off
                self._pii_detector = PIIDetector(PIIConfig(action=PIIAction.BLOCK))
                self.pii_config = self._pii_detector.config
            else:
                self._pii_detector.config.action = PIIAction.BLOCK
                if self.pii_config:
                    self.pii_config.action = PIIAction.BLOCK
            logger.info("Guardrails: HIPAA mode — PII detection force-enabled, action=BLOCK")

    # --------------------------------------------------------
    # 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
