"""
Execution Context - Rich Context for AI Agent Actions.

Provides:
- Execution context wrapping
- Context propagation
- Trace correlation
- Environment isolation
"""
import contextvars
import uuid
import time
from dataclasses import dataclass, field
from typing import Any, Optional, Callable, TypeVar
from datetime import datetime, timezone
from functools import wraps
from enum import Enum

T = TypeVar("T")


class ExecutionMode(str, Enum):
    """Execution modes for agent actions."""
    ENFORCE = "enforce"  # Full policy enforcement, block violations
    MONITOR = "monitor"  # Log violations but don't block
    SIMULATE = "simulate"  # Dry-run, no actual execution


class ContextLevel(str, Enum):
    """Context nesting levels."""
    ROOT = "root"
    CHILD = "child"
    NESTED = "nested"


@dataclass
class ExecutionTrace:
    """Trace information for an execution."""
    trace_id: str
    span_id: str
    parent_span_id: Optional[str] = None
    operation_name: str = ""
    start_time: float = 0.0
    end_time: Optional[float] = None
    duration_ms: Optional[float] = None
    status: str = "pending"
    tags: dict = field(default_factory=dict)
    logs: list[dict] = field(default_factory=list)
    
    def finish(self, status: str = "ok"):
        """Mark trace as finished."""
        self.end_time = time.time()
        self.duration_ms = (self.end_time - self.start_time) * 1000
        self.status = status
    
    def add_log(self, message: str, level: str = "info"):
        """Add a log entry to the trace."""
        self.logs.append({
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "level": level,
            "message": message,
        })
    
    def to_dict(self) -> dict:
        return {
            "trace_id": self.trace_id,
            "span_id": self.span_id,
            "parent_span_id": self.parent_span_id,
            "operation_name": self.operation_name,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": self.duration_ms,
            "status": self.status,
            "tags": self.tags,
            "logs": self.logs,
        }


