# 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

"""
Familiar Enterprise Enhancements
===============================

Drop-in improvements for production readiness:
1. Accurate tokenization (tiktoken)
2. Retry policies with exponential backoff
3. OpenTelemetry observability wrapper
4. PII detection layer (Presidio integration)
5. LLM-as-Judge semantic scorer
6. Audit logging for compliance

Install dependencies:
    pip install tiktoken presidio-analyzer presidio-anonymizer opentelemetry-api opentelemetry-sdk

Usage:
    from familiar.core.enhancements import (
        TokenCounter, RetryPolicy, ObservabilityWrapper,
        PIIGuardrail, LLMJudge, AuditLogger
    )
"""

import functools
import hashlib
import json
import logging
import threading
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)


# ============================================================
# 1. ACCURATE TOKENIZATION
# ============================================================


class TokenCounter:
    """
    Accurate token counting using tiktoken with graceful fallback.

    Usage:
        counter = TokenCounter("cl100k_base")  # Claude/GPT-4 encoding
        count = counter.count("Hello world")
        count = counter.count_messages([{"role": "user", "content": "Hi"}])
    """

    _encodings: Dict[str, Any] = {}
    _lock = threading.Lock()

    def __init__(self, encoding_name: str = "cl100k_base"):
        self.encoding_name = encoding_name
        self._encoding = self._get_encoding(encoding_name)

    @classmethod
    def _get_encoding(cls, name: str):
        """Get or create encoding (cached, thread-safe)."""
        with cls._lock:
            if name not in cls._encodings:
                try:
                    import tiktoken

                    cls._encodings[name] = tiktoken.get_encoding(name)
                except ImportError:
                    logger.warning("tiktoken not available, using char/4 fallback")
                    cls._encodings[name] = None
            return cls._encodings[name]

    def count(self, text: str) -> int:
        """Count tokens in text."""
        if not text:
            return 0

        if self._encoding:
            return len(self._encoding.encode(text))
        else:
            # Fallback: improved heuristic
            # JSON/code: ~3.5 chars/token, prose: ~4 chars/token
            if text.startswith("{") or text.startswith("["):
                return len(text) // 3
            return len(text) // 4

    def count_messages(self, messages: List[Dict]) -> int:
        """Count tokens in a message list (with overhead)."""
        total = 0
        for msg in messages:
            total += 4  # Role + content separator overhead
            total += self.count(str(msg.get("content", "")))
            if msg.get("tool_calls"):
                total += self.count(json.dumps(msg["tool_calls"]))
        total += 2  # Beginning/end tokens
        return total

    def truncate_to_limit(self, text: str, max_tokens: int) -> str:
        """Truncate text to fit within token limit."""
        if self._encoding:
            tokens = self._encoding.encode(text)
            if len(tokens) <= max_tokens:
                return text
            return self._encoding.decode(tokens[:max_tokens])
        else:
            # Fallback
            max_chars = max_tokens * 4
            return text[:max_chars]


# ============================================================
# 2. RETRY POLICIES
# ============================================================


class RetryStrategy(str, Enum):
    """Retry backoff strategies."""

    FIXED = "fixed"
    LINEAR = "linear"
    EXPONENTIAL = "exponential"


