"""
Starvex Guard - Pythonic SDK for AI Agent Protection

Provides an elegant, decorator-based API for protecting AI agents.

Example Usage:
    from starvex.guardrails import StarvexGuard
    from starvex.guardrails import BlockPII, TopicRestriction

    # Initialize
    guard = StarvexGuard(
        rules=[
            BlockPII(),
            TopicRestriction(allowed_topics=["support", "billing"], sensitivity=0.8)
        ]
    )

    # Use as a Decorator (Easiest)
    @guard.protect
    def my_agent_chat(user_message):
        return agent_response

    # OR Use explicitly (More control)
    def chat_api(request):
        if not guard.check_input(request.text).passed:
            return "Sorry, I can't discuss that."

        response = agent.run(request.text)

        if not guard.check_output(response).passed:
            return "Error: Safety violation."

        return response
"""

import asyncio
import functools
import logging
import time
import uuid
from typing import (
    Any,
    Callable,
    Dict,
    List,
    Optional,
    TypeVar,
    Union,
    Coroutine,
)
from dataclasses import dataclass, field

from .rules import GuardRule, RuleResult, RuleAction, default_rules

logger = logging.getLogger(__name__)

T = TypeVar("T")


@dataclass
class CheckResult:
    """Result of a guard check with detailed violation information"""

    passed: bool
    blocked_by: Optional[str] = None
    action: Optional[RuleAction] = None
    message: str = ""
    confidence: float = 1.0
    all_results: List[RuleResult] = field(default_factory=list)
    redacted_text: Optional[str] = None
    latency_ms: float = 0.0
    trace_id: str = field(default_factory=lambda: str(uuid.uuid4()))

    @property
    def violation_details(self) -> Dict[str, Any]:
        """Get structured violation details for programmatic access"""
        if self.passed:
            return {"status": "passed", "checks_run": len(self.all_results)}

        details: Dict[str, Any] = {
            "status": "blocked",
            "blocked_by": self.blocked_by,
            "confidence": self.confidence,
            "checks_run": len(self.all_results),
        }

        for r in self.all_results:
            if not r.passed:
                details["violation"] = {
                    "rule": r.rule_name,
                    "action": r.action.value if r.action else None,
                    "message": r.message,
                    "confidence": r.confidence,
                    "details": r.details,
                }
                break

        return details


@dataclass
class GuardedResponse:
    """Response from a guarded function"""

    status: str  # "success", "blocked", "flagged"
    response: Optional[Any] = None
    input_check: Optional[CheckResult] = None
    output_check: Optional[CheckResult] = None
    error: Optional[str] = None
    trace_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    latency_ms: float = 0.0


