import json
from collections.abc import AsyncIterator, Mapping
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, MagicMock

import httpx
import pytest

import litellm
from litellm.llms.bedrock.chat import BedrockConverseLLM
from litellm.llms.bedrock.chat.converse_handler import make_sync_call
from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from tests._support.stream_chunk_size import DEFAULT_CHUNKING_REQUESTS, ROUTER_CHUNK_SIZE_CASES, keys_at_every_depth


def test_encode_model_id_with_inference_profile():
    """
    Test instance profile is properly encoded when used as a model
    """
    test_model = "arn:aws:bedrock:us-east-1:12345678910:application-inference-profile/ujdtmcirjhevpi"
    expected_model = "arn%3Aaws%3Abedrock%3Aus-east-1%3A12345678910%3Aapplication-inference-profile%2Fujdtmcirjhevpi"
    bedrock_converse_llm = BedrockConverseLLM()
    returned_model = bedrock_converse_llm.encode_model_id(test_model)
    assert expected_model == returned_model


class TestBedrockRegionInModelPath:
    """
    Tests for region extraction from bedrock/{region}/{model} path format.

    When a user passes model="bedrock/ap-northeast-1/moonshotai.kimi-k2.5",
    get_llm_provider strips "bedrock/" and passes "ap-northeast-1/moonshotai.kimi-k2.5"
    to the converse handler. The handler must:
    1. Strip the region from modelId (so AWS gets "moonshotai.kimi-k2.5", not "ap-northeast-1%2Fmoonshotai.kimi-k2.5")
    2. Use the extracted region as aws_region_name for the API call
    """

    @pytest.mark.parametrize(
        "model,expected_model_id,expected_region",
        [
            # Region embedded in path — both modelId and region must be extracted
            (
                "ap-northeast-1/moonshotai.kimi-k2.5",
                "moonshotai.kimi-k2.5",
                "ap-northeast-1",
            ),
            (
                "us-east-1/moonshotai.kimi-k2.5",
                "moonshotai.kimi-k2.5",
                "us-east-1",
            ),
            (
                "us-west-2/anthropic.claude-haiku-4-5-20251001-v1:0",
                "anthropic.claude-haiku-4-5-20251001-v1%3A0",
                "us-west-2",
            ),
            # No region in path — modelId unchanged, no region injected
            (
                "moonshotai.kimi-k2.5",
                "moonshotai.kimi-k2.5",
                None,
            ),
            # Cross-region inference prefix (us., eu., ap.) — not a region path segment
            (
                "us.anthropic.claude-haiku-4-5-20251001-v1:0",
                "us.anthropic.claude-haiku-4-5-20251001-v1%3A0",
                None,
            ),
        ],
    )
    def test_region_and_model_id_extraction(
        self, model: str, expected_model_id: str, expected_region: str | None
    ) -> None:
        """
        Verify that completion() correctly extracts both modelId and aws_region_name
        from the bedrock/{region}/{model} path format.
        """
        bedrock_converse_llm = BedrockConverseLLM()
        optional_params: dict = {}

        # Simulate the modelId + region extraction logic from completion()
        _model_for_id = model
        _stripped = _model_for_id
        for rp in ["bedrock/converse/", "bedrock/", "converse/"]:
            if _stripped.startswith(rp):
                _stripped = _stripped[len(rp) :]
                break

        _region_from_model = None
        _potential_region = _stripped.split("/", 1)[0]
        if _potential_region in _get_all_bedrock_regions() and "/" in _stripped:
            _region_from_model = _potential_region
            _stripped = _stripped.split("/", 1)[1]
            _model_for_id = _stripped

        for _nova_prefix in ["nova-2/", "nova/"]:
            if _stripped.startswith(_nova_prefix):
                _model_for_id = _model_for_id.replace(_nova_prefix, "", 1)
                break

        model_id = bedrock_converse_llm.encode_model_id(model_id=_model_for_id)
        if _region_from_model is not None and "aws_region_name" not in optional_params:
            optional_params["aws_region_name"] = _region_from_model

        assert (
            model_id == expected_model_id
        ), f"modelId mismatch for {model!r}: got {model_id!r}, expected {expected_model_id!r}"
        assert (
            optional_params.get("aws_region_name") == expected_region
        ), (
            f"region mismatch for {model!r}: "
            f"got {optional_params.get('aws_region_name')!r}, expected {expected_region!r}"
        )

    def test_explicit_aws_region_name_not_overridden(self):
        """
        If aws_region_name is already set in optional_params, the region in the
        model path must NOT override it.
        """
        bedrock_converse_llm = BedrockConverseLLM()
        optional_params = {"aws_region_name": "eu-west-1"}
        model = "ap-northeast-1/moonshotai.kimi-k2.5"

        _model_for_id = model
        _stripped = model
        _region_from_model = None
        _potential_region = _stripped.split("/", 1)[0]
        if _potential_region in _get_all_bedrock_regions() and "/" in _stripped:
            _region_from_model = _potential_region
            _stripped = _stripped.split("/", 1)[1]
            _model_for_id = _stripped

        model_id = bedrock_converse_llm.encode_model_id(model_id=_model_for_id)
        if _region_from_model is not None and "aws_region_name" not in optional_params:
            optional_params["aws_region_name"] = _region_from_model

        # modelId is still correctly stripped
        assert model_id == "moonshotai.kimi-k2.5"
        # explicitly set region is preserved
        assert optional_params["aws_region_name"] == "eu-west-1"