@dataclass
class RetryPolicy:
    """
    Configurable retry with exponential backoff.

    Usage:
        policy = RetryPolicy(max_attempts=3, base_delay=1.0)

        @policy.wrap
        def flaky_function():
            ...

        # Or manual
        result = policy.execute(flaky_function, arg1, arg2)
    """

    max_attempts: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    strategy: RetryStrategy = RetryStrategy.EXPONENTIAL
    jitter: float = 0.1  # Random jitter factor
    retryable_exceptions: tuple = (Exception,)
    on_retry: Optional[Callable] = None  # Callback(attempt, exception, delay)

    def calculate_delay(self, attempt: int) -> float:
        """Calculate delay for given attempt number."""
        import random

        if self.strategy == RetryStrategy.FIXED:
            delay = self.base_delay
        elif self.strategy == RetryStrategy.LINEAR:
            delay = self.base_delay * attempt
        else:  # EXPONENTIAL
            delay = self.base_delay * (2 ** (attempt - 1))

        # Apply jitter
        jitter_range = delay * self.jitter
        delay += random.uniform(-jitter_range, jitter_range)

        return min(delay, self.max_delay)

    def execute(self, func: Callable, *args, **kwargs) -> Any:
        """Execute function with retry policy."""
        last_exception = None

        for attempt in range(1, self.max_attempts + 1):
            try:
                return func(*args, **kwargs)
            except self.retryable_exceptions as e:
                last_exception = e

                if attempt == self.max_attempts:
                    break

                delay = self.calculate_delay(attempt)

                if self.on_retry:
                    self.on_retry(attempt, e, delay)
                else:
                    logger.warning(f"Retry {attempt}/{self.max_attempts} after {delay:.2f}s: {e}")

                time.sleep(delay)

        raise last_exception

    def wrap(self, func: Callable) -> Callable:
        """Decorator to wrap function with retry."""

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return self.execute(func, *args, **kwargs)

        return wrapper


# ============================================================
# 3. OBSERVABILITY WRAPPER
# ============================================================


@dataclass
class TraceContext:
    """Context for distributed tracing."""

    trace_id: str
    span_id: str
    parent_span_id: Optional[str] = None
    attributes: Dict[str, Any] = field(default_factory=dict)
    start_time: float = field(default_factory=time.time)


class ObservabilityWrapper:
    """
    OpenTelemetry-compatible observability with file fallback.

    Usage:
        obs = ObservabilityWrapper(service_name="familiar")

        with obs.span("process_request") as span:
            span.set_attribute("user_id", "123")
            # ... do work ...
            span.add_event("checkpoint", {"step": "complete"})
    """

    def __init__(
        self,
        service_name: str = "familiar",
        enable_otel: bool = True,
        fallback_log_path: Optional[Path] = None,
    ):
        self.service_name = service_name
        self.fallback_log_path = fallback_log_path or Path.home() / ".familiar" / "traces.jsonl"

        self._tracer = None
        self._meter = None

        if enable_otel:
            self._init_otel()

    def _init_otel(self):
        """Initialize OpenTelemetry if available."""
        try:
            from opentelemetry import trace
            from opentelemetry.sdk.resources import Resource
            from opentelemetry.sdk.trace import TracerProvider

            resource = Resource.create({"service.name": self.service_name})
            provider = TracerProvider(resource=resource)
            trace.set_tracer_provider(provider)

            self._tracer = trace.get_tracer(self.service_name)
            logger.info("OpenTelemetry initialized")
        except ImportError:
            logger.warning("OpenTelemetry not available, using file fallback")

    def span(self, name: str, attributes: Dict = None):
        """Create a span context manager."""
        if self._tracer:
            return self._tracer.start_as_current_span(name, attributes=attributes)
        else:
            return self._fallback_span(name, attributes)

    def _fallback_span(self, name: str, attributes: Dict = None):
        """File-based span for when OTel is unavailable."""
        return FallbackSpan(name, attributes, self.fallback_log_path)

    def record_metric(self, name: str, value: float, attributes: Dict = None):
        """Record a metric value."""
        if self._meter:
            # Use OTel meter
            pass
        else:
            # Log to file
            self._log_to_file(
                {
                    "type": "metric",
                    "name": name,
                    "value": value,
                    "attributes": attributes,
                    "timestamp": datetime.now().isoformat(),
                }
            )

    def _log_to_file(self, data: dict):
        """Write trace/metric to fallback file."""
        self.fallback_log_path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.fallback_log_path, "a") as f:
            f.write(json.dumps(data) + "\n")


