"""LLM hook system for customizing the pipeline's LLM step."""

from __future__ import annotations

from typing import Callable, Awaitable, Protocol, runtime_checkable


@runtime_checkable
class LLMHook(Protocol):
    """Protocol for LLM hooks — any async callable matching this signature."""

    async def __call__(
        self,
        message: str,
        *,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
    ) -> str: ...


def make_wrapper(
    next_fn: Callable[..., Awaitable[str]],
    *,
    before: Callable | None = None,
    after: Callable | None = None,
) -> LLMHook:
    """Build a hook that wraps an existing chat function.

    Args:
        next_fn: The underlying chat function to call (e.g. ``pipeline.llm.chat``).
        before: Optional async callable ``(message, history, system_prompt) -> (message, history, system_prompt)``
                that transforms inputs before calling *next_fn*.
        after: Optional async callable ``(response) -> response`` that transforms the output.
    """

    async def wrapper(
        message: str,
        *,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
    ) -> str:
        msg, hist, prompt = message, history, system_prompt
        if before is not None:
            msg, hist, prompt = await before(msg, hist, prompt)
        result = await next_fn(msg, history=hist, system_prompt=prompt)
        if after is not None:
            result = await after(result)
        return result

    return wrapper  # type: ignore[return-value]
