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

import pytest

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

from litellm.litellm_core_utils.prompt_templates.common_utils import (
    add_system_prompt_to_messages,
    get_file_ids_from_messages,
    get_format_from_file_id,
    handle_any_messages_to_chat_completion_str_messages_conversion,
    split_concatenated_json_objects,
    update_messages_with_model_file_ids,
)


def test_get_format_from_file_id():
    unified_file_id = (
        "litellm_proxy:application/pdf;unified_id,cbbe3534-8bf8-4386-af00-f5f6b7e370bf"
    )

    format = get_format_from_file_id(unified_file_id)

    assert format == "application/pdf"


def test_update_messages_with_model_file_ids():
    file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCxmYzdmMmVhNS0wZjUwLTQ5ZjYtODljMS03ZTZhNTRiMTIxMzg"
    model_id = "my_model_id"
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "What is in this recording?"},
                {
                    "type": "file",
                    "file": {
                        "file_id": file_id,
                    },
                },
            ],
        },
    ]

    model_file_id_mapping = {file_id: {"my_model_id": "provider_file_id"}}

    updated_messages = update_messages_with_model_file_ids(
        messages, model_id, model_file_id_mapping
    )

    assert updated_messages == [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "What is in this recording?"},
                {
                    "type": "file",
                    "file": {
                        "file_id": "provider_file_id",
                        "format": "application/pdf",
                    },
                },
            ],
        }
    ]


def test_handle_any_messages_to_chat_completion_str_messages_conversion_list():
    # Test with list of messages
    messages = [
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hi there"},
    ]
    result = handle_any_messages_to_chat_completion_str_messages_conversion(messages)
    assert len(result) == 2
    assert result[0] == messages[0]
    assert result[1] == messages[1]


def test_handle_any_messages_to_chat_completion_str_messages_conversion_list_infinite_loop():
    # Test that list handling doesn't cause infinite recursion
    messages = [
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hi there"},
    ]
    # This should complete without stack overflow
    result = handle_any_messages_to_chat_completion_str_messages_conversion(messages)
    assert len(result) == 2
    assert result[0] == messages[0]
    assert result[1] == messages[1]


def test_handle_any_messages_to_chat_completion_str_messages_conversion_dict():
    # Test with single dictionary message
    message = {"role": "user", "content": "Hello"}
    result = handle_any_messages_to_chat_completion_str_messages_conversion(message)
    assert len(result) == 1
    assert result[0]["input"] == json.dumps(message)


def test_handle_any_messages_to_chat_completion_str_messages_conversion_str():
    # Test with string message
    message = "Hello"
    result = handle_any_messages_to_chat_completion_str_messages_conversion(message)
    assert len(result) == 1
    assert result[0]["input"] == message


def test_handle_any_messages_to_chat_completion_str_messages_conversion_other():
    # Test with non-string/dict/list type
    message = 123
    result = handle_any_messages_to_chat_completion_str_messages_conversion(message)
    assert len(result) == 1
    assert result[0]["input"] == "123"


def test_handle_any_messages_to_chat_completion_str_messages_conversion_complex():
    # Test with complex nested structure
    message = {
        "role": "user",
        "content": {"text": "Hello", "metadata": {"timestamp": "2024-01-01"}},
    }
    result = handle_any_messages_to_chat_completion_str_messages_conversion(message)
    assert len(result) == 1
    assert result[0]["input"] == json.dumps(message)


def test_add_system_prompt_to_messages_prepend():
    """Adds system prompt at beginning when no system message exists."""
    messages = [
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hi there"},
    ]
    result = add_system_prompt_to_messages(messages, "You are a helpful assistant.")
    assert result == [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hi there"},
    ]


def test_add_system_prompt_to_messages_empty_prompt_unchanged():
    """Returns messages unchanged when system_prompt is empty."""
    messages = [{"role": "user", "content": "Hello"}]
    assert add_system_prompt_to_messages(messages, "") == messages
    assert add_system_prompt_to_messages(messages, None) == messages


def test_add_system_prompt_to_messages_merge_with_first_system():
    """Merges new prompt into first system message when merge_with_first_system=True."""
    messages = [
        {"role": "system", "content": "Existing system prompt."},
        {"role": "user", "content": "Hello"},
    ]
    result = add_system_prompt_to_messages(
        messages, "You are helpful.", merge_with_first_system=True
    )
    assert result == [
        {"role": "system", "content": "You are helpful.\n\nExisting system prompt."},
        {"role": "user", "content": "Hello"},
    ]


