
"""
Agent Governance SDK - Guard/Wrapper

Provides high-level governance wrappers for agent execution.
This is the primary API end-users will interact with.
"""

import functools
import logging
from datetime import datetime
from typing import Any, Callable, Dict, Optional, TypeVar, Union
from uuid import UUID
from contextlib import contextmanager
import time

from .sdk_types import (
    GovernanceDecision,
    GovernanceHaltError,
    GovernanceResponse,
    SDKConfig,
    StepData,
    StepType,
    TraceMetadata,
    TraceStatus,
)
from .tracer import Tracer
from .utils import TraceLogger, get_logger, generate_step_id, calculate_cost_usd, redact_sensitive_data


logger = get_logger(__name__)
T = TypeVar("T")


class GovernanceGuard:
    """
    High-level guard for agent execution with governance.
    
    This is the main entry point for wrapping agent code with governance.
    """

    def __init__(self, config: SDKConfig):
        self.config = config
        self.tracer = Tracer(config)
        self._logger = TraceLogger(logger)

    def close(self) -> None:
        """Cleanup resources."""
        self.tracer.cleanup()

    @contextmanager
    def govern(self, metadata: Optional[TraceMetadata] = None):
        """
        Context manager for governing an entire agent execution.
        
        Example:
            guard = GovernanceGuard(config)
            with guard.govern(metadata) as trace_id:
                # Your agent code here
                result = some_agent_function()
        """
        with self.tracer.trace_context(metadata) as trace_id:
            yield trace_id

    def with_governance(
        self,
        step_type: StepType = StepType.OBSERVATION,
        model: Optional[str] = None,
        check_policy: bool = True,
        metadata: Optional[Dict[str, Any]] = None,
    ):
        """
        Decorator for wrapping functions with governance.
        
        Args:
            step_type: Type of step this function represents
            model: Model name (if LLM call)
            check_policy: Whether to check policy before execution
            metadata: Additional metadata
            
        Example:
            @guard.with_governance(step_type=StepType.LLM_CALL, model="gpt-4")
            def llm_call(prompt):
                return openai.ChatCompletion.create(...)
        """
        def decorator(func: Callable[..., T]) -> Callable[..., T]:
            @functools.wraps(func)
            def wrapper(*args, **kwargs) -> T:
                # Get trace_id from kwargs or context
                trace_id = kwargs.get("trace_id")
                
                if trace_id is None:
                    # If no trace_id, execute without governance
                    self._logger.warning(
                        f"No trace_id provided for {func.__name__}, "
                        "executing without governance."
                    )
                    return func(*args, **kwargs)

                # Check policy before execution if enabled
                if check_policy and self.config.enable_policy_check:
                    response = self.tracer.check_policy(
                        trace_id=trace_id,
                        step_type=step_type,
                        model=model,
                        metadata=metadata,
                    )
                    
                    if response.decision == GovernanceDecision.HALT:
                        raise GovernanceHaltError(
                            reason=response.reason or "Policy halted execution",
                            policy_id=response.policy_id,
                        )

                # Execute and capture
                start_time = time.time()
                result = None
                error = None
                
                try:
                    result = func(*args, **kwargs)
                    return result
                except Exception as e:
                    error = e
                    raise
                finally:
                    # Record step
                    latency_ms = int((time.time() - start_time) * 1000)
                    
                    # Try to extract useful info from result/error
                    prompt = None
                    response_str = None
                    token_usage = None
                    estimated_cost = None
                    
                    # Extract LLM-related data if available
                    if hasattr(result, "__dict__"):
                        # OpenAI-style response
                        if hasattr(result, "usage"):
                            usage = result.usage
                            token_usage = {
                                "prompt_tokens": usage.prompt_tokens,
                                "completion_tokens": usage.completion_tokens,
                                "total_tokens": usage.total_tokens,
                            }
                            estimated_cost = calculate_cost_usd(
                                token_usage["prompt_tokens"],
                                token_usage["completion_tokens"],
                                model or "unknown",
                            )
                        
                        if hasattr(result, "choices") and len(result.choices) > 0:
                            response_str = result.choices[0].message.content
                    
                    # Record the step
                    self.tracer.record_step(
                        trace_id=trace_id,
                        step_type=step_type,
                        model=model,
                        prompt=prompt,
                        response=response_str,
                        token_usage=token_usage,
                        latency_ms=latency_ms,
                        estimated_cost_usd=estimated_cost,
                        metadata={
                            **(metadata or {}),
                            "function_name": func.__name__,
                            "error": str(error) if error else None,
                        },
                    )

            return wrapper
        return decorator


def create_guard(
    backend_url: str,
    api_key: str,
    agent_id: str,
    environment: str = "production",
    **kwargs,
) -> GovernanceGuard:
    """
    Factory function to create a GovernanceGuard.
    
    Args:
        backend_url: Governance backend URL
        api_key: API key for authentication
        agent_id: Unique identifier for this agent
        environment: Environment name (dev, staging, production)
        **kwargs: Additional SDK config options
        
    Returns:
        Configured GovernanceGuard instance
    """
    from .sdk_types import SDKConfig
    
    config = SDKConfig(
        backend_url=backend_url,
        api_key=api_key,
        agent_id=agent_id,
        environment=environment,
        **kwargs,
    )
    
    return GovernanceGuard(config)


class withGovernance:
    """
    Context manager-based governance wrapper.
    
    Alternative API style for governing agent execution.
    
    Example:
        with withGovernance(config) as guard:
            with guard.govern() as trace_id:
                @guard.with_governance(step_type=StepType.LLM_CALL)
                def my_llm_call():
                    ...
    """
    
    def __init__(self, config: SDKConfig):
        self.guard = GovernanceGuard(config)
    
    def __enter__(self) -> GovernanceGuard:
        return self.guard
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.guard.close()
        return False
