"""Tests for the LLM hook system."""

import pytest
from unittest.mock import AsyncMock

from roohai.hooks import LLMHook, make_wrapper


class TestLLMHookProtocol:
    def test_async_function_conforms(self):
        """A plain async function with the right signature satisfies LLMHook."""

        async def my_hook(message, *, history=None, system_prompt=None):
            return "ok"

        assert isinstance(my_hook, LLMHook)

    def test_callable_class_conforms(self):
        """A class with async __call__ satisfies LLMHook."""

        class MyHook:
            async def __call__(self, message, *, history=None, system_prompt=None):
                return "ok"

        assert isinstance(MyHook(), LLMHook)


class TestMakeWrapper:
    async def test_passthrough(self):
        """With no before/after, wrapper just calls next_fn."""
        next_fn = AsyncMock(return_value="response")
        hook = make_wrapper(next_fn)
        result = await hook("hello", history=None, system_prompt="sys")
        assert result == "response"
        next_fn.assert_awaited_once_with("hello", history=None, system_prompt="sys")

    async def test_before_transform(self):
        """Before callback should transform inputs."""
        next_fn = AsyncMock(return_value="response")

        async def before(msg, hist, prompt):
            return f"prefix: {msg}", hist, prompt

        hook = make_wrapper(next_fn, before=before)
        result = await hook("hello", history=None, system_prompt=None)
        assert result == "response"
        next_fn.assert_awaited_once_with("prefix: hello", history=None, system_prompt=None)

    async def test_after_transform(self):
        """After callback should transform the response."""
        next_fn = AsyncMock(return_value="**bold**")

        async def after(response):
            return response.replace("**", "")

        hook = make_wrapper(next_fn, after=after)
        result = await hook("hello", history=None, system_prompt=None)
        assert result == "bold"

    async def test_before_and_after(self):
        """Both before and after should compose correctly."""
        next_fn = AsyncMock(return_value="response")

        async def before(msg, hist, prompt):
            return msg.upper(), hist, prompt

        async def after(response):
            return f"[{response}]"

        hook = make_wrapper(next_fn, before=before, after=after)
        result = await hook("hello", history=None, system_prompt=None)
        assert result == "[response]"
        next_fn.assert_awaited_once_with("HELLO", history=None, system_prompt=None)