def test_add_system_prompt_to_messages_merge_with_first_system_adds_new_when_no_system():
    """When merge_with_first_system=True but no system message, adds new one at start."""
    messages = [{"role": "user", "content": "Hello"}]
    result = add_system_prompt_to_messages(
        messages, "You are helpful.", merge_with_first_system=True
    )
    assert result == [
        {"role": "system", "content": "You are helpful."},
        {"role": "user", "content": "Hello"},
    ]


def test_add_system_prompt_to_messages_empty_list():
    """Adds system prompt to empty messages list."""
    result = add_system_prompt_to_messages([], "You are helpful.")
    assert result == [{"role": "system", "content": "You are helpful."}]


def test_convert_prefix_message_to_non_prefix_messages():
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        convert_prefix_message_to_non_prefix_messages,
    )

    messages = [
        {"role": "assistant", "content": "value", "prefix": True},
    ]
    result = convert_prefix_message_to_non_prefix_messages(messages)
    assert result == [
        {
            "role": "system",
            "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ",
        },
        {"role": "assistant", "content": "value"},
    ]


# ── split_concatenated_json_objects tests ──


def test_split_concatenated_json_single_object():
    """A single valid JSON object is returned as a one-element list."""
    result = split_concatenated_json_objects('{"location": "Boston"}')
    assert result == [{"location": "Boston"}]


def test_split_concatenated_json_multiple_objects():
    """
    Multiple JSON objects concatenated without separators are split correctly.
    This is the exact pattern from issue #20543 where Bedrock Claude Sonnet 4.5
    returns concatenated JSON in a single tool call arguments string.
    """
    raw = (
        '{"command": ["curl", "-i", "http://localhost:9009"]}'
        '{"command": ["curl", "-i", "http://localhost:9009/robots.txt"]}'
        '{"command": ["curl", "-i", "http://localhost:9009/sitemap.xml"]}'
    )
    result = split_concatenated_json_objects(raw)
    assert len(result) == 3
    assert result[0] == {"command": ["curl", "-i", "http://localhost:9009"]}
    assert result[1] == {"command": ["curl", "-i", "http://localhost:9009/robots.txt"]}
    assert result[2] == {"command": ["curl", "-i", "http://localhost:9009/sitemap.xml"]}


def test_split_concatenated_json_with_whitespace():
    """Objects separated by whitespace are handled correctly."""
    raw = '{"a": 1}  {"b": 2}\n{"c": 3}'
    result = split_concatenated_json_objects(raw)
    assert len(result) == 3
    assert result[0] == {"a": 1}
    assert result[1] == {"b": 2}
    assert result[2] == {"c": 3}


def test_split_concatenated_json_empty_string():
    """Empty or whitespace-only strings return an empty list."""
    assert split_concatenated_json_objects("") == []
    assert split_concatenated_json_objects("   ") == []


def test_split_concatenated_json_non_dict_value():
    """Non-dict JSON values (e.g. arrays, strings) are replaced with {}."""
    result = split_concatenated_json_objects("[1, 2, 3]")
    assert result == [{}]


def test_split_concatenated_json_invalid_raises():
    """Completely invalid JSON raises JSONDecodeError."""
    with pytest.raises(json.JSONDecodeError):
        split_concatenated_json_objects("not json at all")


# ---------------------------------------------------------------------------
# Regression tests for non-OpenAI file content blocks.
#
# `type: "file"` is a public content-block discriminator. Several producers
# (LangChain v1, provider-native shapes, custom user code) emit blocks with
# `type: "file"` but without the OpenAI Chat Completions `file` sub-dict.
# The discovery helpers below are used unconditionally inside
# `AnthropicConfig.validate_environment`, so any crash there surfaces as a
# `500 InternalServerError` before the request is even dispatched.
# ---------------------------------------------------------------------------


def test_get_file_ids_from_messages_skips_langchain_v1_file_block():
    """A LangChain v1 standardized file block must not crash file-id discovery."""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "summarise this PDF"},
                # LangChain v1 shape produced by `_normalize_messages`.
                # No `file` sub-dict: the discriminator is `type: "file"` but
                # the payload lives on `base64`/`mime_type` siblings.
                {
                    "type": "file",
                    "id": "lc_1",
                    "base64": "JVBERi0xLjQK",
                    "mime_type": "application/pdf",
                    "extras": {"file_format": "application/pdf"},
                },
            ],
        }
    ]

    assert get_file_ids_from_messages(messages) == []