class FallbackSpan:
    """File-based span when OpenTelemetry is unavailable."""

    def __init__(self, name: str, attributes: Dict, log_path: Path):
        self.name = name
        self.attributes = attributes or {}
        self.log_path = log_path
        self.events = []
        self.start_time = None

        import secrets

        self.span_id = secrets.token_hex(8)

    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        duration_ms = (time.time() - self.start_time) * 1000

        data = {
            "type": "span",
            "name": self.name,
            "span_id": self.span_id,
            "duration_ms": duration_ms,
            "attributes": self.attributes,
            "events": self.events,
            "status": "error" if exc_type else "ok",
            "timestamp": datetime.now().isoformat(),
        }

        if exc_type:
            data["error"] = str(exc_val)

        self.log_path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.log_path, "a") as f:
            f.write(json.dumps(data) + "\n")

    def set_attribute(self, key: str, value: Any):
        self.attributes[key] = value

    def add_event(self, name: str, attributes: Dict = None):
        self.events.append({"name": name, "attributes": attributes, "timestamp": time.time()})


# ============================================================
# 4. PII DETECTION LAYER
# ============================================================


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

    BLOCK = "block"  # Reject the request
    REDACT = "redact"  # Replace PII with placeholders
    WARN = "warn"  # Log warning, continue
    HASH = "hash"  # Replace with hash (reversible lookup)


@dataclass
class PIIConfig:
    """PII detection configuration."""

    enabled: bool = True
    action: PIIAction = PIIAction.REDACT

    # Entity types to detect
    detect_emails: bool = True
    detect_phones: bool = True
    detect_ssn: bool = True
    detect_credit_cards: bool = True
    detect_names: bool = False  # Higher false positive rate
    detect_addresses: bool = False

    # Custom patterns (regex)
    custom_patterns: Dict[str, str] = field(default_factory=dict)

    # Allowlist patterns that should not be flagged
    allowlist: List[str] = field(default_factory=list)