def _stream_completion_with_spied_iter_bytes(model: str, stream_chunk_size: int | None = None) -> MagicMock:
    mock_response: Final = MagicMock()
    mock_response.status_code = 200
    mock_response.iter_bytes = MagicMock(return_value=iter([]))
    client: Final = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    litellm.completion(
        model=model,
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
        stream_chunk_size=stream_chunk_size,
    )
    return mock_response.iter_bytes


def test_make_sync_call_does_not_rechunk_stream_by_default():
    """Re-chunking the event stream into fixed 1024-byte blocks holds small
    early events in httpx's ByteChunker until 1024 bytes accumulate, delaying
    time-to-first-chunk by the whole generation when Bedrock trickles bytes
    (e.g. buffered tool-use streams)."""
    response = MagicMock()
    response.status_code = 200
    client = MagicMock()
    client.post = MagicMock(return_value=response)

    make_sync_call(
        client=client,
        api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
        headers={},
        data="{}",
        model="anthropic.claude-sonnet-4-6",
        messages=[],
        logging_obj=MagicMock(),
    )

    response.iter_bytes.assert_called_once_with(chunk_size=None)


def test_make_sync_call_honors_explicit_stream_chunk_size():
    response = MagicMock()
    response.status_code = 200
    client = MagicMock()
    client.post = MagicMock(return_value=response)

    make_sync_call(
        client=client,
        api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream",
        headers={},
        data="{}",
        model="anthropic.claude-sonnet-4-6",
        messages=[],
        logging_obj=MagicMock(),
        stream_chunk_size=2048,
    )

    response.iter_bytes.assert_called_once_with(chunk_size=2048)


def _converse_response_body() -> dict:
    return {
        "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}},
        "stopReason": "end_turn",
        "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
    }


def test_converse_completion_forwards_bedrock_response_headers():
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json = MagicMock(return_value=_converse_response_body())
    mock_response.text = json.dumps(_converse_response_body())
    mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-123"})
    client = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    response = litellm.completion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
    )

    assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-123"


def test_converse_streaming_forwards_bedrock_response_headers():
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.iter_bytes = MagicMock(return_value=iter([]))
    mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-456"})
    client = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    response = litellm.completion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
    )

    assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-456"


@pytest.mark.asyncio
async def test_async_converse_completion_forwards_bedrock_response_headers():
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json = MagicMock(return_value=_converse_response_body())
    mock_response.text = json.dumps(_converse_response_body())
    mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-abc"})
    client = AsyncHTTPHandler()
    client.post = AsyncMock(return_value=mock_response)

    response = await litellm.acompletion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
    )

    assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-abc"


@pytest.mark.asyncio
async def test_async_converse_streaming_forwards_bedrock_response_headers():
    async def _no_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]:
        return
        yield b""

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.aiter_bytes = _no_bytes
    mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-def"})
    client = AsyncHTTPHandler()
    client.post = AsyncMock(return_value=mock_response)

    response = await litellm.acompletion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
    )

    assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-def"


def test_completion_plumbs_stream_chunk_size_through_converse() -> None:
    iter_bytes_spy = _stream_completion_with_spied_iter_bytes(
        model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
    )
    iter_bytes_spy.assert_called_once_with(chunk_size=None)

    iter_bytes_spy = _stream_completion_with_spied_iter_bytes(
        model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
        stream_chunk_size=2048,
    )
    iter_bytes_spy.assert_called_once_with(chunk_size=2048)


def _stream_converse_completion_with_spied_client(**request: object) -> tuple[MagicMock, MagicMock]:
    mock_response: Final = MagicMock()
    mock_response.status_code = 200
    mock_response.iter_bytes = MagicMock(return_value=iter([]))
    client: Final = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)

    litellm.completion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
        **request,
    )
    return mock_response.iter_bytes, client.post


def test_completion_stream_chunk_size_reaches_iter_bytes_but_not_converse_body() -> None:
    iter_bytes_spy, post_spy = _stream_converse_completion_with_spied_client(stream_chunk_size=64)

    iter_bytes_spy.assert_called_once_with(chunk_size=64)
    data: Final = post_spy.call_args.kwargs["data"]
    assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data


@pytest.mark.parametrize("request_kwargs", DEFAULT_CHUNKING_REQUESTS)
def test_completion_uses_default_chunking_unless_a_valid_size_is_requested(
    request_kwargs: Mapping[str, object],
) -> None:
    iter_bytes_spy, _ = _stream_converse_completion_with_spied_client(**request_kwargs)

    iter_bytes_spy.assert_called_once_with(chunk_size=None)


