import json
import os
import sys
from unittest.mock import MagicMock

import pytest

sys.path.insert(
    0, os.path.abspath("../../../../..")
)  # Adds the parent directory to the system path

from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig
from litellm.types.llms.openai import OpenAIRealtimeEventTypes


class TestBedrockRealtimeConfig:
    """Test suite for BedrockRealtimeConfig class"""

    def test_initialization(self):
        """Test that BedrockRealtimeConfig initializes with correct defaults"""
        config = BedrockRealtimeConfig()

        assert config is not None
        assert config.max_tokens == 1024
        assert config.temperature == 0.7
        assert config.top_p == 0.9
        assert config.voice_id == "matthew"
        assert config.output_sample_rate_hertz == 24000
        assert config.input_sample_rate_hertz == 16000
        assert config.text_media_type == "text/plain"

    def test_session_configuration_request(self):
        """Test session configuration request generation"""
        config = BedrockRealtimeConfig()

        session_config = config.session_configuration_request("amazon.nova-sonic-v1:0")
        session_dict = json.loads(session_config)

        assert "session_start" in session_dict
        assert "prompt_start" in session_dict

        # Check session start
        session_start = session_dict["session_start"]["event"]["sessionStart"]
        assert session_start["inferenceConfiguration"]["maxTokens"] == 1024
        assert session_start["inferenceConfiguration"]["temperature"] == 0.7

        # Check prompt start
        prompt_start = session_dict["prompt_start"]["event"]["promptStart"]
        assert prompt_start["audioOutputConfiguration"]["voiceId"] == "matthew"
        assert prompt_start["audioOutputConfiguration"]["sampleRateHertz"] == 24000

    def test_session_configuration_with_tools(self):
        """Test session configuration with tools"""
        config = BedrockRealtimeConfig()

        tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "description": "Get weather",
                    "parameters": {
                        "type": "object",
                        "properties": {"location": {"type": "string"}},
                    },
                },
            }
        ]

        session_config = config.session_configuration_request(
            "amazon.nova-sonic-v1:0", tools=tools
        )
        session_dict = json.loads(session_config)

        prompt_start = session_dict["prompt_start"]["event"]["promptStart"]
        assert "toolConfiguration" in prompt_start
        assert "tools" in prompt_start["toolConfiguration"]
        assert len(prompt_start["toolConfiguration"]["tools"]) == 1
        assert (
            prompt_start["toolConfiguration"]["tools"][0]["toolSpec"]["name"]
            == "get_weather"
        )

    def test_transform_tools_to_bedrock_format(self):
        """Test OpenAI tool format to Bedrock format transformation"""
        config = BedrockRealtimeConfig()

        openai_tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "description": "Get current weather",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "location": {"type": "string", "description": "City name"}
                        },
                        "required": ["location"],
                    },
                },
            }
        ]

        bedrock_tools = config._transform_tools_to_bedrock_format(openai_tools)

        assert len(bedrock_tools) == 1
        assert bedrock_tools[0]["toolSpec"]["name"] == "get_weather"
        assert bedrock_tools[0]["toolSpec"]["description"] == "Get current weather"
        assert "inputSchema" in bedrock_tools[0]["toolSpec"]

        # Verify the schema is properly JSON stringified
        schema = json.loads(bedrock_tools[0]["toolSpec"]["inputSchema"]["json"])
        assert schema["type"] == "object"
        assert "location" in schema["properties"]

    def test_audio_format_mapping(self):
        """Test audio format to sample rate mapping"""
        config = BedrockRealtimeConfig()

        # Test PCM16 format
        assert config._map_audio_format_to_sample_rate("pcm16", is_output=True) == 24000
        assert (
            config._map_audio_format_to_sample_rate("pcm16", is_output=False) == 16000
        )

        # Test G.711 formats
        assert (
            config._map_audio_format_to_sample_rate("g711_ulaw", is_output=True) == 8000
        )
        assert (
            config._map_audio_format_to_sample_rate("g711_alaw", is_output=False)
            == 8000
        )

    def test_transform_session_update_event(self):
        """Test session.update event transformation"""
        config = BedrockRealtimeConfig()

        session_update = {
            "type": "session.update",
            "session": {
                "temperature": 0.9,
                "voice": "joanna",
                "max_response_output_tokens": 2048,
                "output_audio_format": "pcm16",
            },
        }

        messages = config.transform_session_update_event(session_update)

        assert len(messages) >= 2  # At least session start and prompt start

        # Verify attributes were updated
        assert config.temperature == 0.9
        assert config.voice_id == "joanna"
        assert config.max_tokens == 2048

        # Verify session start message
        session_start = json.loads(messages[0])
        assert (
            session_start["event"]["sessionStart"]["inferenceConfiguration"][
                "temperature"
            ]
            == 0.9
        )

    def test_transform_session_update_with_tools(self):
        """Test session.update with tools"""
        config = BedrockRealtimeConfig()

        session_update = {
            "type": "session.update",
            "session": {
                "tools": [
                    {
                        "type": "function",
                        "function": {
                            "name": "get_time",
                            "description": "Get current time",
                            "parameters": {"type": "object", "properties": {}},
                        },
                    }
                ]
            },
        }

        messages = config.transform_session_update_event(session_update)

        # Find prompt start message
        prompt_start = json.loads(messages[1])
        assert "toolConfiguration" in prompt_start["event"]["promptStart"]

    def test_transform_conversation_item_create_text(self):
        """Test conversation.item.create with text"""
        config = BedrockRealtimeConfig()

        item_create = {
            "type": "conversation.item.create",
            "item": {
                "type": "message",
                "role": "user",
                "content": [{"type": "input_text", "text": "Hello, how are you?"}],
            },
        }

        messages = config.transform_conversation_item_create_event(item_create)

        # Should have content start, text input, and content end
        assert len(messages) == 3

        content_start = json.loads(messages[0])
        assert content_start["event"]["contentStart"]["type"] == "TEXT"
        assert content_start["event"]["contentStart"]["role"] == "USER"

        text_input = json.loads(messages[1])
        assert text_input["event"]["textInput"]["content"] == "Hello, how are you?"

    def test_transform_conversation_item_create_tool_result(self):
        """Test conversation.item.create with tool result"""
        config = BedrockRealtimeConfig()

        tool_result = {
            "type": "conversation.item.create",
            "item": {
                "type": "function_call_output",
                "call_id": "call_123",
                "output": json.dumps({"temperature": 72, "conditions": "sunny"}),
            },
        }

        messages = config.transform_conversation_item_create_event(tool_result)

        # Should have content start, tool result, and content end
        assert len(messages) == 3

        content_start = json.loads(messages[0])
        assert content_start["event"]["contentStart"]["type"] == "TOOL"
        assert content_start["event"]["contentStart"]["role"] == "TOOL"
        assert (
            content_start["event"]["contentStart"]["toolResultInputConfiguration"][
                "toolUseId"
            ]
            == "call_123"
        )

    def test_transform_input_audio_buffer_append(self):
        """Test input_audio_buffer.append transformation"""
        config = BedrockRealtimeConfig()

        audio_append = {
            "type": "input_audio_buffer.append",
            "audio": "base64_audio_data_here",
        }

        messages = config.transform_input_audio_buffer_append_event(audio_append)

        # First call should include content start
        assert len(messages) == 2

        content_start = json.loads(messages[0])
        assert content_start["event"]["contentStart"]["type"] == "AUDIO"
        assert (
            content_start["event"]["contentStart"]["audioInputConfiguration"][
                "sampleRateHertz"
            ]
            == 16000
        )

        audio_input = json.loads(messages[1])
        assert audio_input["event"]["audioInput"]["content"] == "base64_audio_data_here"

    def test_transform_input_audio_buffer_commit(self):
        """Test input_audio_buffer.commit transformation"""
        config = BedrockRealtimeConfig()

        # First append to set the flag
        config._audio_content_started = True

        commit = {"type": "input_audio_buffer.commit"}

        messages = config.transform_input_audio_buffer_commit_event(commit)

        assert len(messages) == 1
        content_end = json.loads(messages[0])
        assert "contentEnd" in content_end["event"]