class StarvexGuard:
    """
    Starvex Guard - Pythonic SDK for AI Agent Protection.

    A clean, decorator-based interface for protecting AI agents
    with composable guardrail rules.

    Example:
        ```python
        from starvex.guardrails import StarvexGuard
        from starvex.guardrails import BlockPII, TopicRestriction

        guard = StarvexGuard(
            rules=[
                BlockPII(),
                TopicRestriction(allowed_topics=["support", "billing"])
            ]
        )

        @guard.protect
        def my_agent(message: str) -> str:
            return llm.generate(message)

        # Now my_agent is protected!
        result = my_agent("Hello, my SSN is 123-45-6789")
        # Returns error message instead of processing PII
        ```
    """

    def __init__(
        self,
        rules: Optional[List[GuardRule]] = None,
        api_key: Optional[str] = None,
        enable_tracing: bool = True,
        on_block: Optional[Callable[[CheckResult], str]] = None,
        on_flag: Optional[Callable[[CheckResult], None]] = None,
        log_level: str = "INFO",
        middleware: Optional[Any] = None,
        mode: str = "enforce",
        enable_semantic_cache: bool = False,
        cache_similarity_threshold: float = 0.95,
        cache_max_size: int = 5000,
        cache_ttl_seconds: float = 1800.0,
    ):
        """
        Initialize StarvexGuard.

        Args:
            rules: List of GuardRule instances to apply
            api_key: Optional Starvex API key for cloud logging
            enable_tracing: Whether to log events to Starvex dashboard
            on_block: Custom handler when content is blocked
            on_flag: Custom handler when content is flagged
            log_level: Logging level
            middleware: Optional GuardMiddleware for rate limiting/budget control
            mode: "enforce" (default) blocks violations, "shadow" logs but allows
            enable_semantic_cache: Cache guard verdicts for semantically similar inputs
            cache_similarity_threshold: Cosine similarity threshold for cache hits (0.0-1.0)
            cache_max_size: Maximum number of cached verdict entries
            cache_ttl_seconds: Time-to-live for cached verdicts (30 min default)
        """
        self.rules = rules if rules is not None else default_rules()
        self.api_key = api_key
        self.enable_tracing = enable_tracing
        self.on_block = on_block or self._default_block_handler
        self.on_flag = on_flag
        self._middleware = middleware
        self.mode = mode

        # Semantic caching
        self._enable_semantic_cache = enable_semantic_cache
        self._cache_similarity_threshold = cache_similarity_threshold
        self._cache_model = None
        self._cache_np = None
        self._verdict_cache: Optional[Any] = None
        self._cached_embeddings: List[Any] = []  # List of (embedding, CheckResult) tuples

        if enable_semantic_cache:
            self._init_semantic_cache(cache_max_size, cache_ttl_seconds)

        # Initialize tracer if API key provided
        self._tracer = None
        if api_key and enable_tracing:
            try:
                from ..observability import Tracer

                self._tracer = Tracer(api_key=api_key)
            except Exception as e:
                logger.warning(f"Failed to initialize tracer: {e}")

        # Configure logging
        logging.basicConfig(level=getattr(logging, log_level.upper()))

        logger.info(f"StarvexGuard initialized with {len(self.rules)} rules (mode={self.mode})")

    def _init_semantic_cache(self, max_size: int, ttl_seconds: float) -> None:
        """Initialize the semantic verdict cache."""
        try:
            from .._internals.cache import get_cache

            self._verdict_cache = get_cache(max_size=max_size, ttl_seconds=ttl_seconds)
            logger.info("Semantic verdict cache initialized")
        except Exception as e:
            logger.warning(f"Failed to initialize verdict cache: {e}")
            self._enable_semantic_cache = False

    def _get_embedding(self, text: str) -> Optional[Any]:
        """Compute embedding for text using sentence-transformers."""
        if self._cache_model is None:
            try:
                from sentence_transformers import SentenceTransformer
                import numpy as np

                self._cache_model = SentenceTransformer("all-MiniLM-L6-v2")
                self._cache_np = np
            except ImportError:
                logger.warning("sentence-transformers not installed, semantic caching disabled")
                self._enable_semantic_cache = False
                return None

        try:
            return self._cache_model.encode([text])[0]
        except Exception as e:
            logger.debug(f"Embedding computation error: {e}")
            return None

    def _check_verdict_cache(self, text: str) -> Optional[CheckResult]:
        """Check if a semantically similar input was already evaluated."""
        if not self._enable_semantic_cache or not self._verdict_cache:
            return None

        # First try exact-match cache
        cached = self._verdict_cache.get(text, model_name="verdict", prefix="guard_verdict")
        if cached is not None:
            logger.debug("Semantic cache hit (exact match)")
            cached_result: CheckResult = cached
            return CheckResult(
                passed=cached_result.passed,
                blocked_by=cached_result.blocked_by,
                action=cached_result.action,
                message=f"[CACHED] {cached_result.message}",
                confidence=cached_result.confidence,
                all_results=cached_result.all_results,
                redacted_text=cached_result.redacted_text,
                latency_ms=0.0,
            )

        # Then try semantic similarity against cached embeddings
        embedding = self._get_embedding(text)
        if embedding is None:
            return None

        np = self._cache_np
        for cached_emb, cached_result in self._cached_embeddings:
            try:
                sim = float(
                    np.dot(cached_emb, embedding)
                    / (np.linalg.norm(cached_emb) * np.linalg.norm(embedding))
                )
                if sim >= self._cache_similarity_threshold:
                    logger.debug(f"Semantic cache hit (similarity={sim:.3f})")
                    return CheckResult(
                        passed=cached_result.passed,
                        blocked_by=cached_result.blocked_by,
                        action=cached_result.action,
                        message=f"[CACHED] {cached_result.message}",
                        confidence=cached_result.confidence,
                        all_results=cached_result.all_results,
                        redacted_text=cached_result.redacted_text,
                        latency_ms=0.0,
                    )
            except Exception:
                continue

        return None

    def _store_verdict_cache(self, text: str, result: CheckResult) -> None:
        """Store a verdict result in the semantic cache."""
        if not self._enable_semantic_cache or not self._verdict_cache:
            return

        # Only cache block results to save memory
        if result.passed and not result.blocked_by:
            return

        # Exact-match cache
        self._verdict_cache.set(text, result, model_name="verdict", prefix="guard_verdict")

        # Store embedding for semantic similarity checks
        embedding = self._get_embedding(text)
        if embedding is not None:
            self._cached_embeddings.append((embedding, result))
            # Cap the embedding list size
            if len(self._cached_embeddings) > 5000:
                self._cached_embeddings = self._cached_embeddings[-2500:]

    def _default_block_handler(self, result: CheckResult) -> str:
        """Default message when content is blocked with actionable details"""
        return self._build_detailed_message(result)

    @staticmethod
    def _build_detailed_message(result: CheckResult) -> str:
        """Build an actionable error message with violation details and suggestions"""
        parts = []

        # Main violation message
        if result.blocked_by == "block_pii":
            entities = []
            suggestions = []
            for r in result.all_results:
                if r.rule_name == "block_pii" and not r.passed:
                    entities = r.details.get("entities", [])
                    break

            entity_str = ", ".join(entities) if entities else "personal information"
            parts.append(f"Blocked: PII detected ({entity_str})")

            # Add position info if available
            for r in result.all_results:
                if r.rule_name == "block_pii" and r.details.get("entities"):
                    break

            suggestions.append("Remove personal information before sending")
            suggestions.append("Use PII redaction mode: BlockPII(redact_instead=True)")
            parts.append(f"Suggestions: {'; '.join(suggestions)}")

        elif result.blocked_by == "block_jailbreak":
            pattern_info = ""
            for r in result.all_results:
                if r.rule_name == "block_jailbreak" and not r.passed:
                    pattern_info = r.details.get("pattern", "")
                    break
            parts.append("Blocked: Jailbreak/prompt injection attempt detected")
            if pattern_info:
                parts.append(f"Matched pattern: {pattern_info[:80]}")
            parts.append(
                "Suggestion: Rephrase your request without instruction overrides or role manipulation"
            )

        elif result.blocked_by == "topic_restriction":
            topic = ""
            for r in result.all_results:
                if r.rule_name == "topic_restriction" and not r.passed:
                    topic = r.message.replace("Topic not allowed: ", "")
                    break
            parts.append(f"Blocked: Topic '{topic}' is restricted")
            parts.append("Suggestion: Rephrase your question within allowed topic areas")

        elif result.blocked_by == "block_toxicity":
            parts.append("Blocked: Toxic or offensive content detected")
            parts.append("Suggestion: Rephrase your message in a respectful manner")

        elif result.blocked_by == "block_competitor":
            matched = ""
            for r in result.all_results:
                if r.rule_name == "block_competitor" and not r.passed:
                    matched = r.details.get("matched", "")
                    break
            parts.append(
                f"Blocked: Competitor mention detected{f' ({matched})' if matched else ''}"
            )
            parts.append("Suggestion: Rephrase without referencing competitor products")

        elif result.blocked_by == "policy_compliance":
            violations = []
            for r in result.all_results:
                if r.rule_name == "policy_compliance" and r.details.get("violations"):
                    violations = r.details["violations"]
                    break
            num_violations = len(violations)
            parts.append(f"Blocked: {num_violations} policy violation(s) detected")
            for v in violations[:3]:
                parts.append(f"  - Policy: {v.get('policy', 'unknown')}")
            parts.append(
                "Suggestion: Ensure your response aligns with configured business policies"
            )

        elif result.blocked_by == "custom_blocklist":
            parts.append("Blocked: Content matches custom blocklist")
            parts.append("Suggestion: Remove blocked phrases or patterns from your input")

        else:
            parts.append(f"Blocked by: {result.blocked_by or 'safety policy'}")
            parts.append("Suggestion: Review your input for policy violations")

        # Add confidence info
        if result.confidence < 1.0:
            parts.append(f"Confidence: {result.confidence:.0%}")

        return " | ".join(parts)

    def check_input(self, text: str, parallel: bool = False) -> CheckResult:
        """
        Check input text against all rules.

        Args:
            text: Input text to check
            parallel: If True, run non-redaction rules in parallel for faster execution

        Returns:
            CheckResult with pass/fail status and details
        """
        # Check semantic cache first
        cached_result = self._check_verdict_cache(text)
        if cached_result is not None:
            return cached_result

        if parallel:
            result = self._check_input_parallel(text)
        else:
            result = self._check_input_sequential(text)

        # Store in cache if blocked (or shadow-blocked)
        self._store_verdict_cache(text, result)
        return result

    def _check_input_sequential(self, text: str) -> CheckResult:
        """Sequential rule checking (preserves redaction order)"""
        start_time = time.time()
        all_results: List[RuleResult] = []
        redacted_text = None
        shadow_blocked_by: Optional[str] = None
        shadow_message: str = ""
        shadow_confidence: float = 1.0
        shadow_action: Optional[RuleAction] = None

        for rule in self.rules:
            try:
                result = rule.check(text)
                all_results.append(result)

                # Handle redaction
                if result.action == RuleAction.REDACT and result.redacted_text:
                    redacted_text = result.redacted_text
                    text = redacted_text  # Use redacted text for subsequent checks

                # Check if blocked
                if not result.passed and result.action == RuleAction.BLOCK:
                    if self.mode == "shadow":
                        # Shadow mode: log what WOULD be blocked, continue checking
                        logger.warning(
                            f"[SHADOW] Would block by {result.rule_name}: "
                            f"{result.message} (confidence={result.confidence:.2f})"
                        )
                        if shadow_blocked_by is None:
                            shadow_blocked_by = result.rule_name
                            shadow_message = result.message
                            shadow_confidence = result.confidence
                            shadow_action = result.action
                        continue
                    latency = (time.time() - start_time) * 1000
                    return CheckResult(
                        passed=False,
                        blocked_by=result.rule_name,
                        action=result.action,
                        message=result.message,
                        confidence=result.confidence,
                        all_results=all_results,
                        redacted_text=redacted_text,
                        latency_ms=latency,
                    )

            except Exception as e:
                logger.error(f"Error in rule {rule.name}: {e}")
                continue

        latency = (time.time() - start_time) * 1000

        # In shadow mode, return passed=True but include shadow block info
        if self.mode == "shadow" and shadow_blocked_by:
            return CheckResult(
                passed=True,
                blocked_by=shadow_blocked_by,
                action=shadow_action,
                message=f"[SHADOW] {shadow_message}",
                confidence=shadow_confidence,
                all_results=all_results,
                redacted_text=redacted_text,
                latency_ms=latency,
            )

        return CheckResult(
            passed=True,
            all_results=all_results,
            redacted_text=redacted_text,
            latency_ms=latency,
            message="All checks passed",
        )

    def _check_input_parallel(self, text: str) -> CheckResult:
        """
        Parallel rule checking using concurrent.futures for 3-5x latency reduction.

        Separates redaction rules (must run first, sequentially) from
        blocking rules (can run in parallel).
        """
        from concurrent.futures import ThreadPoolExecutor, as_completed

        start_time = time.time()
        all_results: List[RuleResult] = []
        redacted_text = None
        shadow_blocked_by: Optional[str] = None
        shadow_message: str = ""
        shadow_confidence: float = 1.0
        shadow_action: Optional[RuleAction] = None

        # Phase 1: Run redaction rules first (sequential, order-dependent)
        redaction_rules = [r for r in self.rules if r.action == RuleAction.REDACT]
        blocking_rules = [r for r in self.rules if r.action != RuleAction.REDACT]

        check_text = text
        for rule in redaction_rules:
            try:
                result = rule.check(check_text)
                all_results.append(result)
                if result.action == RuleAction.REDACT and result.redacted_text:
                    redacted_text = result.redacted_text
                    check_text = redacted_text
            except Exception as e:
                logger.error(f"Error in redaction rule {rule.name}: {e}")

        # Phase 2: Run blocking rules in parallel
        def _run_rule(rule: GuardRule) -> RuleResult:
            return rule.check(check_text)

        with ThreadPoolExecutor(max_workers=min(len(blocking_rules), 8)) as executor:
            future_to_rule = {executor.submit(_run_rule, rule): rule for rule in blocking_rules}

            for future in as_completed(future_to_rule):
                rule = future_to_rule[future]
                try:
                    result = future.result()
                    all_results.append(result)

                    if not result.passed and result.action == RuleAction.BLOCK:
                        if self.mode == "shadow":
                            # Shadow mode: log but don't block
                            logger.warning(
                                f"[SHADOW] Would block by {result.rule_name}: "
                                f"{result.message} (confidence={result.confidence:.2f})"
                            )
                            if shadow_blocked_by is None:
                                shadow_blocked_by = result.rule_name
                                shadow_message = result.message
                                shadow_confidence = result.confidence
                                shadow_action = result.action
                            continue
                        # Cancel remaining futures on first block
                        for f in future_to_rule:
                            f.cancel()
                        latency = (time.time() - start_time) * 1000
                        return CheckResult(
                            passed=False,
                            blocked_by=result.rule_name,
                            action=result.action,
                            message=result.message,
                            confidence=result.confidence,
                            all_results=all_results,
                            redacted_text=redacted_text,
                            latency_ms=latency,
                        )
                except Exception as e:
                    logger.error(f"Error in rule {rule.name}: {e}")

        latency = (time.time() - start_time) * 1000

        # In shadow mode, return passed=True but include shadow block info
        if self.mode == "shadow" and shadow_blocked_by:
            return CheckResult(
                passed=True,
                blocked_by=shadow_blocked_by,
                action=shadow_action,
                message=f"[SHADOW] {shadow_message}",
                confidence=shadow_confidence,
                all_results=all_results,
                redacted_text=redacted_text,
                latency_ms=latency,
            )

        return CheckResult(
            passed=True,
            all_results=all_results,
            redacted_text=redacted_text,
            latency_ms=latency,
            message="All checks passed",
        )

    def check_output(
        self,
        output_text: str,
        input_text: Optional[str] = None,
        parallel: bool = False,
    ) -> CheckResult:
        """
        Check output text against all rules.

        Args:
            output_text: Agent output to check
            input_text: Original input (for context-aware checks)
            parallel: If True, run rules in parallel for faster execution

        Returns:
            CheckResult with pass/fail status and details
        """
        if parallel:
            return self._check_output_parallel(output_text, input_text)
        return self._check_output_sequential(output_text, input_text)

    def _check_output_sequential(
        self,
        output_text: str,
        input_text: Optional[str] = None,
    ) -> CheckResult:
        """Sequential output checking"""
        start_time = time.time()
        all_results: List[RuleResult] = []
        redacted_text = None
        shadow_blocked_by: Optional[str] = None
        shadow_message: str = ""
        shadow_confidence: float = 1.0
        shadow_action: Optional[RuleAction] = None

        for rule in self.rules:
            try:
                if input_text:
                    result = rule.check_output(input_text, output_text)
                else:
                    result = rule.check(output_text)

                all_results.append(result)

                # Handle redaction
                if result.action == RuleAction.REDACT and result.redacted_text:
                    redacted_text = result.redacted_text

                # Check if blocked
                if not result.passed and result.action == RuleAction.BLOCK:
                    if self.mode == "shadow":
                        logger.warning(
                            f"[SHADOW] Output would block by {result.rule_name}: "
                            f"{result.message} (confidence={result.confidence:.2f})"
                        )
                        if shadow_blocked_by is None:
                            shadow_blocked_by = result.rule_name
                            shadow_message = result.message
                            shadow_confidence = result.confidence
                            shadow_action = result.action
                        continue
                    latency = (time.time() - start_time) * 1000
                    return CheckResult(
                        passed=False,
                        blocked_by=result.rule_name,
                        action=result.action,
                        message=result.message,
                        confidence=result.confidence,
                        all_results=all_results,
                        redacted_text=redacted_text,
                        latency_ms=latency,
                    )

            except Exception as e:
                logger.error(f"Error in rule {rule.name}: {e}")
                continue

        latency = (time.time() - start_time) * 1000

        if self.mode == "shadow" and shadow_blocked_by:
            return CheckResult(
                passed=True,
                blocked_by=shadow_blocked_by,
                action=shadow_action,
                message=f"[SHADOW] {shadow_message}",
                confidence=shadow_confidence,
                all_results=all_results,
                redacted_text=redacted_text,
                latency_ms=latency,
            )

        return CheckResult(
            passed=True,
            all_results=all_results,
            redacted_text=redacted_text,
            latency_ms=latency,
            message="All checks passed",
        )

    def _check_output_parallel(
        self,
        output_text: str,
        input_text: Optional[str] = None,
    ) -> CheckResult:
        """Parallel output checking using concurrent.futures"""
        from concurrent.futures import ThreadPoolExecutor, as_completed

        start_time = time.time()
        all_results: List[RuleResult] = []
        redacted_text = None
        shadow_blocked_by: Optional[str] = None
        shadow_message: str = ""
        shadow_confidence: float = 1.0
        shadow_action: Optional[RuleAction] = None

        def _run_rule(rule: GuardRule) -> RuleResult:
            if input_text:
                return rule.check_output(input_text, output_text)
            return rule.check(output_text)

        with ThreadPoolExecutor(max_workers=min(len(self.rules), 8)) as executor:
            future_to_rule = {executor.submit(_run_rule, rule): rule for rule in self.rules}

            for future in as_completed(future_to_rule):
                rule = future_to_rule[future]
                try:
                    result = future.result()
                    all_results.append(result)

                    if result.action == RuleAction.REDACT and result.redacted_text:
                        redacted_text = result.redacted_text

                    if not result.passed and result.action == RuleAction.BLOCK:
                        if self.mode == "shadow":
                            logger.warning(
                                f"[SHADOW] Output would block by {result.rule_name}: "
                                f"{result.message} "
                                f"(confidence={result.confidence:.2f})"
                            )
                            if shadow_blocked_by is None:
                                shadow_blocked_by = result.rule_name
                                shadow_message = result.message
                                shadow_confidence = result.confidence
                                shadow_action = result.action
                            continue
                        for f in future_to_rule:
                            f.cancel()
                        latency = (time.time() - start_time) * 1000
                        return CheckResult(
                            passed=False,
                            blocked_by=result.rule_name,
                            action=result.action,
                            message=result.message,
                            confidence=result.confidence,
                            all_results=all_results,
                            redacted_text=redacted_text,
                            latency_ms=latency,
                        )
                except Exception as e:
                    logger.error(f"Error in rule {rule.name}: {e}")

        latency = (time.time() - start_time) * 1000

        if self.mode == "shadow" and shadow_blocked_by:
            return CheckResult(
                passed=True,
                blocked_by=shadow_blocked_by,
                action=shadow_action,
                message=f"[SHADOW] {shadow_message}",
                confidence=shadow_confidence,
                all_results=all_results,
                redacted_text=redacted_text,
                latency_ms=latency,
            )

        return CheckResult(
            passed=True,
            all_results=all_results,
            redacted_text=redacted_text,
            latency_ms=latency,
            message="All checks passed",
        )

    def protect(
        self,
        func: Optional[Callable[..., T]] = None,
        *,
        check_input: bool = True,
        check_output: bool = True,
    ) -> Union[Callable[..., T], Callable[[Callable[..., T]], Callable[..., T]]]:
        """
        Decorator to protect a function with guardrails.

        Can be used with or without parentheses:

            @guard.protect
            def my_func(message):
                ...

            @guard.protect(check_output=False)
            def my_func(message):
                ...

        Args:
            func: Function to protect (when used without parens)
            check_input: Whether to check input
            check_output: Whether to check output

        Returns:
            Protected function or decorator
        """

        def decorator(fn: Callable[..., T]) -> Callable[..., T]:
            if asyncio.iscoroutinefunction(fn):

                @functools.wraps(fn)
                async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                    return await self._protected_call_async(
                        fn, args, kwargs, check_input, check_output
                    )

                return async_wrapper
            else:

                @functools.wraps(fn)
                def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
                    return self._protected_call_sync(fn, args, kwargs, check_input, check_output)

                return sync_wrapper

        if func is not None:
            return decorator(func)
        return decorator

    def _protected_call_sync(
        self,
        func: Callable,
        args: tuple,
        kwargs: dict,
        check_input: bool,
        check_output: bool,
    ) -> Any:
        """Execute protected sync function"""
        start_time = time.time()
        trace_id = str(uuid.uuid4())

        # Extract input text (first string argument)
        input_text = self._extract_input(args, kwargs)

        # Middleware rate/budget check
        if self._middleware:
            key = self.api_key or "local"
            allowed, reason = self._middleware.before_check(key)
            if not allowed:
                latency_ms = (time.time() - start_time) * 1000
                block_reason = reason or "Rate limit exceeded"
                blocked_by = "rate_limit" if "Rate limit" in block_reason else "budget_limit"
                result = CheckResult(
                    passed=False,
                    blocked_by=blocked_by,
                    action=RuleAction.BLOCK,
                    message=block_reason,
                    confidence=1.0,
                    latency_ms=latency_ms,
                )
                self._log_event(trace_id, input_text, None, "blocked", result)
                self._middleware.after_check(key, latency_ms=latency_ms, blocked=True)
                return self.on_block(result)

        # Check input
        if check_input and input_text:
            input_result = self.check_input(input_text)

            if not input_result.passed:
                self._log_event(trace_id, input_text, None, "blocked", input_result)
                if self._middleware:
                    key = self.api_key or "local"
                    self._middleware.after_check(
                        key, latency_ms=(time.time() - start_time) * 1000, blocked=True
                    )
                return self.on_block(input_result)

            # Log shadow blocks (passed=True but blocked_by is set)
            if self.mode == "shadow" and input_result.blocked_by:
                self._log_event(trace_id, input_text, None, "blocked_shadow", input_result)

            # Use redacted text if available
            if input_result.redacted_text:
                args, kwargs = self._replace_input(args, kwargs, input_result.redacted_text)

        # Execute function
        try:
            result = func(*args, **kwargs)
        except Exception as e:
            logger.error(f"Function execution error: {e}")
            if self._middleware:
                key = self.api_key or "local"
                self._middleware.after_check(
                    key, latency_ms=(time.time() - start_time) * 1000, blocked=False
                )
            raise

        # Check output
        if check_output and result:
            output_text = str(result) if not isinstance(result, str) else result
            output_result = self.check_output(output_text, input_text)

            if not output_result.passed:
                self._log_event(trace_id, input_text, output_text, "blocked", output_result)
                if self._middleware:
                    key = self.api_key or "local"
                    self._middleware.after_check(
                        key, latency_ms=(time.time() - start_time) * 1000, blocked=True
                    )
                return self.on_block(output_result)

            # Log shadow blocks on output
            if self.mode == "shadow" and output_result.blocked_by:
                self._log_event(trace_id, input_text, output_text, "blocked_shadow", output_result)

            # Use redacted output if available
            if output_result.redacted_text:
                result = output_result.redacted_text

        self._log_event(trace_id, input_text, str(result) if result else None, "success", None)
        if self._middleware:
            key = self.api_key or "local"
            self._middleware.after_check(
                key, latency_ms=(time.time() - start_time) * 1000, blocked=False
            )
        return result

    async def _protected_call_async(
        self,
        func: Callable[..., Coroutine],
        args: tuple,
        kwargs: dict,
        check_input: bool,
        check_output: bool,
    ) -> Any:
        """Execute protected async function"""
        start_time = time.time()
        trace_id = str(uuid.uuid4())

        # Extract input text
        input_text = self._extract_input(args, kwargs)

        # Middleware rate/budget check
        if self._middleware:
            key = self.api_key or "local"
            allowed, reason = self._middleware.before_check(key)
            if not allowed:
                latency_ms = (time.time() - start_time) * 1000
                block_reason = reason or "Rate limit exceeded"
                blocked_by = "rate_limit" if "Rate limit" in block_reason else "budget_limit"
                result = CheckResult(
                    passed=False,
                    blocked_by=blocked_by,
                    action=RuleAction.BLOCK,
                    message=block_reason,
                    confidence=1.0,
                    latency_ms=latency_ms,
                )
                self._log_event(trace_id, input_text, None, "blocked", result)
                self._middleware.after_check(key, latency_ms=latency_ms, blocked=True)
                return self.on_block(result)

        # Check input
        if check_input and input_text:
            input_result = self.check_input(input_text)

            if not input_result.passed:
                self._log_event(trace_id, input_text, None, "blocked", input_result)
                if self._middleware:
                    key = self.api_key or "local"
                    self._middleware.after_check(
                        key, latency_ms=(time.time() - start_time) * 1000, blocked=True
                    )
                return self.on_block(input_result)

            # Log shadow blocks
            if self.mode == "shadow" and input_result.blocked_by:
                self._log_event(trace_id, input_text, None, "blocked_shadow", input_result)

            if input_result.redacted_text:
                args, kwargs = self._replace_input(args, kwargs, input_result.redacted_text)

        # Execute function
        try:
            result = await func(*args, **kwargs)
        except Exception as e:
            logger.error(f"Function execution error: {e}")
            if self._middleware:
                key = self.api_key or "local"
                self._middleware.after_check(
                    key, latency_ms=(time.time() - start_time) * 1000, blocked=False
                )
            raise

        # Check output
        if check_output and result:
            output_text = str(result) if not isinstance(result, str) else result
            output_result = self.check_output(output_text, input_text)

            if not output_result.passed:
                self._log_event(trace_id, input_text, output_text, "blocked", output_result)
                if self._middleware:
                    key = self.api_key or "local"
                    self._middleware.after_check(
                        key, latency_ms=(time.time() - start_time) * 1000, blocked=True
                    )
                return self.on_block(output_result)

            # Log shadow blocks on output
            if self.mode == "shadow" and output_result.blocked_by:
                self._log_event(trace_id, input_text, output_text, "blocked_shadow", output_result)

            if output_result.redacted_text:
                result = output_result.redacted_text

        self._log_event(trace_id, input_text, str(result) if result else None, "success", None)
        if self._middleware:
            key = self.api_key or "local"
            self._middleware.after_check(
                key, latency_ms=(time.time() - start_time) * 1000, blocked=False
            )
        return result

    def _extract_input(self, args: tuple, kwargs: dict) -> Optional[str]:
        """Extract input text from function arguments"""
        # Try first positional arg
        if args and isinstance(args[0], str):
            return args[0]

        # Try common kwarg names
        for key in ["message", "text", "input", "prompt", "query", "content"]:
            if key in kwargs and isinstance(kwargs[key], str):
                return kwargs[key]

        return None

    def _replace_input(
        self,
        args: tuple,
        kwargs: dict,
        new_text: str,
    ) -> tuple:
        """Replace input text in function arguments"""
        if args and isinstance(args[0], str):
            return (new_text,) + args[1:], kwargs

        for key in ["message", "text", "input", "prompt", "query", "content"]:
            if key in kwargs and isinstance(kwargs[key], str):
                kwargs = {**kwargs, key: new_text}
                return args, kwargs

        return args, kwargs

    def _log_event(
        self,
        trace_id: str,
        input_text: Optional[str],
        output_text: Optional[str],
        status: str,
        result: Optional[CheckResult],
    ):
        """Log event to tracer if enabled"""
        if self._tracer:
            try:
                self._tracer.log_event(
                    trace_id=trace_id,
                    input_text=input_text or "",
                    output_text=output_text or "",
                    verdict=status,
                    confidence_score=result.confidence if result else 1.0,
                    latency_ms=result.latency_ms if result else 0,
                )
            except Exception as e:
                logger.debug(f"Failed to log event: {e}")

    def add_rule(self, rule: GuardRule) -> None:
        """Add a rule to the guard"""
        self.rules.append(rule)
        logger.info(f"Added rule: {rule.name}")

    def remove_rule(self, rule_name: str) -> bool:
        """Remove a rule by name"""
        original_count = len(self.rules)
        self.rules = [r for r in self.rules if r.name != rule_name]
        return len(self.rules) < original_count

    def get_rules(self) -> List[str]:
        """Get list of active rule names"""
        return [r.name for r in self.rules]

    # Context manager support
    def __enter__(self) -> "StarvexGuard":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        if self._tracer:
            self._tracer.flush()

    async def __aenter__(self) -> "StarvexGuard":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        if self._tracer:
            self._tracer.flush()


# Convenience alias
Guard = StarvexGuard