def test_get_file_ids_from_messages_still_extracts_from_openai_shape():
    """Well-formed OpenAI file blocks still yield their file_id."""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "what is this?"},
                {"type": "file", "file": {"file_id": "file-abc"}},
            ],
        }
    ]

    assert get_file_ids_from_messages(messages) == ["file-abc"]


def test_get_file_ids_from_messages_mixed_shapes():
    """Mixed OpenAI and non-OpenAI file blocks: extract from the former,
    ignore the latter."""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "file", "file": {"file_id": "file-keep"}},
                {
                    "type": "file",
                    "id": "lc_2",
                    "base64": "AAA",
                    "mime_type": "application/pdf",
                },
            ],
        }
    ]

    assert get_file_ids_from_messages(messages) == ["file-keep"]


def test_get_file_ids_from_messages_file_field_not_dict():
    """`file` set to a non-dict value (e.g. stringified payload) must not crash."""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "file", "file": "unexpectedly-a-string"},
            ],
        }
    ]

    assert get_file_ids_from_messages(messages) == []


def test_update_messages_with_model_file_ids_skips_non_openai_file_blocks():
    """`update_messages_with_model_file_ids` is also called on user content
    before provider dispatch. It must tolerate non-OpenAI file blocks the same
    way."""
    langchain_v1_block = {
        "type": "file",
        "id": "lc_3",
        "base64": "AAA",
        "mime_type": "application/pdf",
    }
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "hello"},
                langchain_v1_block,
            ],
        }
    ]

    updated = update_messages_with_model_file_ids(messages, "model-1", {})

    # Messages pass through unchanged when there is no `file` sub-dict to remap.
    assert updated == messages


# Reusable fixture (decodes to: litellm_proxy:application/pdf;unified_id,...;
# target_model_names,gpt-4o;llm_output_file_id,file-ECBPW7ML9g7XHdwGgUPZaM;
# llm_output_file_model_id,...)
UNIFIED_FILE_ID_B64 = (
    "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0"
    "LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1f"
    "b3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRf"
    "ZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFk"
    "MDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My"
)


def test_update_messages_with_model_file_ids_decodes_unified_id_when_mapping_empty():
    """When the mapping is empty (e.g. multi-replica cache miss), the function
    must decode the base64-encoded unified file id and substitute the embedded
    llm_output_file_id — mirroring the Responses-API sibling."""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "What is in this recording?"},
                {
                    "type": "file",
                    "file": {
                        "file_id": UNIFIED_FILE_ID_B64,
                        "format": "audio/wav",
                    },
                },
            ],
        }
    ]

    updated = update_messages_with_model_file_ids(messages, "any-model-id", {})

    assert updated[0]["content"][1]["file"]["file_id"] == "file-ECBPW7ML9g7XHdwGgUPZaM"
    # Customer-supplied format is preserved (this is the field whose absence
    # the misleading error message used to complain about).
    assert updated[0]["content"][1]["file"]["format"] == "audio/wav"


def test_update_messages_with_model_file_ids_mapping_takes_precedence_over_decode():
    """When both mapping and decode would resolve, the mapping must win
    (preserves per-deployment routing precision)."""
    mapping = {UNIFIED_FILE_ID_B64: {"model-A": "mapped-provider-file-id"}}
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "file",
                    "file": {
                        "file_id": UNIFIED_FILE_ID_B64,
                        "format": "application/pdf",
                    },
                },
            ],
        }
    ]

    updated = update_messages_with_model_file_ids(messages, "model-A", mapping)

    assert updated[0]["content"][0]["file"]["file_id"] == "mapped-provider-file-id"


def test_update_messages_with_model_file_ids_non_unified_passes_through():
    """A raw provider id (e.g. gs:// URI or a random string) must be left
    untouched when the mapping doesn't resolve it. The decode fallback must
    not corrupt non-unified ids."""
    raw_id = "gs://my-bucket/uploads/abc-123.wav"
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "file", "file": {"file_id": raw_id, "format": "audio/wav"}},
            ],
        }
    ]

    updated = update_messages_with_model_file_ids(messages, "model-A", {})

    assert updated[0]["content"][0]["file"]["file_id"] == raw_id


