import json
from collections import defaultdict
from typing import Any, cast

import pytest

from agents.agent import Agent
from agents.items import ItemHelpers, ModelResponse, TResponseInputItem
from agents.lifecycle import AgentHooks, RunHooks
from agents.models.interface import Model
from agents.run import Runner
from agents.run_context import AgentHookContext, RunContextWrapper, TContext
from agents.run_internal.run_loop import validate_run_hooks
from agents.tool import Tool, function_tool
from agents.tool_context import ToolContext
from tests.test_agent_llm_hooks import AgentHooksForTests

from .fake_model import FakeModel
from .test_responses import (
    get_function_tool,
    get_function_tool_call,
    get_handoff_tool_call,
    get_text_message,
)


class RunHooksForTests(RunHooks):
    def __init__(self):
        self.events: dict[str, int] = defaultdict(int)
        self.tool_context_ids: list[str] = []

    def reset(self):
        self.events.clear()
        self.tool_context_ids.clear()

    async def on_agent_start(
        self, context: AgentHookContext[TContext], agent: Agent[TContext]
    ) -> None:
        self.events["on_agent_start"] += 1

    async def on_agent_end(
        self, context: RunContextWrapper[TContext], agent: Agent[TContext], output: Any
    ) -> None:
        self.events["on_agent_end"] += 1

    async def on_handoff(
        self,
        context: RunContextWrapper[TContext],
        from_agent: Agent[TContext],
        to_agent: Agent[TContext],
    ) -> None:
        self.events["on_handoff"] += 1

    async def on_tool_start(
        self, context: RunContextWrapper[TContext], agent: Agent[TContext], tool: Tool
    ) -> None:
        self.events["on_tool_start"] += 1

    async def on_tool_end(
        self,
        context: RunContextWrapper[TContext],
        agent: Agent[TContext],
        tool: Tool,
        result: object,
    ) -> None:
        self.events["on_tool_end"] += 1
        if isinstance(context, ToolContext):
            self.tool_context_ids.append(context.tool_call_id)

    async def on_llm_start(
        self,
        context: RunContextWrapper[TContext],
        agent: Agent[TContext],
        system_prompt: str | None,
        input_items: list[TResponseInputItem],
    ) -> None:
        self.events["on_llm_start"] += 1

    async def on_llm_end(
        self,
        context: RunContextWrapper[TContext],
        agent: Agent[TContext],
        response: ModelResponse,
    ) -> None:
        self.events["on_llm_end"] += 1


# Example test using the above hooks
@pytest.mark.asyncio
async def test_async_run_hooks_with_llm():
    hooks = RunHooksForTests()
    model = FakeModel()

    agent = Agent(name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[])
    # Simulate a single LLM call producing an output:
    model.set_next_output([get_text_message("hello")])
    await Runner.run(agent, input="hello", hooks=hooks)
    # Expect one on_agent_start, one on_llm_start, one on_llm_end, and one on_agent_end
    assert hooks.events == {
        "on_agent_start": 1,
        "on_llm_start": 1,
        "on_llm_end": 1,
        "on_agent_end": 1,
    }


# test_sync_run_hook_with_llm()
def test_sync_run_hook_with_llm():
    hooks = RunHooksForTests()
    model = FakeModel()
    agent = Agent(name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[])
    # Simulate a single LLM call producing an output:
    model.set_next_output([get_text_message("hello")])
    Runner.run_sync(agent, input="hello", hooks=hooks)
    # Expect one on_agent_start, one on_llm_start, one on_llm_end, and one on_agent_end
    assert hooks.events == {
        "on_agent_start": 1,
        "on_llm_start": 1,
        "on_llm_end": 1,
        "on_agent_end": 1,
    }


# test_streamed_run_hooks_with_llm():
@pytest.mark.asyncio
async def test_streamed_run_hooks_with_llm():
    hooks = RunHooksForTests()
    model = FakeModel()
    agent = Agent(name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[])
    # Simulate a single LLM call producing an output:
    model.set_next_output([get_text_message("hello")])
    stream = Runner.run_streamed(agent, input="hello", hooks=hooks)

    async for event in stream.stream_events():
        if event.type == "raw_response_event":
            continue
        if event.type == "agent_updated_stream_event":
            print(f"[EVENT] agent_updated → {event.new_agent.name}")
        elif event.type == "run_item_stream_event":
            item = event.item
            if item.type == "tool_call_item":
                print("[EVENT] tool_call_item")
            elif item.type == "tool_call_output_item":
                print(f"[EVENT] tool_call_output_item → {item.output}")
            elif item.type == "message_output_item":
                text = ItemHelpers.text_message_output(item)
                print(f"[EVENT] message_output_item → {text}")

    # Expect one on_agent_start, one on_llm_start, one on_llm_end, and one on_agent_end
    assert hooks.events == {
        "on_agent_start": 1,
        "on_llm_start": 1,
        "on_llm_end": 1,
        "on_agent_end": 1,
    }


