"""Tests for Bedrock AgentCore runtime service integration."""

from unittest.mock import MagicMock, Mock, patch

import pytest
import requests

from bedrock_agentcore_starter_toolkit.services.runtime import (
    BedrockAgentCoreClient,
    HttpBedrockAgentCoreClient,
    LocalBedrockAgentCoreClient,
    _handle_aws_response,
    _handle_streaming_response,
    generate_session_id,
)


class TestBedrockAgentCoreRuntime:
    """Test Bedrock AgentCore runtime service functionality."""

    def test_create_or_update_agent(self, mock_boto3_clients):
        """Test agent creation and update logic."""
        client = BedrockAgentCoreClient("us-west-2")

        # Test create agent (no existing agent_id)
        result = client.create_or_update_agent(
            agent_id=None,
            agent_name="test-agent",
            image_uri="123456789012.dkr.ecr.us-west-2.amazonaws.com/test:latest",
            execution_role_arn="arn:aws:iam::123456789012:role/TestRole",
            request_header_config=None,
        )

        # Verify create was called
        assert result["id"] == "test-agent-id"
        assert result["arn"] == "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id"
        mock_boto3_clients["bedrock_agentcore"].create_agent_runtime.assert_called_once()

        # Test update agent (existing agent_id)
        result = client.create_or_update_agent(
            agent_id="existing-agent-id",
            agent_name="test-agent",
            image_uri="123456789012.dkr.ecr.us-west-2.amazonaws.com/test:latest",
            execution_role_arn="arn:aws:iam::123456789012:role/TestRole",
            request_header_config=None,
        )

        # Verify update was called
        assert result["id"] == "existing-agent-id"
        assert result["arn"] == "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id"
        mock_boto3_clients["bedrock_agentcore"].update_agent_runtime.assert_called_once()

    def test_wait_for_endpoint_ready(self, mock_boto3_clients):
        """Test endpoint readiness polling."""
        client = BedrockAgentCoreClient("us-west-2")

        # Test successful readiness
        endpoint_arn = client.wait_for_agent_endpoint_ready("test-agent-id")
        expected_arn = "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id/endpoint/default"
        assert endpoint_arn == expected_arn
        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime_endpoint.assert_called()

    def test_wait_for_endpoint_ready_resource_not_found(self, mock_boto3_clients):
        """Test endpoint readiness with ResourceNotFoundException (should be handled gracefully)."""
        from botocore.exceptions import ClientError

        client = BedrockAgentCoreClient("us-west-2")

        # Mock ResourceNotFoundException followed by successful response
        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime_endpoint.side_effect = [
            ClientError(
                error_response={"Error": {"Code": "ResourceNotFoundException", "Message": "Endpoint not found"}},
                operation_name="GetAgentRuntimeEndpoint",
            ),
            {
                "status": "READY",
                "agentRuntimeEndpointArn": (
                    "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id/endpoint/default"
                ),
            },
        ]

        with patch("time.sleep"):  # Mock sleep to speed up test
            endpoint_arn = client.wait_for_agent_endpoint_ready("test-agent-id", max_wait=5)
            expected_arn = (
                "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id/endpoint/default"
            )
            assert endpoint_arn == expected_arn

    def test_invoke_endpoint(self, mock_boto3_clients):
        """Test agent invocation."""
        client = BedrockAgentCoreClient("us-west-2")

        response = client.invoke_endpoint(
            agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
            payload='{"message": "Hello"}',
            session_id="test-session-123",
        )

        # Verify invocation was called correctly
        mock_boto3_clients["bedrock_agentcore"].invoke_agent_runtime.assert_called_once_with(
            agentRuntimeArn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
            qualifier="DEFAULT",
            runtimeSessionId="test-session-123",
            payload='{"message": "Hello"}',
        )

        # Verify response structure
        assert "response" in response
        assert response["response"] == [{"data": "test response"}]

    def test_invoke_endpoint_with_custom_headers(self, mock_boto3_clients):
        """Test agent invocation with custom headers using boto3 event handlers."""
        client = BedrockAgentCoreClient("us-west-2")

        custom_headers = {
            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-Context": "production",
            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-User-ID": "123",
        }

        # Mock the event system - use dataplane client for invocations
        mock_events = Mock()
        client.dataplane_client.meta.events = mock_events

        response = client.invoke_endpoint(
            agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
            payload='{"message": "Hello"}',
            session_id="test-session-123",
            custom_headers=custom_headers,
        )

        # Verify invocation was called correctly
        mock_boto3_clients["bedrock_agentcore"].invoke_agent_runtime.assert_called_once_with(
            agentRuntimeArn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
            qualifier="DEFAULT",
            runtimeSessionId="test-session-123",
            payload='{"message": "Hello"}',
        )

        # Verify single event handler was registered for all custom headers
        assert mock_events.register_first.call_count == 1

        # Verify single unregister was called for cleanup
        assert mock_events.unregister.call_count == 1

        # Verify response structure
        assert "response" in response
        assert response["response"] == [{"data": "test response"}]

    def test_invoke_endpoint_with_empty_custom_headers(self, mock_boto3_clients):
        """Test agent invocation with empty custom headers dict."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock the event system - use dataplane client for invocations
        mock_events = Mock()
        client.dataplane_client.meta.events = mock_events

        response = client.invoke_endpoint(
            agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
            payload='{"message": "Hello"}',
            session_id="test-session-123",
            custom_headers={},
        )

        # Verify no event handlers were registered for empty headers
        mock_events.register_first.assert_not_called()
        mock_events.unregister.assert_not_called()

        # Verify response structure
        assert "response" in response
        assert response["response"] == [{"data": "test response"}]

    def test_invoke_endpoint_with_none_custom_headers(self, mock_boto3_clients):
        """Test agent invocation with None custom headers."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock the event system - use dataplane client for invocations
        mock_events = Mock()
        client.dataplane_client.meta.events = mock_events

        response = client.invoke_endpoint(
            agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
            payload='{"message": "Hello"}',
            session_id="test-session-123",
            custom_headers=None,
        )

        # Verify no event handlers were registered for None headers
        mock_events.register_first.assert_not_called()
        mock_events.unregister.assert_not_called()

        # Verify response structure
        assert "response" in response
        assert response["response"] == [{"data": "test response"}]

    def test_invoke_endpoint_headers_cleanup_on_exception(self, mock_boto3_clients):
        """Test custom headers event handlers are cleaned up even when invocation fails."""
        client = BedrockAgentCoreClient("us-west-2")

        custom_headers = {"X-Amzn-Bedrock-AgentCore-Runtime-Custom-Context": "test"}

        # Mock the event system - use dataplane client for invocations
        mock_events = Mock()
        client.dataplane_client.meta.events = mock_events

        # Mock invoke_agent_runtime to raise an exception
        mock_boto3_clients["bedrock_agentcore"].invoke_agent_runtime.side_effect = Exception("Invocation failed")

        try:
            client.invoke_endpoint(
                agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
                payload='{"message": "Hello"}',
                session_id="test-session-123",
                custom_headers=custom_headers,
            )
        except Exception:
            pass  # Expected

        # Verify event handlers were still cleaned up despite the exception
        mock_events.register_first.assert_called_once()
        mock_events.unregister.assert_called_once()

    def test_invoke_endpoint_multiple_custom_headers(self, mock_boto3_clients):
        """Test agent invocation with multiple custom headers."""
        client = BedrockAgentCoreClient("us-west-2")

        custom_headers = {
            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-Context": "production",
            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-User-ID": "user123",
            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-Session": "session456",
            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-Debug": "true",
        }

        # Mock the event system - use dataplane client for invocations
        mock_events = Mock()
        client.dataplane_client.meta.events = mock_events

        response = client.invoke_endpoint(
            agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
            payload='{"message": "Hello"}',
            session_id="test-session-123",
            custom_headers=custom_headers,
        )

        # Verify single event handler was registered for all headers
        assert mock_events.register_first.call_count == 1
        assert mock_events.unregister.call_count == 1

        # Verify response structure
        assert "response" in response
        assert response["response"] == [{"data": "test response"}]

    def test_api_error_handling(self, mock_boto3_clients):
        """Test handling of Bedrock AgentCore API errors."""
        client = BedrockAgentCoreClient("us-west-2")

        # Test basic error handling - simplified version
        assert client.region == "us-west-2"
        assert hasattr(client, "client")
        assert hasattr(client, "dataplane_client")

    def test_generate_session_id(self):
        """Test session ID generation."""
        session_id = generate_session_id()
        assert isinstance(session_id, str)
        assert len(session_id) > 0

        # Test uniqueness
        session_id2 = generate_session_id()
        assert session_id != session_id2

    def test_client_initialization(self, mock_boto3_clients):
        """Test Bedrock AgentCore client initialization."""
        client = BedrockAgentCoreClient("us-west-2")
        assert client.region == "us-west-2"
        assert client.client is not None
        assert client.dataplane_client is not None

    def test_create_agent_with_optional_configs(self, mock_boto3_clients):
        """Test create agent with network and authorizer configs."""
        client = BedrockAgentCoreClient("us-west-2")

        network_config = {"networkMode": "PRIVATE"}
        authorizer_config = {"type": "IAM"}
        protocol_config = {"serverProtocol": "MCP"}
        env_vars = {"ENV1": "HELLO", "ENV2": "WORLD"}

        result = client.create_agent(
            agent_name="test-agent",
            image_uri="123456789012.dkr.ecr.us-west-2.amazonaws.com/test:latest",
            execution_role_arn="arn:aws:iam::123456789012:role/TestRole",
            network_config=network_config,
            authorizer_config=authorizer_config,
            protocol_config=protocol_config,
            env_vars=env_vars,
        )

        assert result["id"] == "test-agent-id"
        assert result["arn"] == "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id"

        # Verify the call included optional configs
        call_args = mock_boto3_clients["bedrock_agentcore"].create_agent_runtime.call_args[1]
        assert call_args["networkConfiguration"] == network_config
        assert call_args["authorizerConfiguration"] == authorizer_config
        assert call_args["protocolConfiguration"] == protocol_config
        assert call_args["environmentVariables"] == env_vars

    def test_create_agent_error_handling(self, mock_boto3_clients):
        """Test create agent error handling."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock an exception
        mock_boto3_clients["bedrock_agentcore"].create_agent_runtime.side_effect = Exception("API Error")

        try:
            client.create_agent(
                agent_name="test-agent",
                image_uri="123456789012.dkr.ecr.us-west-2.amazonaws.com/test:latest",
                execution_role_arn="arn:aws:iam::123456789012:role/TestRole",
            )
            raise AssertionError("Expected exception")
        except Exception as e:
            assert "API Error" in str(e)

    def test_update_agent_with_optional_configs(self, mock_boto3_clients):
        """Test update agent with network and authorizer configs."""
        client = BedrockAgentCoreClient("us-west-2")

        network_config = {"networkMode": "PRIVATE"}
        authorizer_config = {"type": "IAM"}
        protocol_config = {"serverProtocol": "MCP"}
        env_vars = {"ENV1": "HELLO", "ENV2": "WORLD"}

        result = client.update_agent(
            agent_id="existing-agent-id",
            image_uri="123456789012.dkr.ecr.us-west-2.amazonaws.com/test:latest",
            execution_role_arn="arn:aws:iam::123456789012:role/TestRole",
            network_config=network_config,
            authorizer_config=authorizer_config,
            protocol_config=protocol_config,
            env_vars=env_vars,
        )

        assert result["id"] == "existing-agent-id"
        assert result["arn"] == "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id"

        # Verify the call included optional configs
        call_args = mock_boto3_clients["bedrock_agentcore"].update_agent_runtime.call_args[1]
        assert call_args["networkConfiguration"] == network_config
        assert call_args["authorizerConfiguration"] == authorizer_config
        assert call_args["protocolConfiguration"] == protocol_config
        assert call_args["environmentVariables"] == env_vars

    def test_update_agent_error_handling(self, mock_boto3_clients):
        """Test update agent error handling."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock an exception
        mock_boto3_clients["bedrock_agentcore"].update_agent_runtime.side_effect = Exception("Update Error")

        try:
            client.update_agent(
                agent_id="existing-agent-id",
                image_uri="123456789012.dkr.ecr.us-west-2.amazonaws.com/test:latest",
                execution_role_arn="arn:aws:iam::123456789012:role/TestRole",
            )
            raise AssertionError("Expected exception")
        except Exception as e:
            assert "Update Error" in str(e)

    def test_get_agent_runtime(self, mock_boto3_clients):
        """Test get agent runtime."""
        client = BedrockAgentCoreClient("us-west-2")

        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime.return_value = {
            "agentRuntimeId": "test-agent-id",
            "status": "READY",
        }

        result = client.get_agent_runtime("test-agent-id")
        assert result["agentRuntimeId"] == "test-agent-id"
        assert result["status"] == "READY"

        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime.assert_called_once_with(
            agentRuntimeId="test-agent-id"
        )

    def test_get_agent_runtime_endpoint(self, mock_boto3_clients):
        """Test get agent runtime endpoint."""
        client = BedrockAgentCoreClient("us-west-2")

        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime_endpoint.return_value = {
            "agentRuntimeEndpointArn": (
                "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id/endpoint/default"
            ),
            "status": "READY",
        }

        result = client.get_agent_runtime_endpoint("test-agent-id", "DEFAULT")
        assert "agentRuntimeEndpointArn" in result
        assert result["status"] == "READY"

        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime_endpoint.assert_called_once_with(
            agentRuntimeId="test-agent-id", endpointName="DEFAULT"
        )

    def test_invoke_endpoint_with_events_error(self, mock_boto3_clients):
        """Test invoke endpoint with events processing error."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock response that will cause an error when iterating events
        mock_response = {"response": Exception("Event processing error"), "contentType": "application/json"}
        mock_boto3_clients["bedrock_agentcore"].invoke_agent_runtime.return_value = mock_response

        response = client.invoke_endpoint(
            agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-agent-id",
            payload='{"message": "Hello"}',
            session_id="test-session-123",
        )

        # Should handle the error gracefully
        assert "response" in response
        assert len(response["response"]) == 1
        assert "Error reading EventStream" in response["response"][0]

    def test_find_agent_by_name_not_found(self):
        """Test find_agent_by_name when agent not found."""
        from bedrock_agentcore_starter_toolkit.services.runtime import BedrockAgentCoreClient

        with patch("boto3.client") as mock_boto_client:
            mock_client = MagicMock()
            mock_client.list_agent_runtimes.return_value = {
                "agentRuntimes": []  # No agents found
            }
            mock_boto_client.return_value = mock_client

            client = BedrockAgentCoreClient("us-west-2")
            result = client.find_agent_by_name("nonexistent-agent")

            assert result is None

    def test_list_agents_with_pagination(self, mock_boto3_clients):
        """Test listing agents with pagination."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock responses with pagination
        mock_boto3_clients["bedrock_agentcore"].list_agent_runtimes.side_effect = [
            {"agentRuntimes": [{"agentRuntimeId": "agent-1"}], "nextToken": "token-1"},
            {"agentRuntimes": [{"agentRuntimeId": "agent-2"}], "nextToken": None},
        ]

        result = client.list_agents()

        # Verify results contain both pages
        assert len(result) == 2
        assert result[0]["agentRuntimeId"] == "agent-1"
        assert result[1]["agentRuntimeId"] == "agent-2"

        # Verify pagination was handled
        assert mock_boto3_clients["bedrock_agentcore"].list_agent_runtimes.call_count == 2
        # First call without token
        mock_boto3_clients["bedrock_agentcore"].list_agent_runtimes.assert_any_call(maxResults=100)
        # Second call with token
        mock_boto3_clients["bedrock_agentcore"].list_agent_runtimes.assert_any_call(maxResults=100, nextToken="token-1")

    def test_create_agent_with_conflict_and_no_autoupdate(self, mock_boto3_clients):
        """Test create_agent when agent exists but auto_update_on_conflict is False."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock ConflictException
        from botocore.exceptions import ClientError

        mock_boto3_clients["bedrock_agentcore"].create_agent_runtime.side_effect = ClientError(
            {
                "Error": {
                    "Code": "ConflictException",
                    "Message": "Agent already exists",
                }
            },
            "CreateAgentRuntime",
        )

        # Try to create agent without auto_update_on_conflict
        try:
            client.create_agent(
                agent_name="test-agent",
                image_uri="123456789012.dkr.ecr.us-west-2.amazonaws.com/test:latest",
                execution_role_arn="arn:aws:iam::123456789012:role/TestRole",
                auto_update_on_conflict=False,
            )
            raise AssertionError("Should have raised ClientError")
        except ClientError as e:
            assert "ConflictException" in str(e)
            assert "use the --auto-update-on-conflict flag" in str(e)

    def test_wait_for_agent_endpoint_ready_status_failed(self, mock_boto3_clients):
        """Test wait_for_agent_endpoint_ready when endpoint update fails."""
        client = BedrockAgentCoreClient("us-west-2")

        # Set up a sequence of responses to simulate failing status
        mock_responses = [
            {
                "status": "UPDATE_FAILED",
                "failureReason": "Configuration error",
            }
        ]
        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime_endpoint.side_effect = mock_responses

        # Should return timeout message after max wait
        result = client.wait_for_agent_endpoint_ready("test-agent-id", max_wait=1)
        assert "Endpoint is taking longer than 1 seconds to be ready" in result

    def test_wait_for_agent_endpoint_ready_success(self, mock_boto3_clients):
        """Test wait_for_agent_endpoint_ready when endpoint becomes ready."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock successful endpoint response
        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime_endpoint.return_value = {
            "status": "READY",
            "agentRuntimeEndpointArn": "arn:aws:bedrock:us-west-2:123456789012:agent-endpoint/test-id",
        }

        # Should return the endpoint ARN
        result = client.wait_for_agent_endpoint_ready("test-agent-id")
        assert "arn:aws:bedrock:us-west-2:123456789012:agent-endpoint/test-id" == result

    def test_wait_for_agent_endpoint_ready_timeout(self, mock_boto3_clients):
        """Test wait_for_agent_endpoint_ready when max wait time is exceeded."""
        client = BedrockAgentCoreClient("us-west-2")

        # Mock endpoint status response for UPDATING
        mock_boto3_clients["bedrock_agentcore"].get_agent_runtime_endpoint.return_value = {
            "status": "UPDATING",
        }

        # Test with very short max_wait to force timeout
        result = client.wait_for_agent_endpoint_ready("test-agent-id", max_wait=1)
        assert "Endpoint is taking longer than 1 seconds to be ready" in result