def test_update_messages_with_model_file_ids_mapping_miss_falls_back_to_decode():
    """A mapping that exists but doesn't contain this file_id should still
    trigger the decode fallback — covers the case where the hook resolved
    *some* ids but not this one."""
    other_id = "some-other-file-id"
    mapping = {other_id: {"model-A": "other-provider-id"}}
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "file",
                    "file": {"file_id": UNIFIED_FILE_ID_B64, "format": "audio/wav"},
                },
            ],
        }
    ]

    updated = update_messages_with_model_file_ids(messages, "model-A", mapping)

    assert updated[0]["content"][0]["file"]["file_id"] == "file-ECBPW7ML9g7XHdwGgUPZaM"


def test_update_messages_with_model_file_ids_tolerates_non_dict_content_items():
    """Content list items aren't always dicts. text_completion forwards
    token-ids (list of ints, or list of list of ints for batch) through
    this path. The function must skip non-dict items instead of indexing
    into them."""
    messages_token_ids = [{"role": "user", "content": [15496, 995]}]
    messages_token_ids_batch = [{"role": "user", "content": [[15496, 995], [9906, 0]]}]

    # Both should pass through unchanged without raising.
    assert (
        update_messages_with_model_file_ids(messages_token_ids, "model-A", {})
        == messages_token_ids
    )
    assert (
        update_messages_with_model_file_ids(messages_token_ids_batch, "model-A", {})
        == messages_token_ids_batch
    )


class TestExtractFileDataBareStr:
    """``extract_file_data`` used to accept bare ``str`` values and ``open()``
    them server-side. When the helper runs inside a proxy request handler the
    value is attacker-controlled, so the open() call was a textbook arbitrary
    local file read. Lock the new contract: bare ``str`` is rejected with a
    clear migration message; ``pathlib.Path`` is still accepted for SDK
    ergonomics because it's a Python-level type that HTTP form values can't
    fabricate."""

    def test_rejects_bare_str(self):
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            extract_file_data,
        )

        with pytest.raises(ValueError, match="does not accept bare str inputs"):
            extract_file_data("/etc/passwd")

    def test_accepts_pathlib_path(self):
        import tempfile
        from pathlib import Path

        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            extract_file_data,
        )

        content = b"hello"
        with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
            f.write(content)
            tmp_path = Path(f.name)

        try:
            extracted = extract_file_data(tmp_path)
            assert extracted.get("content") == content
        finally:
            os.unlink(str(tmp_path))

    def test_accepts_bytes(self):
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            extract_file_data,
        )

        extracted = extract_file_data(b"raw bytes content")
        assert extracted.get("content") == b"raw bytes content"

    def test_accepts_tuple(self):
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            extract_file_data,
        )

        extracted = extract_file_data(("foo.txt", b"raw bytes content"))
        assert extracted.get("filename") == "foo.txt"
        assert extracted.get("content") == b"raw bytes content"


