"""
Core trait definitions and enumerations.

This module defines the foundational types for the trait system:
- Trait enum with all available traits
- TraitDefinition for trait metadata
- Core interfaces and type aliases
"""

from __future__ import annotations

import weakref
from collections.abc import Callable
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING, Any, TypeVar

if TYPE_CHECKING:
    pass

__all__ = ["Trait", "TraitDefinition", "TraitImpl", "TraitValidator"]

# Type variables for trait system
T = TypeVar("T")
TraitValidator = Callable[[Any], bool]
TraitImpl = type[T]


class Trait(str, Enum):
    """
    Enumeration of all available traits in the LionAGI system.

    Each trait represents a specific behavior or capability that can be
    composed into domain models. Traits are implemented as Protocols
    for optimal type safety and performance.
    """

    # Core identity and lifecycle traits
    IDENTIFIABLE = "identifiable"  # Has unique ID and identity methods
    TEMPORAL = "temporal"  # Has creation/modification timestamps
    AUDITABLE = "auditable"  # Tracks changes and emits audit events
    HASHABLE = "hashable"  # Provides stable hashing behavior

    # Behavior and operation traits
    OPERABLE = "operable"  # Supports operations and transformations
    OBSERVABLE = "observable"  # Emits events for state changes
    VALIDATABLE = "validatable"  # Supports validation and constraint checking
    SERIALIZABLE = "serializable"  # Can be serialized/deserialized

    # Advanced composition traits
    COMPOSABLE = "composable"  # Can be composed with other models
    EXTENSIBLE = "extensible"  # Supports dynamic extension/plugins
    CACHEABLE = "cacheable"  # Provides caching and memoization
    INDEXABLE = "indexable"  # Can be indexed and searched

    # Performance and optimization traits
    LAZY = "lazy"  # Supports lazy loading and evaluation
    STREAMING = "streaming"  # Supports streaming updates
    PARTIAL = "partial"  # Supports partial/incremental construction

    # Security and capability traits
    SECURED = "secured"  # Has security policies and access control
    CAPABILITY_AWARE = (
        "capability_aware"  # Participates in capability-based security
    )


@dataclass(frozen=True, slots=True)
class TraitDefinition:
    """
    Metadata definition for a specific trait implementation.

    This immutable dataclass captures all metadata needed to track
    and manage trait implementations within the system.
    """

    trait: Trait
    protocol_type: type[Any]  # Protocol type
    implementation_type: type[Any]
    dependencies: frozenset[Trait] = field(default_factory=frozenset)
    version: str = "1.0.0"
    description: str = ""

    # Performance tracking
    registration_time: float = field(default=0.0)
    validation_checks: int = field(default=0)

    # Weak reference to avoid circular dependencies
    _weak_impl_ref: weakref.ReferenceType[type[Any]] = field(
        init=False, repr=False
    )

    def __post_init__(self) -> None:
        """Initialize weak reference to implementation."""

        def cleanup_callback(ref: weakref.ReferenceType[type[Any]]) -> None:
            pass

        object.__setattr__(
            self,
            "_weak_impl_ref",
            weakref.ref(self.implementation_type, cleanup_callback),
        )

    @property
    def is_alive(self) -> bool:
        """Check if the implementation type is still alive."""
        return self._weak_impl_ref() is not None

    def validate_dependencies(self, available_traits: set[Trait]) -> bool:
        """
        Validate that all trait dependencies are satisfied.

        Args:
            available_traits: Set of traits available on the target type

        Returns:
            True if all dependencies are satisfied
        """
        return self.dependencies.issubset(available_traits)

    def get_dependency_graph(self) -> dict[Trait, set[Trait]]:
        """
        Get the dependency graph for this trait.

        Returns:
            Mapping of trait to its direct dependencies
        """
        return {self.trait: set(self.dependencies)}


# Default trait definitions with zero dependencies
DEFAULT_TRAIT_DEFINITIONS: dict[Trait, TraitDefinition] = {}


def _initialize_default_definitions() -> None:
    """Initialize default trait definitions (called at module load)."""

    from .protocols import (
        Auditable,
        Cacheable,
        CapabilityAware,
        Composable,
        Extensible,
        Hashable,
        Identifiable,
        Indexable,
        Lazy,
        Observable,
        Operable,
        Partial,
        Secured,
        Serializable,
        Streaming,
        Temporal,
        Validatable,
    )

    # Map traits to their protocol types
    _protocol_mapping = {
        Trait.IDENTIFIABLE: Identifiable,
        Trait.TEMPORAL: Temporal,
        Trait.AUDITABLE: Auditable,
        Trait.HASHABLE: Hashable,
        Trait.OPERABLE: Operable,
        Trait.OBSERVABLE: Observable,
        Trait.VALIDATABLE: Validatable,
        Trait.SERIALIZABLE: Serializable,
        Trait.COMPOSABLE: Composable,
        Trait.EXTENSIBLE: Extensible,
        Trait.CACHEABLE: Cacheable,
        Trait.INDEXABLE: Indexable,
        Trait.LAZY: Lazy,
        Trait.STREAMING: Streaming,
        Trait.PARTIAL: Partial,
        Trait.SECURED: Secured,
        Trait.CAPABILITY_AWARE: CapabilityAware,
    }

    _trait_dependencies = {
        # Core traits - no dependencies
        Trait.IDENTIFIABLE: frozenset(),
        Trait.TEMPORAL: frozenset(),
        Trait.HASHABLE: frozenset(),
        # Behavioral traits
        Trait.AUDITABLE: frozenset({Trait.IDENTIFIABLE, Trait.TEMPORAL}),
        Trait.OPERABLE: frozenset({Trait.IDENTIFIABLE}),
        Trait.OBSERVABLE: frozenset({Trait.IDENTIFIABLE}),
        Trait.VALIDATABLE: frozenset(),
        Trait.SERIALIZABLE: frozenset({Trait.IDENTIFIABLE}),
        # Advanced composition traits
        Trait.COMPOSABLE: frozenset({Trait.IDENTIFIABLE}),
        Trait.EXTENSIBLE: frozenset({Trait.IDENTIFIABLE}),
        Trait.CACHEABLE: frozenset({Trait.HASHABLE}),
        Trait.INDEXABLE: frozenset({Trait.IDENTIFIABLE, Trait.HASHABLE}),
        # Performance traits
        Trait.LAZY: frozenset({Trait.IDENTIFIABLE}),
        Trait.STREAMING: frozenset({Trait.OBSERVABLE}),
        Trait.PARTIAL: frozenset({Trait.VALIDATABLE}),
        # Security traits
        Trait.SECURED: frozenset({Trait.IDENTIFIABLE}),
        Trait.CAPABILITY_AWARE: frozenset({Trait.SECURED, Trait.IDENTIFIABLE}),
    }

    # Create default definitions with dependencies
    for trait, protocol_type in _protocol_mapping.items():
        DEFAULT_TRAIT_DEFINITIONS[trait] = TraitDefinition(
            trait=trait,
            protocol_type=protocol_type,
            implementation_type=object,  # Placeholder
            dependencies=_trait_dependencies.get(trait, frozenset()),
            description=f"Default definition for {trait.name} trait",
        )


_initialize_default_definitions()
