"""
Execution context management for orchestration patterns.

This module provides context storage for pattern execution, enabling
sharing of data, configuration, and state across components.
"""

from typing import Any, Dict, Optional
from pydantic import BaseModel, Field
from datetime import datetime


class Context(BaseModel):
    """
    Execution context for pattern and agent operations.

    This Pydantic model provides a simple key-value store for sharing
    data, configuration, and temporary state during execution.

    :param data: Dict[str, Any] - Context data storage.
    :param created_at: datetime - Timestamp when context was created.
    :param updated_at: datetime - Timestamp of last update.
    """

    data: Dict[str, Any] = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)

    def set(self, key: str, value: Any) -> None:
        """
        Set a context value and update timestamp.

        :param key: str - Context key.
        :param value: Any - Value to store.
        """
        self.data[key] = value
        self.updated_at = datetime.utcnow()

    def get(self, key: str, default: Any = None) -> Any:
        """
        Get a context value with optional default.

        :param key: str - Context key.
        :param default: Any - Default value if key not found.
        :return: Any - Stored value or default.
        """
        return self.data.get(key, default)

    def has(self, key: str) -> bool:
        """
        Check if a key exists in context.

        :param key: str - Context key.
        :return: bool - True if key exists, False otherwise.
        """
        return key in self.data

    def remove(self, key: str) -> None:
        """
        Remove a key from context and update timestamp.

        :param key: str - Context key to remove.
        """
        if key in self.data:
            del self.data[key]
            self.updated_at = datetime.utcnow()

    def update(self, data: Dict[str, Any]) -> None:
        """
        Merge data into context and update timestamp.

        :param data: Dict[str, Any] - Data to merge into context.
        """
        self.data.update(data)
        self.updated_at = datetime.utcnow()

    def clear(self) -> None:
        """Clear all context data and update timestamp."""
        self.data.clear()
        self.updated_at = datetime.utcnow()

    class Config:
        json_encoders = {datetime: lambda v: v.isoformat()}


class ContextManager:
    """
    Manager for multiple execution contexts.

    Provides centralized management of multiple context instances,
    enabling isolated execution environments for different patterns or agents.
    """

    def __init__(self):
        """Initialize the context manager with empty context registry."""
        self._contexts: Dict[str, Context] = {}

    def create_context(
        self, context_id: str, initial_data: Optional[Dict[str, Any]] = None
    ) -> Context:
        """
        Create and register a new context.

        :param context_id: str - Unique identifier for the context.
        :param initial_data: Optional[Dict[str, Any]] - Initial data for the context.
        :return: Context - The newly created context instance.
        """
        context = Context(data=initial_data or {})
        self._contexts[context_id] = context
        return context

    def get_context(self, context_id: str) -> Optional[Context]:
        """
        Retrieve a context by ID.

        :param context_id: str - Context identifier.
        :return: Optional[Context] - The context if found, None otherwise.
        """
        return self._contexts.get(context_id)

    def delete_context(self, context_id: str) -> None:
        """
        Delete a context by ID.

        :param context_id: str - Context identifier to delete.
        """
        if context_id in self._contexts:
            del self._contexts[context_id]

    def list_contexts(self) -> list:
        """
        List all registered context IDs.

        :return: list - List of context identifiers.
        """
        return list(self._contexts.keys())

    def clear_all(self) -> None:
        """Clear all registered contexts."""
        self._contexts.clear()