class TestHttpBedrockAgentCoreClient:
    """Test HttpBedrockAgentCoreClient functionality."""

    def test_invoke_endpoint_success(self):
        """Test successful endpoint invocation with bearer token."""
        client = HttpBedrockAgentCoreClient("us-west-2")

        # Mock successful HTTP response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"data: response content\n\n"
        mock_response.text = "data: response content\n\n"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with patch("requests.post", return_value=mock_response) as mock_post:
            result = client.invoke_endpoint(
                agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-id",
                payload='{"message": "hello"}',  # JSON string as it comes from invoke_bedrock_agentcore
                session_id="test-session-123",
                bearer_token="test-bearer-token",
            )

            # Verify request was made correctly
            mock_post.assert_called_once()
            call_args = mock_post.call_args

            # Check URL
            expected_url = "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/arn%3Aaws%3Abedrock_agentcore%3Aus-west-2%3A123456789012%3Aagent-runtime%2Ftest-id/invocations"
            assert call_args[0][0] == expected_url

            # Check headers
            headers = call_args[1]["headers"]
            assert headers["Authorization"] == "Bearer test-bearer-token"
            assert headers["Content-Type"] == "application/json"
            assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] == "test-session-123"

            # Check payload - should now send the payload directly, not wrapped
            body = call_args[1]["json"]
            assert body == {"message": "hello"}

            # Check query params
            params = call_args[1]["params"]
            assert params["qualifier"] == "DEFAULT"

            # Check timeout
            assert call_args[1]["timeout"] == 900

            # Verify response
            assert result["response"] == "data: response content\n\n"

    def test_invoke_endpoint_with_custom_qualifier(self):
        """Test invocation with custom endpoint qualifier."""
        client = HttpBedrockAgentCoreClient("us-east-1")

        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"test content"
        mock_response.text = "test content"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with patch("requests.post", return_value=mock_response) as mock_post:
            client.invoke_endpoint(
                agent_arn="arn:aws:bedrock_agentcore:us-east-1:123456789012:agent-runtime/test-id",
                payload='"test payload"',  # JSON string as it would come from invoke_bedrock_agentcore
                session_id="session-456",
                bearer_token="token-123",
                endpoint_name="CUSTOM",
            )

            # Verify custom qualifier was used
            call_args = mock_post.call_args
            params = call_args[1]["params"]
            assert params["qualifier"] == "CUSTOM"

    def test_invoke_endpoint_http_error(self):
        """Test handling of HTTP errors."""
        client = HttpBedrockAgentCoreClient("us-west-2")

        # Mock HTTP error response
        mock_response = Mock()
        mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError("404 Not Found")

        with patch("requests.post", return_value=mock_response):
            with pytest.raises(requests.exceptions.HTTPError):
                client.invoke_endpoint(
                    agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/nonexistent",
                    payload='{"test": "data"}',
                    session_id="session-123",
                    bearer_token="token-456",
                )

    def test_invoke_endpoint_connection_error(self):
        """Test handling of connection errors."""
        client = HttpBedrockAgentCoreClient("us-west-2")

        with patch("requests.post", side_effect=requests.exceptions.ConnectionError("Connection failed")):
            with pytest.raises(requests.exceptions.ConnectionError):
                client.invoke_endpoint(
                    agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-id",
                    payload='{"test": "data"}',
                    session_id="session-123",
                    bearer_token="token-456",
                )

    def test_invoke_endpoint_timeout(self):
        """Test handling of request timeout."""
        client = HttpBedrockAgentCoreClient("us-west-2")

        with patch("requests.post", side_effect=requests.exceptions.Timeout("Request timed out")):
            with pytest.raises(requests.exceptions.Timeout):
                client.invoke_endpoint(
                    agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-id",
                    payload='{"test": "data"}',
                    session_id="session-123",
                    bearer_token="token-456",
                )

    def test_invoke_endpoint_empty_response(self):
        """Test handling of empty response."""
        client = HttpBedrockAgentCoreClient("us-west-2")

        # Mock empty response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b""  # Empty content
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with patch("requests.post", return_value=mock_response):
            with pytest.raises(ValueError, match="Empty response from agent endpoint"):
                client.invoke_endpoint(
                    agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-id",
                    payload='{"test": "data"}',
                    session_id="session-123",
                    bearer_token="token-456",
                )

    def test_url_encoding_special_characters(self):
        """Test proper URL encoding of agent ARN with special characters."""
        client = HttpBedrockAgentCoreClient("us-west-2")

        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"test"
        mock_response.text = "test"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with patch("requests.post", return_value=mock_response) as mock_post:
            # ARN with special characters that need encoding
            complex_arn = "arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-id:with:colons"

            client.invoke_endpoint(
                agent_arn=complex_arn, payload='{"test": "data"}', session_id="session-123", bearer_token="token-456"
            )

            # Verify URL encoding
            call_args = mock_post.call_args
            url = call_args[0][0]
            # Colons should be encoded as %3A
            assert "%3A" in url
            assert "with%3Acolons" in url

    def test_payload_types(self):
        """Test different payload types are handled correctly."""
        client = HttpBedrockAgentCoreClient("us-west-2")

        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"response"
        mock_response.text = "response"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        # Test cases as JSON strings (as they come from invoke_bedrock_agentcore)
        test_cases = [
            ('{"message": "hello"}', {"message": "hello"}),  # Valid JSON dict
            ('"simple string"', "simple string"),  # Valid JSON string
            ('["list", "payload"]', ["list", "payload"]),  # Valid JSON list
            ("42", 42),  # Valid JSON number
            ("invalid json string", {"payload": "invalid json string"}),  # Invalid JSON - fallback
        ]

        with patch("requests.post", return_value=mock_response) as mock_post:
            for payload_input, expected_body in test_cases:
                client.invoke_endpoint(
                    agent_arn="arn:aws:bedrock_agentcore:us-west-2:123456789012:agent-runtime/test-id",
                    payload=payload_input,
                    session_id="session-123",
                    bearer_token="token-456",
                )

                # Verify payload was parsed and sent correctly
                call_args = mock_post.call_args
                body = call_args[1]["json"]
                assert body == expected_body

                mock_post.reset_mock()


