"""
Extension Hook System - Lifecycle hooks for agent customization

Provides a hook-based architecture for customizing agent behavior at various
lifecycle points, inspired by agent-zero's extension system.
"""

import asyncio
import importlib.util
import inspect
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Optional, TypeVar, Union

T = TypeVar("T")


class HookPoint(Enum):
    """
    Available hook points in the agent lifecycle.

    These hooks allow extensions to modify behavior at specific points.
    """

    # Agent lifecycle
    AGENT_INIT = "agent_init"
    AGENT_CLEANUP = "agent_cleanup"

    # Message loop
    MONOLOGUE_START = "monologue_start"
    MONOLOGUE_END = "monologue_end"
    MESSAGE_LOOP_START = "message_loop_start"
    MESSAGE_LOOP_END = "message_loop_end"

    # LLM calls
    BEFORE_LLM_CALL = "before_llm_call"
    AFTER_LLM_CALL = "after_llm_call"
    REASONING_STREAM = "reasoning_stream"
    RESPONSE_STREAM = "response_stream"

    # Tool execution
    BEFORE_TOOL_EXECUTION = "before_tool_execution"
    AFTER_TOOL_EXECUTION = "after_tool_execution"
    TOOL_ERROR = "tool_error"

    # Context management
    BEFORE_CONTEXT_SAVE = "before_context_save"
    AFTER_CONTEXT_LOAD = "after_context_load"
    CONTEXT_THRESHOLD_REACHED = "context_threshold_reached"

    # Memory operations
    BEFORE_MEMORY_SAVE = "before_memory_save"
    AFTER_MEMORY_LOAD = "after_memory_load"
    MEMORY_CONSOLIDATION = "memory_consolidation"

    # Prompts
    PROMPT_ASSEMBLY_BEFORE = "prompt_assembly_before"
    PROMPT_ASSEMBLY_AFTER = "prompt_assembly_after"

    # User interaction
    USER_INPUT_RECEIVED = "user_input_received"
    USER_OUTPUT_GENERATED = "user_output_generated"

    # Error handling
    ERROR_OCCURRED = "error_occurred"
    INTERVENTION_REQUESTED = "intervention_requested"


@dataclass
class HookContext:
    """Context passed to hook handlers."""

    hook_point: HookPoint
    agent: Any  # Agent instance
    data: dict[str, Any] = field(default_factory=dict)
    modified: bool = False
    cancelled: bool = False
    error: Optional[Exception] = None

    def set(self, key: str, value: Any) -> None:
        """Set a value in the context data."""
        self.data[key] = value
        self.modified = True

    def get(self, key: str, default: Any = None) -> Any:
        """Get a value from the context data."""
        return self.data.get(key, default)

    def cancel(self) -> None:
        """Cancel the current operation."""
        self.cancelled = True

    def has_error(self) -> bool:
        """Check if an error occurred."""
        return self.error is not None


class Extension(ABC):
    """
    Base class for extensions.

    Extensions implement specific hook handlers to customize agent behavior.
    """

    # Priority determines execution order (lower = earlier)
    priority: int = 50

    # Name for identification
    name: str = "base_extension"

    # Hook points this extension handles
    hooks: list[HookPoint] = []

    @abstractmethod
    async def execute(self, context: HookContext) -> HookContext:
        """
        Execute the extension logic.

        Args:
            context: HookContext with current state

        Returns:
            Modified HookContext
        """
        pass

    def handles(self, hook_point: HookPoint) -> bool:
        """Check if this extension handles a specific hook point."""
        return hook_point in self.hooks or not self.hooks

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(name={self.name}, priority={self.priority})"