class TestBedrockRealtimeResponseTransformation:
    """Test suite for response transformation"""

    def test_transform_session_start_response(self):
        """Test sessionStart response transformation"""
        config = BedrockRealtimeConfig()
        logging_obj = MagicMock()
        logging_obj.litellm_trace_id = "trace_123"

        bedrock_message = {
            "event": {
                "sessionStart": {
                    "inferenceConfiguration": {"maxTokens": 1024, "temperature": 0.7}
                }
            }
        }

        result = config.transform_realtime_response(
            json.dumps(bedrock_message),
            "amazon.nova-sonic-v1:0",
            logging_obj,
            realtime_response_transform_input={
                "session_configuration_request": None,
                "current_output_item_id": None,
                "current_response_id": None,
                "current_conversation_id": None,
                "current_delta_chunks": [],
                "current_item_chunks": [],
                "current_delta_type": None,
            },
        )

        assert len(result["response"]) == 1
        assert result["response"][0]["type"] == "session.created"
        assert result["response"][0]["session"]["id"] == "trace_123"
        assert "model" in result["response"][0]["session"]

    def test_transform_text_output_response(self):
        """Test textOutput response transformation"""
        config = BedrockRealtimeConfig()
        logging_obj = MagicMock()
        logging_obj.litellm_trace_id = "trace_123"

        # First create a content start to initialize IDs
        content_start_message = {
            "event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}
        }

        result1 = config.transform_realtime_response(
            json.dumps(content_start_message),
            "amazon.nova-sonic-v1:0",
            logging_obj,
            realtime_response_transform_input={
                "session_configuration_request": json.dumps({"configured": True}),
                "current_output_item_id": None,
                "current_response_id": None,
                "current_conversation_id": None,
                "current_delta_chunks": [],
                "current_item_chunks": [],
                "current_delta_type": None,
            },
        )

        # Now send text output
        text_output_message = {"event": {"textOutput": {"content": "Hello, world!"}}}

        result2 = config.transform_realtime_response(
            json.dumps(text_output_message),
            "amazon.nova-sonic-v1:0",
            logging_obj,
            realtime_response_transform_input={
                "session_configuration_request": json.dumps({"configured": True}),
                "current_output_item_id": result1["current_output_item_id"],
                "current_response_id": result1["current_response_id"],
                "current_conversation_id": result1["current_conversation_id"],
                "current_delta_chunks": result1["current_delta_chunks"],
                "current_item_chunks": [],
                "current_delta_type": result1["current_delta_type"],
            },
        )

        # Check for text delta
        text_deltas = [
            msg for msg in result2["response"] if msg["type"] == "response.text.delta"
        ]
        assert len(text_deltas) == 1
        assert text_deltas[0]["delta"] == "Hello, world!"

        # Check that delta chunks are accumulated
        assert len(result2["current_delta_chunks"]) == 1

    def test_transform_audio_output_response(self):
        """Test audioOutput response transformation"""
        config = BedrockRealtimeConfig()
        logging_obj = MagicMock()
        logging_obj.litellm_trace_id = "trace_123"

        # First create a content start for audio
        content_start_message = {
            "event": {"contentStart": {"role": "ASSISTANT", "type": "AUDIO"}}
        }

        result1 = config.transform_realtime_response(
            json.dumps(content_start_message),
            "amazon.nova-sonic-v1:0",
            logging_obj,
            realtime_response_transform_input={
                "session_configuration_request": json.dumps({"configured": True}),
                "current_output_item_id": None,
                "current_response_id": None,
                "current_conversation_id": None,
                "current_delta_chunks": [],
                "current_item_chunks": [],
                "current_delta_type": None,
            },
        )

        # Now send audio output
        audio_output_message = {
            "event": {"audioOutput": {"content": "base64_audio_content"}}
        }

        result2 = config.transform_realtime_response(
            json.dumps(audio_output_message),
            "amazon.nova-sonic-v1:0",
            logging_obj,
            realtime_response_transform_input={
                "session_configuration_request": json.dumps({"configured": True}),
                "current_output_item_id": result1["current_output_item_id"],
                "current_response_id": result1["current_response_id"],
                "current_conversation_id": result1["current_conversation_id"],
                "current_delta_chunks": [],
                "current_item_chunks": [],
                "current_delta_type": result1["current_delta_type"],
            },
        )

        # Check for audio delta
        audio_deltas = [
            msg for msg in result2["response"] if msg["type"] == "response.audio.delta"
        ]
        assert len(audio_deltas) == 1
        assert audio_deltas[0]["delta"] == "base64_audio_content"

    def test_transform_tool_use_response(self):
        """Test toolUse response transformation"""
        config = BedrockRealtimeConfig()
        logging_obj = MagicMock()
        logging_obj.litellm_trace_id = "trace_123"

        tool_use_message = {
            "event": {
                "toolUse": {
                    "toolUseId": "tool_call_123",
                    "toolName": "get_weather",
                    "input": json.dumps({"location": "San Francisco"}),
                }
            }
        }

        result = config.transform_realtime_response(
            json.dumps(tool_use_message),
            "amazon.nova-sonic-v1:0",
            logging_obj,
            realtime_response_transform_input={
                "session_configuration_request": json.dumps({"configured": True}),
                "current_output_item_id": "item_123",
                "current_response_id": "resp_123",
                "current_conversation_id": "conv_123",
                "current_delta_chunks": [],
                "current_item_chunks": [],
                "current_delta_type": "text",
            },
        )

        # Check for function call event
        assert len(result["response"]) == 1
        function_call = result["response"][0]
        assert function_call["type"] == "response.function_call_arguments.done"
        assert function_call["call_id"] == "tool_call_123"
        assert function_call["name"] == "get_weather"

        # Verify arguments are properly formatted
        args = json.loads(function_call["arguments"])
        assert args["location"] == "San Francisco"

    def test_transform_content_end_text(self):
        """Test contentEnd for text response"""
        config = BedrockRealtimeConfig()
        logging_obj = MagicMock()
        logging_obj.litellm_trace_id = "trace_123"

        # Create some delta chunks first
        delta_chunks = [
            {"delta": "Hello, ", "type": "response.text.delta"},
            {"delta": "world!", "type": "response.text.delta"},
        ]

        content_end_message = {"event": {"contentEnd": {}}}

        result = config.transform_realtime_response(
            json.dumps(content_end_message),
            "amazon.nova-sonic-v1:0",
            logging_obj,
            realtime_response_transform_input={
                "session_configuration_request": json.dumps({"configured": True}),
                "current_output_item_id": "item_123",
                "current_response_id": "resp_123",
                "current_conversation_id": "conv_123",
                "current_delta_chunks": delta_chunks,
                "current_item_chunks": [],
                "current_delta_type": "text",
            },
        )

        # Should have text.done, content_part.done, and output_item.done
        assert len(result["response"]) == 3

        text_done = [
            msg for msg in result["response"] if msg["type"] == "response.text.done"
        ][0]
        assert text_done["text"] == "Hello, world!"

        # Delta chunks should be reset
        assert result["current_delta_chunks"] is None

    def test_transform_prompt_end_response(self):
        """Test promptEnd response transformation"""
        config = BedrockRealtimeConfig()
        logging_obj = MagicMock()
        logging_obj.litellm_trace_id = "trace_123"

        prompt_end_message = {"event": {"promptEnd": {}}}

        result = config.transform_realtime_response(
            json.dumps(prompt_end_message),
            "amazon.nova-sonic-v1:0",
            logging_obj,
            realtime_response_transform_input={
                "session_configuration_request": json.dumps({"configured": True}),
                "current_output_item_id": "item_123",
                "current_response_id": "resp_123",
                "current_conversation_id": "conv_123",
                "current_delta_chunks": [],
                "current_item_chunks": [],
                "current_delta_type": "text",
            },
        )

        # Should have response.done
        assert len(result["response"]) == 1
        assert result["response"][0]["type"] == "response.done"
        assert result["response"][0]["response"]["status"] == "completed"

        # State should be reset
        assert result["current_output_item_id"] is None
        assert result["current_response_id"] is None
        assert result["current_delta_type"] is None

    def test_event_id_uniqueness(self):
        """Test that all event_ids are unique"""
        config = BedrockRealtimeConfig()
        logging_obj = MagicMock()
        logging_obj.litellm_trace_id = "trace_123"

        # Create a sequence of messages
        content_start = {
            "event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}
        }
        text_output1 = {"event": {"textOutput": {"content": "Hello"}}}
        text_output2 = {"event": {"textOutput": {"content": " world"}}}

        all_events = []
        state = {
            "session_configuration_request": json.dumps({"configured": True}),
            "current_output_item_id": None,
            "current_response_id": None,
            "current_conversation_id": None,
            "current_delta_chunks": [],
            "current_item_chunks": [],
            "current_delta_type": None,
        }

        # Process all messages
        for msg in [content_start, text_output1, text_output2]:
            result = config.transform_realtime_response(
                json.dumps(msg),
                "amazon.nova-sonic-v1:0",
                logging_obj,
                realtime_response_transform_input=state,
            )
            all_events.extend(result["response"])
            # Update state for next iteration
            state.update(
                {
                    "current_output_item_id": result["current_output_item_id"],
                    "current_response_id": result["current_response_id"],
                    "current_conversation_id": result["current_conversation_id"],
                    "current_delta_chunks": result["current_delta_chunks"],
                    "current_delta_type": result["current_delta_type"],
                }
            )

        # Check all event_ids are unique
        event_ids = [event["event_id"] for event in all_events if "event_id" in event]
        assert len(event_ids) == len(set(event_ids)), "Event IDs should be unique"

    def test_response_id_consistency(self):
        """Test that response_id remains consistent across related events"""
        config = BedrockRealtimeConfig()
        logging_obj = MagicMock()
        logging_obj.litellm_trace_id = "trace_123"

        # Create a sequence of messages
        content_start = {
            "event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}
        }
        text_output = {"event": {"textOutput": {"content": "Hello"}}}

        all_events = []
        state = {
            "session_configuration_request": json.dumps({"configured": True}),
            "current_output_item_id": None,
            "current_response_id": None,
            "current_conversation_id": None,
            "current_delta_chunks": [],
            "current_item_chunks": [],
            "current_delta_type": None,
        }

        # Process messages
        for msg in [content_start, text_output]:
            result = config.transform_realtime_response(
                json.dumps(msg),
                "amazon.nova-sonic-v1:0",
                logging_obj,
                realtime_response_transform_input=state,
            )
            all_events.extend(result["response"])
            state.update(
                {
                    "current_output_item_id": result["current_output_item_id"],
                    "current_response_id": result["current_response_id"],
                    "current_conversation_id": result["current_conversation_id"],
                    "current_delta_chunks": result["current_delta_chunks"],
                    "current_delta_type": result["current_delta_type"],
                }
            )

        # Check all response_ids are the same
        response_ids = [
            event["response_id"] for event in all_events if "response_id" in event
        ]
        assert len(set(response_ids)) == 1, "Response IDs should be consistent"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