@dataclass
class ExecutionContext:
    """
    Context for agent execution.
    
    Carries all contextual information needed for:
    - Policy evaluation
    - Audit logging
    - Trace correlation
    - Environment isolation
    """
    # Identity
    agent_id: str
    session_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    
    # Execution settings
    mode: ExecutionMode = ExecutionMode.ENFORCE
    environment: str = "production"
    
    # Tracing
    trace_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    span_id: str = field(default_factory=lambda: str(uuid.uuid4())[:16])
    parent_span_id: Optional[str] = None
    
    # Context level
    level: ContextLevel = ContextLevel.ROOT
    depth: int = 0
    
    # Metadata
    user_id: Optional[str] = None
    tenant_id: Optional[str] = None
    correlation_id: Optional[str] = None
    
    # Timing
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    timeout_seconds: Optional[float] = None
    deadline: Optional[datetime] = None
    
    # Custom context
    attributes: dict = field(default_factory=dict)
    baggage: dict = field(default_factory=dict)  # Propagated across boundaries
    
    # Internal state
    _traces: list = field(default_factory=list)
    _active_trace: Optional[ExecutionTrace] = None
    
    def is_dry_run(self) -> bool:
        """Check if this is a dry-run (simulate mode)."""
        return self.mode == ExecutionMode.SIMULATE
    
    def should_enforce(self) -> bool:
        """Check if policy enforcement is active."""
        return self.mode == ExecutionMode.ENFORCE
    
    def should_block_violations(self) -> bool:
        """Check if violations should block execution."""
        return self.mode == ExecutionMode.ENFORCE
    
    def is_expired(self) -> bool:
        """Check if context deadline has passed."""
        if not self.deadline:
            return False
        return datetime.now(timezone.utc) > self.deadline
    
    def remaining_time_seconds(self) -> Optional[float]:
        """Get remaining time before deadline."""
        if not self.deadline:
            return None
        delta = self.deadline - datetime.now(timezone.utc)
        return max(0, delta.total_seconds())
    
    def start_span(self, operation_name: str, tags: Optional[dict] = None) -> ExecutionTrace:
        """Start a new trace span."""
        trace = ExecutionTrace(
            trace_id=self.trace_id,
            span_id=str(uuid.uuid4())[:16],
            parent_span_id=self._active_trace.span_id if self._active_trace else self.span_id,
            operation_name=operation_name,
            start_time=time.time(),
            tags=tags or {},
        )
        self._traces.append(trace)
        self._active_trace = trace
        return trace
    
    def end_span(self, status: str = "ok"):
        """End the current trace span."""
        if self._active_trace:
            self._active_trace.finish(status)
            self._active_trace = None
    
    def child_context(self, operation: str = "") -> "ExecutionContext":
        """Create a child context for nested operations."""
        return ExecutionContext(
            agent_id=self.agent_id,
            session_id=self.session_id,
            mode=self.mode,
            environment=self.environment,
            trace_id=self.trace_id,
            span_id=str(uuid.uuid4())[:16],
            parent_span_id=self.span_id,
            level=ContextLevel.CHILD if self.level == ContextLevel.ROOT else ContextLevel.NESTED,
            depth=self.depth + 1,
            user_id=self.user_id,
            tenant_id=self.tenant_id,
            correlation_id=self.correlation_id,
            deadline=self.deadline,
            attributes=dict(self.attributes),
            baggage=dict(self.baggage),
        )
    
    def with_mode(self, mode: ExecutionMode) -> "ExecutionContext":
        """Create a new context with different execution mode."""
        ctx = self.child_context()
        ctx.mode = mode
        return ctx
    
    def set_attribute(self, key: str, value: Any):
        """Set a context attribute."""
        self.attributes[key] = value
    
    def get_attribute(self, key: str, default: Any = None) -> Any:
        """Get a context attribute."""
        return self.attributes.get(key, default)
    
    def set_baggage(self, key: str, value: str):
        """Set baggage (propagated across service boundaries)."""
        self.baggage[key] = value
    
    def get_baggage(self, key: str, default: str = "") -> str:
        """Get baggage value."""
        return self.baggage.get(key, default)
    
    def to_headers(self) -> dict:
        """Export context to HTTP headers for propagation."""
        headers = {
            "X-Vaikora-Trace-ID": self.trace_id,
            "X-Vaikora-Span-ID": self.span_id,
            "X-Vaikora-Agent-ID": self.agent_id,
            "X-Vaikora-Session-ID": self.session_id,
            "X-Vaikora-Mode": self.mode.value,
        }
        
        if self.parent_span_id:
            headers["X-Vaikora-Parent-Span-ID"] = self.parent_span_id
        if self.correlation_id:
            headers["X-Vaikora-Correlation-ID"] = self.correlation_id
        if self.tenant_id:
            headers["X-Vaikora-Tenant-ID"] = self.tenant_id
        
        # Add baggage
        for key, value in self.baggage.items():
            headers[f"X-Vaikora-Baggage-{key}"] = value
        
        return headers
    
    @classmethod
    def from_headers(cls, headers: dict, agent_id: str) -> "ExecutionContext":
        """Create context from HTTP headers."""
        mode_str = headers.get("X-Vaikora-Mode", "enforce")
        try:
            mode = ExecutionMode(mode_str)
        except ValueError:
            mode = ExecutionMode.ENFORCE
        
        ctx = cls(
            agent_id=agent_id,
            session_id=headers.get("X-Vaikora-Session-ID", str(uuid.uuid4())),
            mode=mode,
            trace_id=headers.get("X-Vaikora-Trace-ID", str(uuid.uuid4())),
            span_id=str(uuid.uuid4())[:16],
            parent_span_id=headers.get("X-Vaikora-Span-ID"),
            correlation_id=headers.get("X-Vaikora-Correlation-ID"),
            tenant_id=headers.get("X-Vaikora-Tenant-ID"),
            level=ContextLevel.CHILD,
        )
        
        # Extract baggage
        for key, value in headers.items():
            if key.startswith("X-Vaikora-Baggage-"):
                baggage_key = key[18:]  # Remove prefix
                ctx.baggage[baggage_key] = value
        
        return ctx
    
    def to_dict(self) -> dict:
        """Export context to dictionary."""
        return {
            "agent_id": self.agent_id,
            "session_id": self.session_id,
            "mode": self.mode.value,
            "environment": self.environment,
            "trace_id": self.trace_id,
            "span_id": self.span_id,
            "parent_span_id": self.parent_span_id,
            "level": self.level.value,
            "depth": self.depth,
            "user_id": self.user_id,
            "tenant_id": self.tenant_id,
            "correlation_id": self.correlation_id,
            "created_at": self.created_at.isoformat(),
            "attributes": self.attributes,
            "baggage": self.baggage,
            "trace_count": len(self._traces),
        }