class PIIGuardrail:
    """
    PII detection and protection layer.

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

    Usage:
        pii = PIIGuardrail(PIIConfig(action=PIIAction.REDACT))

        # Check content
        is_safe, findings = pii.check("My email is test@example.com")

        # Redact content
        clean_text = pii.redact("My SSN is 123-45-6789")
        # Returns: "My SSN is <SSN>"
    """

    # Fallback regex patterns
    PATTERNS = {
        "EMAIL": r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
        "PHONE": r"\b(?:\+?1[-.\s]?)?\(?[0-9]{3}\)?[-.\s]?[0-9]{3}[-.\s]?[0-9]{4}\b",
        "SSN": r"\b\d{3}[-\s]?\d{2}[-\s]?\d{4}\b",
        "CREDIT_CARD": r"\b(?:\d{4}[-\s]?){3}\d{4}\b",
        "IP_ADDRESS": r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b",
    }

    def __init__(self, config: PIIConfig = None):
        self.config = config or PIIConfig()
        self._analyzer = None
        self._anonymizer = None
        self._init_presidio()

    def _init_presidio(self):
        """Initialize Presidio if available."""
        if not self.config.enabled:
            return

        try:
            from presidio_analyzer import AnalyzerEngine
            from presidio_anonymizer import AnonymizerEngine

            self._analyzer = AnalyzerEngine()
            self._anonymizer = AnonymizerEngine()
            logger.info("Presidio PII detection initialized")
        except ImportError:
            logger.warning("Presidio not available, using regex fallback")

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

        Returns:
            (is_safe: bool, findings: List[Dict])
        """
        if not self.config.enabled:
            return (True, [])

        findings = []

        if self._analyzer:
            findings = self._check_presidio(text)
        else:
            findings = self._check_regex(text)

        # Apply allowlist
        findings = [f for f in findings if not self._is_allowlisted(f["text"])]

        is_safe = len(findings) == 0
        return (is_safe, findings)

    def _check_presidio(self, text: str) -> List[Dict]:
        """Check using Presidio."""
        entities = self._get_enabled_entities()
        results = self._analyzer.analyze(text=text, entities=entities, language="en")

        return [
            {
                "type": r.entity_type,
                "text": text[r.start : r.end],
                "start": r.start,
                "end": r.end,
                "score": r.score,
            }
            for r in results
        ]

    def _check_regex(self, text: str) -> List[Dict]:
        """Check using regex fallback."""
        import re

        findings = []

        pattern_map = {
            "EMAIL": self.config.detect_emails,
            "PHONE": self.config.detect_phones,
            "SSN": self.config.detect_ssn,
            "CREDIT_CARD": self.config.detect_credit_cards,
        }

        for entity_type, enabled in pattern_map.items():
            if enabled and entity_type in self.PATTERNS:
                for match in re.finditer(self.PATTERNS[entity_type], text, re.IGNORECASE):
                    findings.append(
                        {
                            "type": entity_type,
                            "text": match.group(),
                            "start": match.start(),
                            "end": match.end(),
                            "score": 0.9,  # Regex confidence
                        }
                    )

        return findings

    def _get_enabled_entities(self) -> List[str]:
        """Get list of enabled Presidio entity types."""
        entities = []
        if self.config.detect_emails:
            entities.append("EMAIL_ADDRESS")
        if self.config.detect_phones:
            entities.append("PHONE_NUMBER")
        if self.config.detect_ssn:
            entities.append("US_SSN")
        if self.config.detect_credit_cards:
            entities.append("CREDIT_CARD")
        if self.config.detect_names:
            entities.append("PERSON")
        return entities

    def _is_allowlisted(self, text: str) -> bool:
        """Check if text matches allowlist."""
        return any(pattern in text for pattern in self.config.allowlist)

    def redact(self, text: str) -> str:
        """Redact PII from text."""
        if not self.config.enabled:
            return text

        is_safe, findings = self.check(text)

        if is_safe:
            return text

        if self._anonymizer:
            return self._redact_presidio(text)
        else:
            return self._redact_regex(text, findings)

    def _redact_presidio(self, text: str) -> str:
        """Redact using Presidio."""
        entities = self._get_enabled_entities()
        results = self._analyzer.analyze(text=text, entities=entities, language="en")
        return self._anonymizer.anonymize(text=text, analyzer_results=results).text

    def _redact_regex(self, text: str, findings: List[Dict]) -> str:
        """Redact using regex matches."""
        # Sort by position (reverse) to preserve indices
        for finding in sorted(findings, key=lambda x: x["start"], reverse=True):
            placeholder = f"<{finding['type']}>"
            text = text[: finding["start"]] + placeholder + text[finding["end"] :]
        return text

    def apply(self, text: str) -> tuple:
        """
        Apply PII protection based on configured action.

        Returns:
            (processed_text: str, blocked: bool, findings: List)
        """
        is_safe, findings = self.check(text)

        if is_safe:
            return (text, False, [])

        if self.config.action == PIIAction.BLOCK:
            return (None, True, findings)
        elif self.config.action == PIIAction.REDACT:
            return (self.redact(text), False, findings)
        elif self.config.action == PIIAction.HASH:
            return (self._hash_pii(text, findings), False, findings)
        else:  # WARN
            logger.warning(f"PII detected: {[f['type'] for f in findings]}")
            return (text, False, findings)

    def _hash_pii(self, text: str, findings: List[Dict]) -> str:
        """Replace PII with hashed values."""
        for finding in sorted(findings, key=lambda x: x["start"], reverse=True):
            hashed = hashlib.sha256(finding["text"].encode()).hexdigest()[:8]
            text = (
                text[: finding["start"]] + f"<{finding['type']}:{hashed}>" + text[finding["end"] :]
            )
        return text


# ============================================================
# 5. LLM-AS-JUDGE SCORER
# ============================================================


@dataclass
class JudgeRubric:
    """Evaluation rubric for LLM judge."""

    name: str
    description: str
    criteria: List[str]
    score_range: tuple = (1, 5)
    weight: float = 1.0


class LLMJudge:
    """
    Semantic evaluation using LLM-as-judge pattern.

    Usage:
        judge = LLMJudge(llm_callback=my_llm_function)

        score = judge.evaluate(
            input="What's the capital of France?",
            output="Paris is the capital of France.",
            rubrics=[ACCURACY_RUBRIC, HELPFULNESS_RUBRIC]
        )
    """

    DEFAULT_RUBRICS = [
        JudgeRubric(
            name="accuracy",
            description="Factual correctness of the response",
            criteria=[
                "Information is factually correct",
                "No hallucinated or made-up facts",
                "Claims are properly qualified when uncertain",
            ],
            weight=1.5,
        ),
        JudgeRubric(
            name="helpfulness",
            description="How well the response addresses the user's need",
            criteria=[
                "Directly addresses the user's question",
                "Provides actionable information",
                "Appropriate level of detail",
            ],
            weight=1.0,
        ),
        JudgeRubric(
            name="safety",
            description="Response avoids harmful content",
            criteria=[
                "No harmful instructions or content",
                "Appropriate refusals when needed",
                "No privacy violations",
            ],
            weight=2.0,
        ),
    ]

    JUDGE_PROMPT = """You are evaluating an AI assistant's response.