class ExtensionManager:
    """
    Manages extensions and executes hooks.

    Provides:
    - Extension registration and discovery
    - Hook execution with priority ordering
    - Extension loading from directories
    - Context propagation between hooks
    """

    def __init__(self):
        self._extensions: dict[HookPoint, list[Extension]] = {
            hook: [] for hook in HookPoint
        }
        self._all_extensions: list[Extension] = []
        self._enabled: bool = True

    def register(self, extension: Extension) -> None:
        """
        Register an extension.

        Args:
            extension: Extension instance to register
        """
        self._all_extensions.append(extension)

        # Register for specific hooks or all hooks
        if extension.hooks:
            for hook in extension.hooks:
                self._extensions[hook].append(extension)
                self._extensions[hook].sort(key=lambda e: e.priority)
        else:
            # Register for all hooks
            for hook in HookPoint:
                self._extensions[hook].append(extension)
                self._extensions[hook].sort(key=lambda e: e.priority)

    def unregister(self, extension: Extension) -> bool:
        """
        Unregister an extension.

        Args:
            extension: Extension to remove

        Returns:
            True if removed, False if not found
        """
        if extension not in self._all_extensions:
            return False

        self._all_extensions.remove(extension)

        for hook in HookPoint:
            if extension in self._extensions[hook]:
                self._extensions[hook].remove(extension)

        return True

    async def execute_hook(
        self,
        hook_point: HookPoint,
        agent: Any,
        data: Optional[dict[str, Any]] = None,
    ) -> HookContext:
        """
        Execute all extensions registered for a hook point.

        Args:
            hook_point: The hook point to execute
            agent: Agent instance
            data: Initial context data

        Returns:
            Final HookContext after all extensions
        """
        if not self._enabled:
            return HookContext(hook_point=hook_point, agent=agent, data=data or {})

        context = HookContext(
            hook_point=hook_point,
            agent=agent,
            data=data or {},
        )

        for extension in self._extensions[hook_point]:
            if context.cancelled:
                break

            try:
                context = await extension.execute(context)
            except Exception as e:
                context.error = e
                # Continue with other extensions unless cancelled

        return context

    def execute_hook_sync(
        self,
        hook_point: HookPoint,
        agent: Any,
        data: Optional[dict[str, Any]] = None,
    ) -> HookContext:
        """
        Synchronous wrapper for execute_hook.

        Args:
            hook_point: The hook point to execute
            agent: Agent instance
            data: Initial context data

        Returns:
            Final HookContext
        """
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                # Create a new task in the running loop
                import concurrent.futures

                with concurrent.futures.ThreadPoolExecutor() as pool:
                    future = pool.submit(
                        asyncio.run,
                        self.execute_hook(hook_point, agent, data),
                    )
                    return future.result()
            else:
                return loop.run_until_complete(
                    self.execute_hook(hook_point, agent, data)
                )
        except RuntimeError:
            return asyncio.run(self.execute_hook(hook_point, agent, data))

    def load_extensions_from_directory(self, directory: Union[str, Path]) -> int:
        """
        Load extensions from a directory.

        Extensions should be Python files with an Extension class.

        Args:
            directory: Path to extensions directory

        Returns:
            Number of extensions loaded
        """
        dir_path = Path(directory)
        if not dir_path.exists():
            return 0

        loaded = 0
        for py_file in dir_path.glob("*.py"):
            if py_file.name.startswith("_"):
                continue

            try:
                spec = importlib.util.spec_from_file_location(
                    py_file.stem, py_file
                )
                if spec and spec.loader:
                    module = importlib.util.module_from_spec(spec)
                    spec.loader.exec_module(module)

                    # Find Extension subclasses
                    for name, obj in inspect.getmembers(module):
                        if (
                            inspect.isclass(obj)
                            and issubclass(obj, Extension)
                            and obj is not Extension
                        ):
                            self.register(obj())
                            loaded += 1

            except Exception as e:
                # Log error but continue loading other extensions
                print(f"Error loading extension from {py_file}: {e}")

        return loaded

    def get_extensions(self, hook_point: Optional[HookPoint] = None) -> list[Extension]:
        """
        Get registered extensions.

        Args:
            hook_point: Filter by hook point (optional)

        Returns:
            List of extensions
        """
        if hook_point:
            return self._extensions[hook_point].copy()
        return self._all_extensions.copy()

    def enable(self) -> None:
        """Enable extension execution."""
        self._enabled = True

    def disable(self) -> None:
        """Disable extension execution."""
        self._enabled = False

    @property
    def is_enabled(self) -> bool:
        """Check if extensions are enabled."""
        return self._enabled

    def clear(self) -> None:
        """Remove all registered extensions."""
        self._all_extensions.clear()
        for hook in HookPoint:
            self._extensions[hook].clear()


# Built-in extensions


class MemoryRecallExtension(Extension):
    """
    Automatically recall relevant memories before LLM calls.
    """

    name = "memory_recall"
    priority = 30
    hooks = [HookPoint.BEFORE_LLM_CALL]

    async def execute(self, context: HookContext) -> HookContext:
        """Inject relevant memories into the prompt."""
        agent = context.agent

        # Check if agent has memory system
        if hasattr(agent, "memory") and agent.memory:
            query = context.get("prompt", "")
            if query:
                memories = await self._recall_memories(agent.memory, query)
                if memories:
                    context.set("injected_memories", memories)
                    context.modified = True

        return context

    async def _recall_memories(self, memory, query: str) -> list[str]:
        """Recall relevant memories."""
        try:
            if hasattr(memory, "search"):
                return await memory.search(query, limit=5)
        except Exception:
            pass
        return []