class TestUnpackLegacyDefs:
    """Cover the public ``unpack_legacy_defs`` helper directly so the no-op
    branches (non-dict input, schema with no legacy/OpenAPI defs) are exercised
    without needing a provider-specific entry point.
    """

    @pytest.mark.parametrize(
        "value",
        [None, [], "string-not-a-dict", 42, 1.5, True, set(), tuple()],
    )
    def test_non_dict_returns_unchanged_no_op(self, value):
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            unpack_legacy_defs,
        )

        # Should never raise; returns the input unchanged.
        assert unpack_legacy_defs(value) is value
        assert unpack_legacy_defs(value, copy=True) is value

    def test_dict_without_legacy_defs_is_no_op(self):
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            unpack_legacy_defs,
        )

        schema = {
            "type": "object",
            "properties": {"a": {"$ref": "#/$defs/A"}},
            "$defs": {"A": {"type": "string"}},
        }
        snapshot = json.loads(json.dumps(schema))

        # No `definitions` and no `components.schemas` -> early return, no work.
        out = unpack_legacy_defs(schema)
        assert out is schema
        assert schema == snapshot, "schema mutated despite no legacy defs"

    def test_components_with_no_schemas_block_is_no_op(self):
        """``components`` without a ``schemas`` sub-key must not be popped."""
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            unpack_legacy_defs,
        )

        schema = {
            "type": "object",
            "properties": {"a": {"type": "string"}},
            "components": {"securitySchemes": {"foo": "bar"}},
        }
        snapshot = json.loads(json.dumps(schema))

        unpack_legacy_defs(schema)
        assert schema == snapshot, "components without schemas was incorrectly popped"

    def test_legitimate_schema_within_budget_succeeds(self):
        """A flat schema with many distinct ``$ref``s into small targets must
        inline cleanly under the default budget -- the budget rejects bombs,
        not legitimately-shaped schemas.
        """
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            unpack_legacy_defs,
        )

        n = 200
        schema = {
            "type": "object",
            "properties": {f"f{i}": {"$ref": f"#/definitions/T{i}"} for i in range(n)},
            "definitions": {f"T{i}": {"type": "string"} for i in range(n)},
        }

        out = unpack_legacy_defs(schema)
        assert "definitions" not in out
        for i in range(n):
            assert out["properties"][f"f{i}"] == {"type": "string"}

    # Schema-bomb amplification vectors. ``max_inlined_bytes`` is the universal
    # measure of expansion: every other dimension (ref count, node count,
    # scalar size) reduces to bytes-on-the-wire, so a single byte budget
    # closes all three vectors at once.

    def test_rejects_fan_out_bomb(self):
        """Each level multiplies refs (cycle detection only stops re-entry
        along the *same* path). Must trip the byte budget."""
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            unpack_legacy_defs,
        )

        depth, fanout = 12, 2  # 2**12 = 4096 leaves
        definitions = {
            f"L{i}": {
                "type": "object",
                "properties": {
                    f"x{j}": {"$ref": f"#/definitions/L{i + 1}"} for j in range(fanout)
                },
            }
            for i in range(depth)
        }
        definitions[f"L{depth}"] = {"type": "string"}
        schema = {
            "type": "object",
            "properties": {"root": {"$ref": "#/definitions/L0"}},
            "definitions": definitions,
        }

        with pytest.raises(ValueError, match="byte budget"):
            unpack_legacy_defs(schema, max_inlined_bytes=100_000)

    def test_rejects_target_amplification_bomb(self):
        """Few refs each deep-copying one large target -- bounded total
        expanded bytes catches it even though ref count is small."""
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            unpack_legacy_defs,
        )

        big = {
            "type": "object",
            "properties": {f"p{i}": {"type": "string"} for i in range(100)},
        }
        schema = {
            "type": "object",
            "properties": {f"r{i}": {"$ref": "#/definitions/Big"} for i in range(50)},
            "definitions": {"Big": big},
        }

        with pytest.raises(ValueError, match="byte budget"):
            unpack_legacy_defs(schema, max_inlined_bytes=10_000)

    def test_rejects_scalar_byte_amplification_bomb(self):
        """Many ``$ref``s to a target containing one large scalar (e.g. a
        long ``description``, ``const`` value, or ``enum`` entry). A
        node-counter would treat this as 1 node per resolution and miss it;
        a byte budget catches the actual wire-size amplification.
        """
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            unpack_legacy_defs,
        )

        big_description = "x" * 100_000  # 100KB string
        schema = {
            "type": "object",
            "properties": {f"r{i}": {"$ref": "#/definitions/Big"} for i in range(50)},
            "definitions": {
                "Big": {"type": "string", "description": big_description},
            },
        }
        # 50 refs * ~100KB string == ~5MB cumulative; 1MB budget trips.
        with pytest.raises(ValueError, match="byte budget"):
            unpack_legacy_defs(schema, max_inlined_bytes=1_000_000)

    def test_budget_does_not_trip_for_legitimate_large_schema(self):
        """An OpenAPI-derived tool with ~50 small targets must inline cleanly
        under the default ``max_inlined_bytes`` budget."""
        from litellm.litellm_core_utils.prompt_templates.common_utils import (
            unpack_legacy_defs,
        )

        schema = {
            "type": "object",
            "properties": {
                f"r{i}": {"$ref": f"#/components/schemas/T{i}"} for i in range(50)
            },
            "components": {
                "schemas": {
                    f"T{i}": {
                        "type": "object",
                        "properties": {f"p{j}": {"type": "string"} for j in range(5)},
                    }
                    for i in range(50)
                }
            },
        }

        out = unpack_legacy_defs(schema)
        assert "components" not in out
        assert out["properties"]["r0"]["properties"]["p0"] == {"type": "string"}