class TestLocalBedrockAgentCoreClient:
    """Test LocalBedrockAgentCoreClient functionality."""

    def test_initialization(self):
        """Test LocalBedrockAgentCoreClient initialization."""
        endpoint = "http://localhost:8080"
        client = LocalBedrockAgentCoreClient(endpoint)

        assert client.endpoint == endpoint

    def test_invoke_endpoint_success(self):
        """Test successful endpoint invocation."""
        client = LocalBedrockAgentCoreClient("http://localhost:8080")

        # Mock successful HTTP response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"test response content"
        mock_response.text = "test response content"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with (
            patch("requests.post", return_value=mock_response) as mock_post,
            patch(
                "bedrock_agentcore_starter_toolkit.services.runtime._handle_http_response",
                return_value={"response": "test response"},
            ) as mock_handle,
        ):
            result = client.invoke_endpoint(
                session_id="test-session-123", payload='{"message": "hello"}', workload_access_token="test-token-456"
            )

            # Verify request was made correctly
            mock_post.assert_called_once()
            call_args = mock_post.call_args

            # Check URL
            expected_url = "http://localhost:8080/invocations"
            assert call_args[0][0] == expected_url

            # Check headers - need to import the constants
            from bedrock_agentcore.runtime.models import ACCESS_TOKEN_HEADER, SESSION_HEADER

            headers = call_args[1]["headers"]
            assert headers["Content-Type"] == "application/json"
            assert headers[ACCESS_TOKEN_HEADER] == "test-token-456"
            assert headers[SESSION_HEADER] == "test-session-123"

            # Check payload
            body = call_args[1]["json"]
            assert body == {"message": "hello"}

            # Check timeout
            assert call_args[1]["timeout"] == 900
            assert call_args[1]["stream"] is True

            # Verify response handling
            mock_handle.assert_called_once_with(mock_response)
            assert result == {"response": "test response"}

    def test_invoke_endpoint_with_non_json_payload(self):
        """Test invocation with non-JSON payload."""
        client = LocalBedrockAgentCoreClient("http://localhost:9090")

        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"response"
        mock_response.text = "response"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with (
            patch("requests.post", return_value=mock_response) as mock_post,
            patch(
                "bedrock_agentcore_starter_toolkit.services.runtime._handle_http_response",
                return_value={"response": "wrapped"},
            ),
        ):
            # Test with invalid JSON string
            client.invoke_endpoint(
                session_id="session-456", payload="invalid json string", workload_access_token="token-123"
            )

            # Verify payload was wrapped
            call_args = mock_post.call_args
            body = call_args[1]["json"]
            assert body == {"payload": "invalid json string"}

    def test_invoke_endpoint_with_custom_headers(self):
        """Test endpoint invocation with custom headers."""
        client = LocalBedrockAgentCoreClient("http://localhost:8080")

        custom_headers = {
            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-Context": "local",
            "X-Amzn-Bedrock-AgentCore-Runtime-Custom-Debug": "true",
        }

        # Mock successful HTTP response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"local response with headers"
        mock_response.text = "local response with headers"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with (
            patch("requests.post", return_value=mock_response) as mock_post,
            patch(
                "bedrock_agentcore_starter_toolkit.services.runtime._handle_http_response",
                return_value={"response": "local response with custom headers"},
            ) as mock_handle,
        ):
            result = client.invoke_endpoint(
                session_id="test-session-123",
                payload='{"message": "hello"}',
                workload_access_token="test-token-456",
                custom_headers=custom_headers,
            )

            # Verify request was made correctly
            mock_post.assert_called_once()
            call_args = mock_post.call_args

            # Check URL
            expected_url = "http://localhost:8080/invocations"
            assert call_args[0][0] == expected_url

            # Check headers - need to import the constants
            from bedrock_agentcore.runtime.models import ACCESS_TOKEN_HEADER, SESSION_HEADER

            headers = call_args[1]["headers"]
            assert headers["Content-Type"] == "application/json"
            assert headers[ACCESS_TOKEN_HEADER] == "test-token-456"
            assert headers[SESSION_HEADER] == "test-session-123"

            # Verify custom headers were added
            assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Custom-Context"] == "local"
            assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Custom-Debug"] == "true"

            # Verify response handling
            mock_handle.assert_called_once_with(mock_response)
            assert result == {"response": "local response with custom headers"}

    def test_invoke_endpoint_with_empty_custom_headers(self):
        """Test endpoint invocation with empty custom headers dict."""
        client = LocalBedrockAgentCoreClient("http://localhost:8080")

        # Mock successful HTTP response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"local response without headers"
        mock_response.text = "local response without headers"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with (
            patch("requests.post", return_value=mock_response) as mock_post,
            patch(
                "bedrock_agentcore_starter_toolkit.services.runtime._handle_http_response",
                return_value={"response": "local response"},
            ) as mock_handle,
        ):
            result = client.invoke_endpoint(
                session_id="test-session-123",
                payload='{"message": "hello"}',
                workload_access_token="test-token-456",
                custom_headers={},
            )

            # Verify request was made correctly
            mock_post.assert_called_once()
            call_args = mock_post.call_args

            # Check headers - should only have default headers
            from bedrock_agentcore.runtime.models import ACCESS_TOKEN_HEADER, SESSION_HEADER

            headers = call_args[1]["headers"]
            assert headers["Content-Type"] == "application/json"
            assert headers[ACCESS_TOKEN_HEADER] == "test-token-456"
            assert headers[SESSION_HEADER] == "test-session-123"

            # Verify no custom headers were added
            custom_header_keys = [k for k in headers.keys() if k.startswith("X-Amzn-Bedrock-AgentCore-Runtime-Custom-")]
            assert len(custom_header_keys) == 0

            # Verify response handling
            mock_handle.assert_called_once_with(mock_response)
            assert result == {"response": "local response"}

    def test_invoke_endpoint_with_none_custom_headers(self):
        """Test endpoint invocation with None custom headers."""
        client = LocalBedrockAgentCoreClient("http://localhost:8080")

        # Mock successful HTTP response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.content = b"local response without headers"
        mock_response.text = "local response without headers"
        mock_response.raise_for_status.return_value = None
        mock_response.headers = {"content-type": "application/json"}

        with (
            patch("requests.post", return_value=mock_response) as mock_post,
            patch(
                "bedrock_agentcore_starter_toolkit.services.runtime._handle_http_response",
                return_value={"response": "local response"},
            ) as mock_handle,
        ):
            result = client.invoke_endpoint(
                session_id="test-session-123",
                payload='{"message": "hello"}',
                workload_access_token="test-token-456",
                custom_headers=None,
            )

            # Verify request was made correctly
            mock_post.assert_called_once()
            call_args = mock_post.call_args

            # Check headers - should only have default headers
            from bedrock_agentcore.runtime.models import ACCESS_TOKEN_HEADER, SESSION_HEADER

            headers = call_args[1]["headers"]
            assert headers["Content-Type"] == "application/json"
            assert headers[ACCESS_TOKEN_HEADER] == "test-token-456"
            assert headers[SESSION_HEADER] == "test-session-123"

            # Verify no custom headers were added
            custom_header_keys = [k for k in headers.keys() if k.startswith("X-Amzn-Bedrock-AgentCore-Runtime-Custom-")]
            assert len(custom_header_keys) == 0

            # Verify response handling
            mock_handle.assert_called_once_with(mock_response)
            assert result == {"response": "local response"}