# Context variable for current execution context
_current_context: contextvars.ContextVar[Optional[ExecutionContext]] = contextvars.ContextVar(
    "vaikora_context", default=None
)


def get_current_context() -> Optional[ExecutionContext]:
    """Get the current execution context."""
    return _current_context.get()


def set_current_context(ctx: ExecutionContext) -> contextvars.Token:
    """Set the current execution context."""
    return _current_context.set(ctx)


def reset_context(token: contextvars.Token):
    """Reset context to previous value."""
    _current_context.reset(token)


class context_scope:
    """Context manager for execution context scope."""
    
    def __init__(self, ctx: ExecutionContext):
        self.ctx = ctx
        self.token: Optional[contextvars.Token] = None
    
    def __enter__(self) -> ExecutionContext:
        self.token = set_current_context(self.ctx)
        return self.ctx
    
    def __exit__(self, *args):
        if self.token:
            reset_context(self.token)


def with_context(
    mode: ExecutionMode = ExecutionMode.ENFORCE,
    timeout_seconds: Optional[float] = None,
) -> Callable:
    """
    Decorator to run function with execution context.
    
    Usage:
        @with_context(mode=ExecutionMode.MONITOR)
        async def my_function():
            ctx = get_current_context()
            ...
    """
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @wraps(func)
        async def async_wrapper(*args, **kwargs) -> T:
            # Try to get existing context or create new one
            existing = get_current_context()
            if existing:
                ctx = existing.child_context()
            else:
                # Need agent_id from somewhere
                agent_id = kwargs.pop("_agent_id", "unknown")
                ctx = ExecutionContext(agent_id=agent_id)
            
            ctx.mode = mode
            if timeout_seconds:
                ctx.timeout_seconds = timeout_seconds
                ctx.deadline = datetime.now(timezone.utc) + timedelta(seconds=timeout_seconds)
            
            with context_scope(ctx):
                trace = ctx.start_span(func.__name__)
                try:
                    result = await func(*args, **kwargs)
                    ctx.end_span("ok")
                    return result
                except Exception as e:
                    ctx.end_span("error")
                    trace.add_log(str(e), "error")
                    raise
        
        @wraps(func)
        def sync_wrapper(*args, **kwargs) -> T:
            existing = get_current_context()
            if existing:
                ctx = existing.child_context()
            else:
                agent_id = kwargs.pop("_agent_id", "unknown")
                ctx = ExecutionContext(agent_id=agent_id)
            
            ctx.mode = mode
            if timeout_seconds:
                ctx.timeout_seconds = timeout_seconds
                ctx.deadline = datetime.now(timezone.utc) + timedelta(seconds=timeout_seconds)
            
            with context_scope(ctx):
                trace = ctx.start_span(func.__name__)
                try:
                    result = func(*args, **kwargs)
                    ctx.end_span("ok")
                    return result
                except Exception as e:
                    ctx.end_span("error")
                    trace.add_log(str(e), "error")
                    raise
        
        import asyncio
        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        return sync_wrapper
    
    return decorator


# Import timedelta for deadline calculations
from datetime import timedelta