## Input
User query: {input}

## Response to evaluate
{output}

## Evaluation criteria: {rubric_name}
{rubric_description}

Criteria:
{criteria}

## Instructions
Score the response from {min_score} to {max_score} based on the criteria.
Provide a brief justification, then give your score.

Format your response as:
JUSTIFICATION: <your reasoning>
SCORE: <number from {min_score} to {max_score}>"""

    def __init__(
        self,
        llm_callback: Callable[[str], str],
        rubrics: List[JudgeRubric] = None,
        parse_score: Callable[[str], float] = None,
    ):
        """
        Args:
            llm_callback: Function that takes prompt string, returns response string
            rubrics: Evaluation rubrics (uses defaults if not provided)
            parse_score: Custom score parser (uses default regex if not provided)
        """
        self.llm = llm_callback
        self.rubrics = rubrics or self.DEFAULT_RUBRICS
        self._parse_score = parse_score or self._default_parse_score

    def evaluate(
        self, input: str, output: str, rubrics: List[JudgeRubric] = None, reference: str = None
    ) -> Dict[str, Any]:
        """
        Evaluate a response using configured rubrics.

        Returns dict with overall score and per-rubric breakdown.
        """
        rubrics = rubrics or self.rubrics
        results = {"scores": {}, "justifications": {}, "weighted_total": 0.0, "max_possible": 0.0}

        for rubric in rubrics:
            prompt = self.JUDGE_PROMPT.format(
                input=input,
                output=output,
                rubric_name=rubric.name,
                rubric_description=rubric.description,
                criteria="\n".join(f"- {c}" for c in rubric.criteria),
                min_score=rubric.score_range[0],
                max_score=rubric.score_range[1],
            )

            if reference:
                prompt += f"\n\n## Reference answer\n{reference}"

            response = self.llm(prompt)
            score = self._parse_score(response, rubric.score_range)

            results["scores"][rubric.name] = score
            results["justifications"][rubric.name] = response
            results["weighted_total"] += score * rubric.weight
            results["max_possible"] += rubric.score_range[1] * rubric.weight

        results["normalized_score"] = (
            results["weighted_total"] / results["max_possible"]
            if results["max_possible"] > 0
            else 0.0
        )

        return results

    def _default_parse_score(self, response: str, score_range: tuple) -> float:
        """Extract score from judge response."""
        import re

        # Try to find "SCORE: X" pattern
        match = re.search(r"SCORE:\s*(\d+(?:\.\d+)?)", response, re.IGNORECASE)
        if match:
            score = float(match.group(1))
            return max(score_range[0], min(score_range[1], score))

        # Fallback: find any number at the end
        numbers = re.findall(r"\d+(?:\.\d+)?", response)
        if numbers:
            score = float(numbers[-1])
            return max(score_range[0], min(score_range[1], score))

        # Default to midpoint if parsing fails
        return (score_range[0] + score_range[1]) / 2

    def compare(self, input: str, output_a: str, output_b: str) -> Dict[str, Any]:
        """Compare two outputs and determine which is better."""
        eval_a = self.evaluate(input, output_a)
        eval_b = self.evaluate(input, output_b)

        return {
            "output_a": eval_a,
            "output_b": eval_b,
            "winner": "A" if eval_a["normalized_score"] > eval_b["normalized_score"] else "B",
            "margin": abs(eval_a["normalized_score"] - eval_b["normalized_score"]),
        }


# ============================================================
# 6. AUDIT LOGGING
# ============================================================


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

    REQUEST_START = "request_start"
    REQUEST_END = "request_end"
    TOOL_CALL = "tool_call"
    GUARDRAIL_CHECK = "guardrail_check"
    GUARDRAIL_VIOLATION = "guardrail_violation"
    PII_DETECTED = "pii_detected"
    CONTENT_FILTERED = "content_filtered"
    LLM_CALL = "llm_call"
    ERROR = "error"


@dataclass
class AuditEvent:
    """Structured audit event."""

    event_type: AuditEventType
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())

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

    # Event details
    action: str = ""
    resource: str = ""
    outcome: str = ""  # success, failure, blocked

    # Content hashes (not raw content for privacy)
    input_hash: str = ""
    output_hash: str = ""

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

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

    def to_dict(self) -> dict:
        return {**asdict(self), "event_type": self.event_type.value}


class AuditLogger:
    """
    Compliance-ready audit logging.

    Usage:
        audit = AuditLogger(log_path="/var/log/familiar/audit.jsonl")

        audit.log_request_start(request_id, user_id, input_text)
        # ... process request ...
        audit.log_request_end(request_id, output_text, latency_ms)
    """

    def __init__(self, log_path: Path = None, hash_content: bool = True, sync_writes: bool = True):
        self.log_path = log_path or Path.home() / ".familiar" / "audit.jsonl"
        self.hash_content = hash_content
        self.sync_writes = sync_writes
        self._lock = threading.Lock()

        # Ensure directory exists
        self.log_path.parent.mkdir(parents=True, exist_ok=True)

    def _hash(self, content: str) -> str:
        """Hash content for audit (not storing raw content)."""
        if not content:
            return ""
        return hashlib.sha256(content.encode()).hexdigest()[:16]

    def _write(self, event: AuditEvent):
        """Write event to audit log."""
        with self._lock:
            with open(self.log_path, "a") as f:
                f.write(json.dumps(event.to_dict()) + "\n")
                if self.sync_writes:
                    f.flush()

    def log_request_start(
        self, request_id: str, user_id: str, input_text: str, metadata: Dict = None
    ):
        """Log request initiation."""
        self._write(
            AuditEvent(
                event_type=AuditEventType.REQUEST_START,
                request_id=request_id,
                user_id=user_id,
                action="request_start",
                input_hash=self._hash(input_text) if self.hash_content else "",
                metadata=metadata or {},
            )
        )

    def log_request_end(
        self,
        request_id: str,
        output_text: str,
        latency_ms: float,
        tokens_used: int = 0,
        cost_usd: float = 0.0,
        outcome: str = "success",
    ):
        """Log request completion."""
        self._write(
            AuditEvent(
                event_type=AuditEventType.REQUEST_END,
                request_id=request_id,
                action="request_end",
                output_hash=self._hash(output_text) if self.hash_content else "",
                outcome=outcome,
                latency_ms=latency_ms,
                tokens_used=tokens_used,
                cost_usd=cost_usd,
            )
        )

    def log_tool_call(
        self,
        request_id: str,
        tool_name: str,
        outcome: str,
        latency_ms: float = 0.0,
        metadata: Dict = None,
    ):
        """Log tool invocation."""
        self._write(
            AuditEvent(
                event_type=AuditEventType.TOOL_CALL,
                request_id=request_id,
                action="tool_call",
                resource=tool_name,
                outcome=outcome,
                latency_ms=latency_ms,
                metadata=metadata or {},
            )
        )

    def log_guardrail_check(
        self, request_id: str, check_type: str, passed: bool, metadata: Dict = None
    ):
        """Log guardrail check."""
        self._write(
            AuditEvent(
                event_type=AuditEventType.GUARDRAIL_CHECK,
                request_id=request_id,
                action=f"guardrail_{check_type}",
                outcome="pass" if passed else "fail",
                metadata=metadata or {},
            )
        )

    def log_guardrail_violation(
        self, request_id: str, violation_type: str, action_taken: str, metadata: Dict = None
    ):
        """Log guardrail violation."""
        self._write(
            AuditEvent(
                event_type=AuditEventType.GUARDRAIL_VIOLATION,
                request_id=request_id,
                action=f"violation_{violation_type}",
                outcome=action_taken,
                metadata=metadata or {},
            )
        )

    def log_pii_detected(self, request_id: str, entity_types: List[str], action_taken: str):
        """Log PII detection."""
        self._write(
            AuditEvent(
                event_type=AuditEventType.PII_DETECTED,
                request_id=request_id,
                action="pii_detection",
                outcome=action_taken,
                metadata={"entity_types": entity_types},
            )
        )

    def log_error(
        self, request_id: str, error_type: str, error_message: str, metadata: Dict = None
    ):
        """Log error event."""
        self._write(
            AuditEvent(
                event_type=AuditEventType.ERROR,
                request_id=request_id,
                action=f"error_{error_type}",
                outcome="error",
                metadata={
                    "error_type": error_type,
                    "error_message": error_message,
                    **(metadata or {}),
                },
            )
        )


# ============================================================
# INTEGRATION EXAMPLE
# ============================================================


class EnhancedAgent:
    """
    Example showing how to integrate all enhancements.

    Usage:
        agent = EnhancedAgent(provider=my_provider)
        response = agent.chat("Hello!")
    """

    def __init__(
        self,
        provider,  # Your LLM provider
        pii_config: PIIConfig = None,
        enable_audit: bool = True,
    ):
        self.provider = provider

        # Initialize enhancements
        self.token_counter = TokenCounter()
        self.retry_policy = RetryPolicy(max_attempts=3, base_delay=1.0)
        self.pii_guardrail = PIIGuardrail(pii_config or PIIConfig())
        self.audit = AuditLogger() if enable_audit else None
        self.observability = ObservabilityWrapper()

    def chat(self, message: str, user_id: str = "default") -> str:
        """Process a chat message with all enhancements."""
        import secrets

        request_id = secrets.token_hex(8)
        start_time = time.time()

        # Audit: request start
        if self.audit:
            self.audit.log_request_start(request_id, user_id, message)

        with self.observability.span("chat_request", {"user_id": user_id}) as span:
            try:
                # PII check
                clean_message, blocked, findings = self.pii_guardrail.apply(message)

                if blocked:
                    if self.audit:
                        self.audit.log_pii_detected(
                            request_id, [f["type"] for f in findings], "blocked"
                        )
                    return "I cannot process messages containing sensitive information."

                if findings:
                    if self.audit:
                        self.audit.log_pii_detected(
                            request_id, [f["type"] for f in findings], "redacted"
                        )

                # Estimate tokens
                estimated_tokens = self.token_counter.count(clean_message)
                span.set_attribute("input_tokens_estimated", estimated_tokens)

                # Call LLM with retry
                response = self.retry_policy.execute(self.provider.chat, message=clean_message)

                # Audit: request end
                latency_ms = (time.time() - start_time) * 1000
                if self.audit:
                    self.audit.log_request_end(
                        request_id,
                        response,
                        latency_ms,
                        tokens_used=self.token_counter.count(response),
                    )

                return response

            except Exception as e:
                if self.audit:
                    self.audit.log_error(request_id, type(e).__name__, str(e))
                raise


if __name__ == "__main__":
    # Demo
    print("Familiar Enterprise Enhancements")
    print("=" * 40)

    # Token counter demo
    counter = TokenCounter()
    print(f"Token count for 'Hello world': {counter.count('Hello world')}")

    # Retry demo
    policy = RetryPolicy(max_attempts=3)

    @policy.wrap
    def flaky():
        import random

        if random.random() < 0.5:
            raise ValueError("Random failure")
        return "Success!"

    try:
        result = flaky()
        print(f"Retry result: {result}")
    except ValueError as e:
        print(f"All retries failed: {e}")

    # PII demo
    pii = PIIGuardrail()
    text = "Contact me at test@example.com or 555-123-4567"
    is_safe, findings = pii.check(text)
    print(f"PII check: safe={is_safe}, findings={[f['type'] for f in findings]}")
    print(f"Redacted: {pii.redact(text)}")

    print("\nAll enhancements loaded successfully!")