# test_async_run_hooks_with_agent_hooks_with_llm
@pytest.mark.asyncio
async def test_async_run_hooks_with_agent_hooks_with_llm():
    hooks = RunHooksForTests()
    agent_hooks = AgentHooksForTests()
    model = FakeModel()

    agent = Agent(
        name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[], hooks=agent_hooks
    )
    # Simulate a single LLM call producing an output:
    model.set_next_output([get_text_message("hello")])
    await Runner.run(agent, input="hello", hooks=hooks)
    # Expect one on_agent_start, one on_llm_start, one on_llm_end, and one on_agent_end
    assert hooks.events == {
        "on_agent_start": 1,
        "on_llm_start": 1,
        "on_llm_end": 1,
        "on_agent_end": 1,
    }
    # Expect one on_start, one on_llm_start, one on_llm_end, and one on_end
    assert agent_hooks.events == {"on_start": 1, "on_llm_start": 1, "on_llm_end": 1, "on_end": 1}


@pytest.mark.asyncio
async def test_run_hooks_llm_error_non_streaming(monkeypatch):
    hooks = RunHooksForTests()
    model = FakeModel()
    agent = Agent(name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[])

    async def boom(*args, **kwargs):
        raise RuntimeError("boom")

    monkeypatch.setattr(FakeModel, "get_response", boom, raising=True)

    with pytest.raises(RuntimeError, match="boom"):
        await Runner.run(agent, input="hello", hooks=hooks)

    # Current behavior is that hooks will not fire on LLM failure
    assert hooks.events["on_agent_start"] == 1
    assert hooks.events["on_llm_start"] == 1
    assert hooks.events["on_llm_end"] == 0
    assert hooks.events["on_agent_end"] == 0


class DummyAgentHooks(AgentHooks):
    """Agent-scoped hooks used to verify runtime validation."""


@pytest.mark.asyncio
async def test_runner_run_rejects_agent_hooks():
    model = FakeModel()
    agent = Agent(name="A", model=model)
    hooks = cast(RunHooks, DummyAgentHooks())

    with pytest.raises(TypeError, match="Run hooks must be instances of RunHooks"):
        await Runner.run(agent, input="hello", hooks=hooks)


def test_runner_run_streamed_rejects_agent_hooks():
    model = FakeModel()
    agent = Agent(name="A", model=model)
    hooks = cast(RunHooks, DummyAgentHooks())

    with pytest.raises(TypeError, match="Run hooks must be instances of RunHooks"):
        Runner.run_streamed(agent, input="hello", hooks=hooks)


def test_validate_run_hooks_rejects_non_hook_objects() -> None:
    with pytest.raises(TypeError, match="Received object"):
        validate_run_hooks(object())


class BoomModel(Model):
    async def get_response(self, *a, **k):
        raise AssertionError("get_response should not be called in streaming test")

    async def stream_response(self, *a, **k):
        yield {"foo": "bar"}
        raise RuntimeError("stream blew up")


@pytest.mark.asyncio
async def test_streamed_run_hooks_llm_error(monkeypatch):
    """
    Verify that when the streaming path raises, we still emit on_llm_start
    but do NOT emit on_llm_end (current behavior), and the exception propagates.
    """
    hooks = RunHooksForTests()
    agent = Agent(name="A", model=BoomModel(), tools=[get_function_tool("f", "res")], handoffs=[])

    stream = Runner.run_streamed(agent, input="hello", hooks=hooks)

    # Consuming the stream should surface the exception
    with pytest.raises(RuntimeError, match="stream blew up"):
        async for _ in stream.stream_events():
            pass

    # Current behavior: success-only on_llm_end; ensure starts fired but ends did not.
    assert hooks.events["on_agent_start"] == 1
    assert hooks.events["on_llm_start"] == 1
    assert hooks.events["on_llm_end"] == 0
    assert hooks.events["on_agent_end"] == 0


class RunHooksWithTurnInput(RunHooks):
    """Run hooks that capture turn_input from on_agent_start."""

    def __init__(self):
        self.captured_turn_inputs: list[list[Any]] = []

    async def on_agent_start(
        self, context: AgentHookContext[TContext], agent: Agent[TContext]
    ) -> None:
        self.captured_turn_inputs.append(list(context.turn_input))


@pytest.mark.asyncio
async def test_run_hooks_receives_turn_input_string():
    """Test that on_agent_start receives turn_input when input is a string."""
    hooks = RunHooksWithTurnInput()
    model = FakeModel()
    agent = Agent(name="test", model=model)

    model.set_next_output([get_text_message("response")])
    await Runner.run(agent, input="hello world", hooks=hooks)

    assert len(hooks.captured_turn_inputs) == 1
    turn_input = hooks.captured_turn_inputs[0]
    assert len(turn_input) == 1
    assert turn_input[0]["content"] == "hello world"
    assert turn_input[0]["role"] == "user"


