# Copyright (c) Microsoft. All rights reserved.

"""Tests for AGUIChatClient."""

import json
from collections.abc import AsyncGenerator, Awaitable, MutableSequence
from typing import Any

from agent_framework import (
    ChatOptions,
    ChatResponse,
    ChatResponseUpdate,
    Content,
    Message,
    ResponseStream,
    tool,
)
from pytest import MonkeyPatch

from agent_framework_ag_ui._client import AGUIChatClient
from agent_framework_ag_ui._http_service import AGUIHttpService


class StubAGUIChatClient(AGUIChatClient):
    """Testable wrapper exposing protected helpers."""

    @property
    def http_service(self) -> AGUIHttpService:
        """Expose http service for monkeypatching."""
        return self._http_service

    def extract_state_from_messages(self, messages: list[Message]) -> tuple[list[Message], dict[str, Any] | None]:
        """Expose state extraction helper."""
        return self._extract_state_from_messages(messages)

    def convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
        """Expose message conversion helper."""
        return self._convert_messages_to_agui_format(messages)

    def get_thread_id(self, options: dict[str, Any]) -> str:
        """Expose thread id helper."""
        return self._get_thread_id(options)

    def inner_get_response(
        self, *, messages: MutableSequence[Message], options: dict[str, Any], stream: bool = False
    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
        """Proxy to protected response call."""
        return self._inner_get_response(messages=messages, options=options, stream=stream)


class TestAGUIChatClient:
    """Test suite for AGUIChatClient."""

    async def test_client_initialization(self) -> None:
        """Test client initialization."""
        client = StubAGUIChatClient(endpoint="http://localhost:8888/")

        assert client.http_service is not None
        assert client.http_service.endpoint.startswith("http://localhost:8888")

    async def test_client_context_manager(self) -> None:
        """Test client as async context manager."""
        async with StubAGUIChatClient(endpoint="http://localhost:8888/") as client:
            assert client is not None

    async def test_extract_state_from_messages_no_state(self) -> None:
        """Test state extraction when no state is present."""
        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        messages = [
            Message(role="user", contents=["Hello"]),
            Message(role="assistant", contents=["Hi there"]),
        ]

        result_messages, state = client.extract_state_from_messages(messages)

        assert result_messages == messages
        assert state is None

    async def test_extract_state_from_messages_with_state(self) -> None:
        """Test state extraction from last message."""
        import base64

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")

        state_data = {"key": "value", "count": 42}
        state_json = json.dumps(state_data)
        state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")

        messages = [
            Message(role="user", contents=["Hello"]),
            Message(
                role="user",
                contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
            ),
        ]

        result_messages, state = client.extract_state_from_messages(messages)

        assert len(result_messages) == 1
        assert result_messages[0].text == "Hello"
        assert state == state_data

    async def test_extract_state_invalid_json(self) -> None:
        """Test state extraction with invalid JSON."""
        import base64

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")

        invalid_json = "not valid json"
        state_b64 = base64.b64encode(invalid_json.encode("utf-8")).decode("utf-8")

        messages = [
            Message(
                role="user",
                contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
            ),
        ]

        result_messages, state = client.extract_state_from_messages(messages)

        assert result_messages == messages
        assert state is None

    async def test_convert_messages_to_agui_format(self) -> None:
        """Test message conversion to AG-UI format."""
        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        messages = [
            Message(role="user", contents=["What is the weather?"]),
            Message(role="assistant", contents=["Let me check."], message_id="msg_123"),
        ]

        agui_messages = client.convert_messages_to_agui_format(messages)

        assert len(agui_messages) == 2
        assert agui_messages[0]["role"] == "user"
        assert agui_messages[0]["content"] == "What is the weather?"
        assert agui_messages[1]["role"] == "assistant"
        assert agui_messages[1]["content"] == "Let me check."
        assert agui_messages[1]["id"] == "msg_123"

    async def test_get_thread_id_from_metadata(self) -> None:
        """Test thread ID extraction from metadata."""
        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        chat_options = ChatOptions(metadata={"thread_id": "existing_thread_123"})

        thread_id = client.get_thread_id(chat_options)

        assert thread_id == "existing_thread_123"

    async def test_get_thread_id_generation(self) -> None:
        """Test automatic thread ID generation."""
        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        chat_options = ChatOptions()

        thread_id = client.get_thread_id(chat_options)

        assert thread_id.startswith("thread_")
        assert len(thread_id) > 7

    async def test_get_response_streaming(self, monkeypatch: MonkeyPatch) -> None:
        """Test streaming response method."""
        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": " world"},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            for event in mock_events:
                yield event

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        messages = [Message(role="user", contents=["Test message"])]
        chat_options = ChatOptions()

        updates: list[ChatResponseUpdate] = []
        async for update in client._inner_get_response(messages=messages, stream=True, options=chat_options):
            updates.append(update)

        assert len(updates) == 4
        assert updates[0].additional_properties is not None
        assert updates[0].additional_properties["thread_id"] == "thread_1"

        first_content = updates[1].contents[0]
        second_content = updates[2].contents[0]
        assert first_content.type == "text"
        assert second_content.type == "text"
        assert first_content.text == "Hello"
        assert second_content.text == " world"

    async def test_get_response_non_streaming(self, monkeypatch: MonkeyPatch) -> None:
        """Test non-streaming response method."""
        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Complete response"},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            for event in mock_events:
                yield event

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        messages = [Message(role="user", contents=["Test message"])]
        chat_options = {}

        response = await client.inner_get_response(messages=messages, options=chat_options)

        assert response is not None
        assert len(response.messages) > 0
        assert "Complete response" in response.text

    async def test_tool_handling(self, monkeypatch: MonkeyPatch) -> None:
        """Test that client tool metadata is sent to server.

        Client tool metadata (name, description, schema) is sent to server for planning.
        When server requests a client function, function invocation mixin
        intercepts and executes it locally. This matches .NET AG-UI implementation.
        """
        from agent_framework import tool

        @tool
        def test_tool(param: str) -> str:
            """Test tool."""
            return "result"

        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            # Client tool metadata should be sent to server
            tools: list[dict[str, Any]] | None = kwargs.get("tools")
            assert tools is not None
            assert len(tools) == 1
            tool_entry = tools[0]
            assert tool_entry["name"] == "test_tool"
            assert tool_entry["description"] == "Test tool."
            assert "parameters" in tool_entry
            for event in mock_events:
                yield event

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        messages = [Message(role="user", contents=["Test with tools"])]
        chat_options = ChatOptions(tools=[test_tool])

        response = await client.inner_get_response(messages=messages, options=chat_options)

        assert response is not None

    async def test_server_tool_calls_unwrapped_after_invocation(self, monkeypatch: MonkeyPatch) -> None:
        """Ensure server-side tool calls are exposed as FunctionCallContent after processing."""

        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_time_zone"},
            {"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"location": "Seattle"}'},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            for event in mock_events:
                yield event

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        messages = [Message(role="user", contents=["Test server tool execution"])]

        updates: list[ChatResponseUpdate] = []
        async for update in client.get_response(messages, stream=True):
            updates.append(update)

        function_calls = [
            content for update in updates for content in update.contents if content.type == "function_call"
        ]
        assert function_calls
        assert function_calls[0].name == "get_time_zone"

        assert not any(content.type == "server_function_call" for update in updates for content in update.contents)

    async def test_server_tool_calls_not_executed_locally(self, monkeypatch: MonkeyPatch) -> None:
        """Server tools should not trigger local function invocation even when client tools exist."""

        @tool
        def client_tool() -> str:
            """Client tool stub."""
            return "client"

        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "get_time_zone"},
            {"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"location": "Seattle"}'},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            for event in mock_events:
                yield event

        async def fake_auto_invoke(*args: object, **kwargs: Any) -> None:
            function_call = kwargs.get("function_call_content") or args[0]
            raise AssertionError(f"Unexpected local execution of server tool: {getattr(function_call, 'name', '?')}")

        monkeypatch.setattr("agent_framework._tools._auto_invoke_function", fake_auto_invoke)

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        messages = [Message(role="user", contents=["Test server tool execution"])]

        async for _ in client.get_response(
            messages, stream=True, options={"tool_choice": "auto", "tools": [client_tool]}
        ):
            pass

    async def test_state_transmission(self, monkeypatch: MonkeyPatch) -> None:
        """Test state is properly transmitted to server."""
        import base64

        state_data = {"user_id": "123", "session": "abc"}
        state_json = json.dumps(state_data)
        state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")

        messages = [
            Message(role="user", contents=["Hello"]),
            Message(
                role="user",
                contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
            ),
        ]

        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            assert kwargs.get("state") == state_data
            for event in mock_events:
                yield event

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        chat_options = ChatOptions()

        response = await client.inner_get_response(messages=messages, options=chat_options)

        assert response is not None

    async def test_extract_state_from_empty_messages(self) -> None:
        """Empty messages list returns empty list and None state."""
        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        result_messages, state = client.extract_state_from_messages([])
        assert result_messages == []
        assert state is None

    async def test_register_server_tool_non_dict_config(self) -> None:
        """Non-dict function_invocation_configuration is a no-op."""
        client = StubAGUIChatClient(
            endpoint="http://localhost:8888/",
            function_invocation_configuration=None,  # type: ignore[arg-type]
        )
        # Should not raise
        client._register_server_tool_placeholder("some_tool")

    async def test_non_streaming_response(self, monkeypatch: MonkeyPatch) -> None:
        """Non-streaming path collects updates into ChatResponse."""
        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "msg_1", "delta": "Hello"},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            for event in mock_events:
                yield event

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        messages = [Message(role="user", contents=["Test"])]
        response = await client.inner_get_response(messages=messages, options={}, stream=False)

        assert response is not None
        assert len(response.messages) > 0

    async def test_client_tool_sets_additional_properties(self, monkeypatch: MonkeyPatch) -> None:
        """Client tool content gets agui_thread_id additional property."""

        @tool
        def my_tool(param: str) -> str:
            """My tool."""
            return "result"

        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "TOOL_CALL_START", "toolCallId": "call_1", "toolName": "my_tool"},
            {"type": "TOOL_CALL_ARGS", "toolCallId": "call_1", "delta": '{"param": "test"}'},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            for event in mock_events:
                yield event

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        messages = [Message(role="user", contents=["Test"])]
        updates: list[ChatResponseUpdate] = []
        async for update in client._inner_get_response(messages=messages, stream=True, options={"tools": [my_tool]}):
            updates.append(update)

        # Find the function_call content - it should have agui_thread_id
        found = False
        for update in updates:
            for content in update.contents:
                if content.type == "function_call" and content.name == "my_tool":
                    assert content.additional_properties is not None
                    assert "agui_thread_id" in content.additional_properties
                    found = True
                    break
        assert found, "Expected to find function_call content for my_tool"

    async def test_interrupt_options_transmission(self, monkeypatch: MonkeyPatch) -> None:
        """Interrupt option fields are forwarded to the HTTP service."""
        available_interrupts = [{"id": "req_1", "type": "request_info"}]
        resume_payload = {"interrupts": [{"id": "req_1", "value": "approved"}]}

        mock_events = [
            {"type": "RUN_STARTED", "threadId": "thread_1", "runId": "run_1"},
            {"type": "RUN_FINISHED", "threadId": "thread_1", "runId": "run_1"},
        ]

        async def mock_post_run(*args: object, **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]:
            assert kwargs.get("available_interrupts") == available_interrupts
            assert kwargs.get("resume") == resume_payload
            for event in mock_events:
                yield event

        client = StubAGUIChatClient(endpoint="http://localhost:8888/")
        monkeypatch.setattr(client.http_service, "post_run", mock_post_run)

        messages = [Message(role="user", contents=["continue"])]
        options = {
            "available_interrupts": available_interrupts,
            "resume": resume_payload,
        }

        response = await client.inner_get_response(messages=messages, options=options)
        assert response is not None