class TestHandleStreamingResponse:
    """Test _handle_streaming_response functionality."""

    def test_handle_streaming_response_with_data_lines(self):
        """Test streaming response with data: prefixed lines."""
        # Mock response object with JSON data chunks
        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'data: "Hello from agent"',
            b'data: "This is a streaming response"',
            b'data: "Final chunk"',
        ]

        # Mock console to capture print calls
        with patch("bedrock_agentcore_starter_toolkit.services.runtime.console") as mock_console:
            result = _handle_streaming_response(mock_response)

            # Verify result structure - function returns empty dict for streaming
            assert result == {}

            # Verify console.print was called for each JSON chunk + final newline
            assert mock_console.print.call_count == 4  # 3 chunks + 1 final newline
            mock_console.print.assert_any_call("Hello from agent", end="")
            mock_console.print.assert_any_call("This is a streaming response", end="")
            mock_console.print.assert_any_call("Final chunk", end="")
            mock_console.print.assert_any_call()  # Final newline call

    def test_handle_aws_response_byte_parsing(self):
        """Test _handle_aws_response properly parses byte strings."""
        # Test with byte string in response
        response = {
            "response": [b'"Hello from agent"', b'"Another message"'],
            "ResponseMetadata": {"RequestId": "test-123"},
        }

        result = _handle_aws_response(response)

        # Verify bytes were decoded and JSON parsed
        assert result["response"] == ["Hello from agent", "Another message"]
        assert result["ResponseMetadata"]["RequestId"] == "test-123"

        # Test with non-JSON bytes
        response = {
            "response": [b"plain text message"],
        }

        result = _handle_aws_response(response)
        assert result["response"] == ["plain text message"]

        # Test with mixed content
        response = {
            "response": [b'"JSON string"', "regular string", b"plain bytes"],
        }

        result = _handle_aws_response(response)
        assert result["response"] == ["JSON string", "regular string", "plain bytes"]