async def _astream_converse_completion_with_spied_client(**request: object) -> tuple[MagicMock, AsyncMock]:
    async def _no_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]:
        return
        yield b""

    mock_response: Final = MagicMock()
    mock_response.status_code = 200
    mock_response.aiter_bytes = MagicMock(return_value=_no_bytes())
    aiter_bytes_spy: Final = mock_response.aiter_bytes
    client: Final = AsyncHTTPHandler()
    client.post = AsyncMock(return_value=mock_response)

    await litellm.acompletion(
        model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
        aws_access_key_id="fake",
        aws_secret_access_key="fake",
        aws_region_name="us-east-1",
        **request,
    )
    return aiter_bytes_spy, client.post


@pytest.mark.asyncio
async def test_acompletion_stream_chunk_size_reaches_aiter_bytes_but_not_converse_body() -> None:
    aiter_bytes_spy, post_spy = await _astream_converse_completion_with_spied_client(stream_chunk_size=64)

    aiter_bytes_spy.assert_called_once_with(chunk_size=64)
    data: Final = post_spy.call_args.kwargs["data"]
    assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data


@pytest.mark.asyncio
@pytest.mark.parametrize("request_kwargs", DEFAULT_CHUNKING_REQUESTS)
async def test_acompletion_uses_default_chunking_unless_a_valid_size_is_requested(
    request_kwargs: Mapping[str, object],
) -> None:
    aiter_bytes_spy, _ = await _astream_converse_completion_with_spied_client(**request_kwargs)

    aiter_bytes_spy.assert_called_once_with(chunk_size=None)


CONVERSE_DEPLOYMENT: Final = MappingProxyType(
    {
        "model": "bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
        "aws_access_key_id": "fake",
        "aws_secret_access_key": "fake",
        "aws_region_name": "us-east-1",
    }
)


@pytest.mark.parametrize("deployment_extras,expected_chunk_size", ROUTER_CHUNK_SIZE_CASES)
def test_router_deployment_stream_chunk_size_reaches_iter_bytes(
    deployment_extras: Mapping[str, object], expected_chunk_size: int | None
) -> None:
    mock_response: Final = MagicMock()
    mock_response.status_code = 200
    mock_response.iter_bytes = MagicMock(return_value=iter([]))
    client: Final = HTTPHandler()
    client.post = MagicMock(return_value=mock_response)
    router: Final = litellm.Router(
        model_list=[{"model_name": "converse-chunked", "litellm_params": {**CONVERSE_DEPLOYMENT, **deployment_extras}}]
    )

    router.completion(
        model="converse-chunked",
        messages=[{"role": "user", "content": "hi"}],
        stream=True,
        client=client,
    )

    mock_response.iter_bytes.assert_called_once_with(chunk_size=expected_chunk_size)
    data: Final = client.post.call_args.kwargs["data"]
    assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data


@pytest.mark.parametrize("stream", [True, False], ids=["stream", "non_stream"])
def test_converse_rejects_non_int_stream_chunk_size_before_calling_bedrock(stream: bool) -> None:
    send: Final = MagicMock(return_value=httpx.Response(200))
    client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(send)))

    with pytest.raises(litellm.BadRequestError):
        litellm.completion(
            model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
            messages=[{"role": "user", "content": "hi"}],
            stream=stream,
            client=client,
            aws_access_key_id="fake",
            aws_secret_access_key="fake",
            aws_region_name="us-east-1",
            stream_chunk_size="sixty-four",
        )

    send.assert_not_called()


def _bedrock_error_response(status_code: int, request_id: str) -> httpx.Response:
    return httpx.Response(
        status_code=status_code,
        headers={
            "x-amzn-RequestId": request_id,
            "x-amzn-ErrorType": "InternalServerException",
        },
        text=json.dumps({"message": "Amazon Bedrock is unable to process your request."}),
        request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"),
    )


def test_converse_completion_error_forwards_bedrock_response_headers():
    error_response = _bedrock_error_response(500, "req-err-123")
    client = HTTPHandler()
    client.post = MagicMock(
        side_effect=httpx.HTTPStatusError(
            "server error",
            request=error_response.request,
            response=error_response,
        )
    )

    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
        litellm.completion(
            model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
            messages=[{"role": "user", "content": "hi"}],
            client=client,
            aws_access_key_id="fake",
            aws_secret_access_key="fake",
            aws_region_name="us-east-1",
        )

    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-123"


@pytest.mark.asyncio
async def test_async_converse_completion_error_forwards_bedrock_response_headers():
    error_response = _bedrock_error_response(500, "req-err-456")
    client = AsyncHTTPHandler()
    client.post = AsyncMock(
        side_effect=httpx.HTTPStatusError(
            "server error",
            request=error_response.request,
            response=error_response,
        )
    )

    with pytest.raises(litellm.ServiceUnavailableError) as exc_info:
        await litellm.acompletion(
            model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0",
            messages=[{"role": "user", "content": "hi"}],
            client=client,
            aws_access_key_id="fake",
            aws_secret_access_key="fake",
            aws_region_name="us-east-1",
        )

    assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-456"