"""Lifecycle hook protocol and fire helpers for the Swarm framework.

Lifecycle hooks are observational — they cannot alter execution flow.
They fire at LLM and tool call boundaries for logging, metrics, and tracing.
"""

import logging
from typing import Any, Protocol, runtime_checkable

logger = logging.getLogger(__name__)


@runtime_checkable
class LifecycleHook(Protocol):
    """Observational hook — cannot alter execution flow.

    Implement any subset of the four methods. Unimplemented methods should
    be async no-ops. The runner will call each method at the appropriate
    boundary and swallow any exceptions so hooks never crash the runner.

    ``response`` parameter contract:
        - DialogueRunner (streaming): ``SwarmAssistantMessage`` on success, ``None`` on failure.
        - AgentRunner (batch): ``ChatCompletion`` on success, ``None`` on failure.
    """

    async def on_llm_start(
        self, agent: Any, messages: list[Any], context_variables: dict[str, Any]
    ) -> None: ...

    async def on_llm_end(
        self,
        agent: Any,
        response: Any,
        error: BaseException | None = None,
    ) -> None: ...

    async def on_tool_start(
        self,
        agent: Any,
        tool_call_id: str,
        tool_name: str,
        arguments: dict[str, Any],
    ) -> None: ...

    async def on_tool_end(
        self,
        agent: Any,
        tool_call_id: str,
        tool_name: str,
        result: Any,
        error: BaseException | None = None,
    ) -> None: ...


# ---------------------------------------------------------------------------
# Fire helpers — iterate hooks in try/except, never crash the runner
# ---------------------------------------------------------------------------


async def fire_hook_on_llm_start(
    agent: Any,
    messages: list[Any],
    context_variables: dict[str, Any],
) -> None:
    """Fire ``on_llm_start`` on every valid hook attached to *agent*."""
    for hook in getattr(agent, "hooks", None) or []:
        if not isinstance(hook, LifecycleHook):
            logger.warning("Skipping invalid hook: %s", type(hook).__name__)
            continue
        try:
            await hook.on_llm_start(agent, messages, context_variables)
        except Exception:
            logger.exception("Hook %s.on_llm_start failed", type(hook).__name__)


async def fire_hook_on_llm_end(
    agent: Any,
    response: Any,
    *,
    error: BaseException | None = None,
) -> None:
    """Fire ``on_llm_end`` on every valid hook attached to *agent*."""
    for hook in getattr(agent, "hooks", None) or []:
        if not isinstance(hook, LifecycleHook):
            logger.warning("Skipping invalid hook: %s", type(hook).__name__)
            continue
        try:
            await hook.on_llm_end(agent, response, error=error)
        except Exception:
            logger.exception("Hook %s.on_llm_end failed", type(hook).__name__)


async def fire_hook_on_tool_start(
    agent: Any,
    tool_call_id: str,
    tool_name: str,
    arguments: dict[str, Any],
) -> None:
    """Fire ``on_tool_start`` on every valid hook attached to *agent*."""
    for hook in getattr(agent, "hooks", None) or []:
        if not isinstance(hook, LifecycleHook):
            logger.warning("Skipping invalid hook: %s", type(hook).__name__)
            continue
        try:
            await hook.on_tool_start(agent, tool_call_id, tool_name, arguments)
        except Exception:
            logger.exception("Hook %s.on_tool_start failed", type(hook).__name__)


async def fire_hook_on_tool_end(
    agent: Any,
    tool_call_id: str,
    tool_name: str,
    result: Any,
    *,
    error: BaseException | None = None,
) -> None:
    """Fire ``on_tool_end`` on every valid hook attached to *agent*."""
    for hook in getattr(agent, "hooks", None) or []:
        if not isinstance(hook, LifecycleHook):
            logger.warning("Skipping invalid hook: %s", type(hook).__name__)
            continue
        try:
            await hook.on_tool_end(agent, tool_call_id, tool_name, result, error=error)
        except Exception:
            logger.exception("Hook %s.on_tool_end failed", type(hook).__name__)