@pytest.mark.asyncio
async def test_run_hooks_receives_turn_input_list():
    """Test that on_agent_start receives turn_input when input is a list."""
    hooks = RunHooksWithTurnInput()
    model = FakeModel()
    agent = Agent(name="test", model=model)

    input_items: list[Any] = [
        {"role": "user", "content": "first message"},
        {"role": "user", "content": "second message"},
    ]

    model.set_next_output([get_text_message("response")])
    await Runner.run(agent, input=input_items, hooks=hooks)

    assert len(hooks.captured_turn_inputs) == 1
    turn_input = hooks.captured_turn_inputs[0]
    assert len(turn_input) == 2
    assert turn_input[0]["content"] == "first message"
    assert turn_input[1]["content"] == "second message"


@pytest.mark.asyncio
async def test_run_hooks_receives_turn_input_streamed():
    """Test that on_agent_start receives turn_input in streamed mode."""
    hooks = RunHooksWithTurnInput()
    model = FakeModel()
    agent = Agent(name="test", model=model)

    model.set_next_output([get_text_message("response")])
    result = Runner.run_streamed(agent, input="streamed input", hooks=hooks)
    async for _ in result.stream_events():
        pass

    assert len(hooks.captured_turn_inputs) == 1
    turn_input = hooks.captured_turn_inputs[0]
    assert len(turn_input) == 1
    assert turn_input[0]["content"] == "streamed input"


@pytest.mark.asyncio
async def test_run_hooks_count_tool_and_handoff_invocations():
    hooks = RunHooksForTests()
    model = FakeModel()

    agent_1 = Agent(name="test_1", model=model)
    agent_2 = Agent(
        name="test_2",
        model=model,
        handoffs=[agent_1],
        tools=[get_function_tool("some_function", "result")],
    )

    model.add_multiple_turn_outputs(
        [
            [get_function_tool_call("some_function", json.dumps({"a": "b"}))],
            [get_text_message("a_message"), get_handoff_tool_call(agent_1)],
            [get_text_message("done")],
        ]
    )
    await Runner.run(agent_2, input="user_message", hooks=hooks)

    assert hooks.events["on_tool_start"] == 1
    assert hooks.events["on_tool_end"] == 1
    assert hooks.events["on_handoff"] == 1
    assert hooks.events["on_agent_start"] == 2
    assert hooks.events["on_agent_end"] == 1
    assert len(hooks.tool_context_ids) == 1


@pytest.mark.asyncio
async def test_streamed_run_hooks_count_tool_and_handoff_invocations():
    hooks = RunHooksForTests()
    model = FakeModel()

    agent_1 = Agent(name="test_1", model=model)
    agent_2 = Agent(
        name="test_2",
        model=model,
        handoffs=[agent_1],
        tools=[get_function_tool("some_function", "result")],
    )

    model.add_multiple_turn_outputs(
        [
            [
                get_function_tool_call("some_function", json.dumps({"a": "b"})),
                get_function_tool_call("some_function", json.dumps({"a": "b"})),
            ],
            [get_text_message("a_message"), get_handoff_tool_call(agent_1)],
            [get_text_message("done")],
        ]
    )
    stream = Runner.run_streamed(agent_2, input="user_message", hooks=hooks)
    async for _ in stream.stream_events():
        pass

    assert hooks.events["on_tool_start"] == 2
    assert hooks.events["on_tool_end"] == 2
    assert hooks.events["on_handoff"] == 1
    assert hooks.events["on_agent_start"] == 2
    assert hooks.events["on_agent_end"] == 1
    assert len(hooks.tool_context_ids) == 2


@pytest.mark.asyncio
async def test_tool_end_hooks_receive_raw_function_tool_result():
    class RecordingRunHooks(RunHooks):
        def __init__(self):
            self.result: object | None = None

        async def on_tool_end(
            self,
            context: RunContextWrapper[Any],
            agent: Agent[Any],
            tool: Tool,
            result: object,
        ) -> None:
            self.result = result

    class RecordingAgentHooks(AgentHooks):
        def __init__(self):
            self.result: object | None = None

        async def on_tool_end(
            self,
            context: RunContextWrapper[Any],
            agent: Agent[Any],
            tool: Tool,
            result: object,
        ) -> None:
            self.result = result

    metadata_result: dict[str, object] = {"status": "ok", "count": 1}

    @function_tool
    def get_metadata() -> dict[str, object]:
        return metadata_result

    run_hooks = RecordingRunHooks()
    agent_hooks = RecordingAgentHooks()
    model = FakeModel()
    agent = Agent(name="test", model=model, tools=[get_metadata], hooks=agent_hooks)

    model.add_multiple_turn_outputs(
        [
            [get_function_tool_call("get_metadata", "{}")],
            [get_text_message("done")],
        ]
    )

    await Runner.run(agent, input="user_message", hooks=run_hooks)

    assert run_hooks.result is metadata_result
    assert agent_hooks.result is metadata_result