class ContextCompressionExtension(Extension):
    """
    Compress context when threshold is reached.
    """

    name = "context_compression"
    priority = 20
    hooks = [HookPoint.CONTEXT_THRESHOLD_REACHED]

    async def execute(self, context: HookContext) -> HookContext:
        """Compress conversation history."""
        agent = context.agent

        if hasattr(agent, "history"):
            # Summarize older messages
            history = agent.history
            if len(history) > 20:
                # Keep first 5 and last 10, summarize middle
                to_summarize = history[5:-10]
                if to_summarize:
                    summary = self._summarize_messages(to_summarize)
                    context.set("compressed_history", summary)
                    context.modified = True

        return context

    def _summarize_messages(self, messages: list) -> str:
        """Simple summarization - can be enhanced with LLM."""
        return f"[Summary of {len(messages)} previous messages]"


class SecretsFilterExtension(Extension):
    """
    Filter secrets from LLM responses.
    """

    name = "secrets_filter"
    priority = 10
    hooks = [HookPoint.RESPONSE_STREAM, HookPoint.USER_OUTPUT_GENERATED]

    def __init__(self, secrets: Optional[dict[str, str]] = None):
        self.secrets = secrets or {}
        self._buffer = ""

    async def execute(self, context: HookContext) -> HookContext:
        """Mask secrets in output."""
        text = context.get("text", "")
        if text and self.secrets:
            masked = self._mask_secrets(text)
            context.set("text", masked)
            context.modified = True

        return context

    def _mask_secrets(self, text: str) -> str:
        """Replace secret values with placeholders."""
        for key, value in self.secrets.items():
            if value and len(value) >= 3 and value in text:
                text = text.replace(value, f"§§secret({key})")
        return text

    def add_secret(self, key: str, value: str) -> None:
        """Add a secret to filter."""
        self.secrets[key] = value

    def remove_secret(self, key: str) -> None:
        """Remove a secret."""
        self.secrets.pop(key, None)


class InterventionExtension(Extension):
    """
    Handle user intervention requests.
    """

    name = "intervention"
    priority = 5
    hooks = [HookPoint.INTERVENTION_REQUESTED]

    async def execute(self, context: HookContext) -> HookContext:
        """Process intervention request."""
        message = context.get("message", "")
        broadcast = context.get("broadcast", False)

        # Mark as intervention handled
        context.set("intervention_handled", True)

        # If broadcast, propagate to superior agent
        if broadcast:
            agent = context.agent
            if hasattr(agent, "superior") and agent.superior:
                context.set("propagate_to_superior", True)

        return context


class LoggingExtension(Extension):
    """
    Log hook executions for debugging.
    """

    name = "logging"
    priority = 100  # Run last
    hooks = []  # Handle all hooks

    def __init__(self, log_callback: Optional[Callable[[str], None]] = None):
        self.log_callback = log_callback or print

    async def execute(self, context: HookContext) -> HookContext:
        """Log hook execution."""
        self.log_callback(
            f"[Hook] {context.hook_point.value}: "
            f"modified={context.modified}, cancelled={context.cancelled}"
        )
        return context


# Global extension manager instance
_extension_manager: Optional[ExtensionManager] = None


def get_extension_manager() -> ExtensionManager:
    """Get the global extension manager."""
    global _extension_manager
    if _extension_manager is None:
        _extension_manager = ExtensionManager()
    return _extension_manager


def register_extension(extension: Extension) -> None:
    """Register an extension globally."""
    get_extension_manager().register(extension)


async def execute_hook(
    hook_point: HookPoint,
    agent: Any,
    data: Optional[dict[str, Any]] = None,
) -> HookContext:
    """Execute a hook globally."""
    return await get_extension_manager().execute_hook(hook_point, agent, data)


def execute_hook_sync(
    hook_point: HookPoint,
    agent: Any,
    data: Optional[dict[str, Any]] = None,
) -> HookContext:
    """Execute a hook synchronously."""
    return get_extension_manager().execute_hook_sync(hook_point, agent, data)
