"""Standard LangChain interface tests"""

import base64
import warnings
from typing import Any, Literal, Optional, Type

import httpx
import pytest
import yaml  # type: ignore[import-untyped]
from langchain_core.exceptions import OutputParserException
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessageChunk,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_core.tools import BaseTool, tool
from langchain_tests.integration_tests import ChatModelIntegrationTests
from pydantic import BaseModel, Field
from typing_extensions import Annotated, TypedDict
from vcr import VCR

from langchain_aws import ChatBedrockConverse


class TestBedrockStandard(ChatModelIntegrationTests):
    @property
    def chat_model_class(self) -> Type[BaseChatModel]:
        return ChatBedrockConverse

    @property
    def chat_model_params(self) -> dict:
        return {"model": "us.anthropic.claude-haiku-4-5-20251001-v1:0"}

    @property
    def standard_chat_model_params(self) -> dict:
        return {"temperature": 0, "max_tokens": 100, "stop": []}

    @property
    def supports_image_inputs(self) -> bool:
        return True


class TestBedrockMistralStandard(ChatModelIntegrationTests):
    @property
    def chat_model_class(self) -> Type[BaseChatModel]:
        return ChatBedrockConverse

    @property
    def chat_model_params(self) -> dict:
        return {"model": "mistral.mistral-large-2402-v1:0"}

    @property
    def standard_chat_model_params(self) -> dict:
        return {"temperature": 0, "max_tokens": 100, "stop": []}

    @property
    def has_tool_choice(self) -> bool:
        return False

    @pytest.mark.xfail(reason="Human messages following AI messages not supported.")
    def test_tool_message_histories_list_content(
        self, model: BaseChatModel, my_adder_tool: BaseTool
    ) -> None:
        super().test_tool_message_histories_list_content(model, my_adder_tool)


class TestBedrockNovaStandard(ChatModelIntegrationTests):
    @property
    def chat_model_class(self) -> Type[BaseChatModel]:
        return ChatBedrockConverse

    @property
    def chat_model_params(self) -> dict:
        return {"model": "us.amazon.nova-pro-v1:0"}

    @property
    def standard_chat_model_params(self) -> dict:
        return {"max_tokens": 300, "stop": []}

    @pytest.mark.xfail(reason="Tool choice 'Any' not supported.")
    def test_structured_few_shot_examples(
        self, model: BaseChatModel, my_adder_tool: BaseTool
    ) -> None:
        super().test_structured_few_shot_examples(model, my_adder_tool)

    @pytest.mark.xfail(reason="Human messages following AI messages not supported.")
    def test_tool_message_histories_list_content(
        self, model: BaseChatModel, my_adder_tool: BaseTool
    ) -> None:
        super().test_tool_message_histories_list_content(model, my_adder_tool)


class TestBedrockCohereStandard(ChatModelIntegrationTests):
    @property
    def chat_model_class(self) -> Type[BaseChatModel]:
        return ChatBedrockConverse

    @property
    def chat_model_params(self) -> dict:
        return {"model": "cohere.command-r-plus-v1:0"}

    @property
    def standard_chat_model_params(self) -> dict:
        return {"temperature": 0, "max_tokens": 100, "stop": []}

    @property
    def has_tool_choice(self) -> bool:
        return False

    @pytest.mark.xfail(reason="Cohere models don't support tool_choice.")
    def test_structured_few_shot_examples(
        self, model: BaseChatModel, my_adder_tool: BaseTool
    ) -> None:
        pass

    @pytest.mark.xfail(reason="Cohere models don't support tool_choice.")
    def test_unicode_tool_call_integration(
        self,
        model: BaseChatModel,
        *,
        tool_choice: Optional[str] = None,
        force_tool_call: bool = False,
    ) -> None:
        pass

    @pytest.mark.xfail(reason="Generates invalid tool call.")
    def test_tool_calling_with_no_arguments(self, model: BaseChatModel) -> None:
        pass


class TestBedrockMetaStandard(ChatModelIntegrationTests):
    @property
    def chat_model_class(self) -> Type[BaseChatModel]:
        return ChatBedrockConverse

    @property
    def chat_model_params(self) -> dict:
        return {"model": "us.meta.llama3-2-90b-instruct-v1:0"}

    @property
    def standard_chat_model_params(self) -> dict:
        return {"temperature": 0.1, "max_tokens": 100, "stop": []}

    @property
    def has_tool_choice(self) -> bool:
        return False

    @pytest.mark.xfail(reason="Meta models don't support tool_choice.")
    def test_structured_few_shot_examples(
        self, model: BaseChatModel, my_adder_tool: BaseTool
    ) -> None:
        pass

    @pytest.mark.xfail(reason="Meta models don't support tool_choice.")
    def test_unicode_tool_call_integration(
        self,
        model: BaseChatModel,
        *,
        tool_choice: Optional[str] = None,
        force_tool_call: bool = False,
    ) -> None:
        pass

    # TODO: This needs investigation, if this is a bug with Bedrock or Llama models,
    # but this test consistently seem to return single quoted input values {input: '3'}
    # instead of {input: 3} failing the test. Upon checking with tools with non-numeric
    # inputs, tool calling seems to work as expected with Bedrock and Llama models.
    # Same problem with tool_calling_async, below.
    @pytest.mark.xfail(
        reason="Bedrock Meta models tend to return string values for integer inputs ."
    )
    def test_tool_calling(self, model: BaseChatModel) -> None:
        super().test_tool_calling(model)

    @pytest.mark.xfail(
        reason="Bedrock Meta models tend to return string values for integer inputs ."
    )
    async def test_tool_calling_async(self, model: BaseChatModel) -> None:
        await super().test_tool_calling_async(model)

    @pytest.mark.xfail(reason="Meta models don't support tool_choice.")
    def test_tool_calling_with_no_arguments(self, model: BaseChatModel) -> None:
        pass

    @pytest.mark.xfail(
        reason="Human messages following AI messages not supported by Bedrock."
    )
    def test_tool_message_histories_list_content(
        self, model: BaseChatModel, my_adder_tool: BaseTool
    ) -> None:
        super().test_tool_message_histories_list_content(model, my_adder_tool)


def test_multiple_system_messages_anthropic() -> None:
    model = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", temperature=0
    )

    system1 = SystemMessage(content="You are a helpful assistant.")
    system2 = SystemMessage(content="Always respond in a concise manner.")
    human = HumanMessage(content="Hello")
    response = model.invoke([system1, system2, human])

    assert isinstance(response, AIMessage)
    assert isinstance(response.content, str)


class ClassifyQuery(BaseModel):
    """Classify a query."""

    query_type: Literal["cat", "dog"] = Field(
        description="Classify a query as related to cats or dogs."
    )


def test_structured_output_snake_case() -> None:
    model = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", temperature=0
    )

    chat = model.with_structured_output(ClassifyQuery)
    for chunk in chat.stream("How big are cats?"):
        assert isinstance(chunk, ClassifyQuery)


def test_tool_calling_snake_case() -> None:
    model = ChatBedrockConverse(model="us.anthropic.claude-sonnet-4-5-20250929-v1:0")

    def classify_query(query_type: Literal["cat", "dog"]) -> None:
        pass

    chat = model.bind_tools([classify_query], tool_choice="any")
    response = chat.invoke("How big are cats?")
    assert isinstance(response, AIMessage)
    assert len(response.tool_calls) == 1
    tool_call = response.tool_calls[0]
    assert tool_call["name"] == "classify_query"
    assert tool_call["args"] == {"query_type": "cat"}

    full = None
    for chunk in chat.stream("How big are cats?"):
        full = chunk if full is None else full + chunk  # type: ignore[assignment]
    assert isinstance(full, AIMessageChunk)
    assert len(full.tool_calls) == 1
    tool_call = full.tool_calls[0]
    assert tool_call["name"] == "classify_query"
    assert tool_call["args"] == {"query_type": "cat"}

    # Also test for response metadata, though this is not relevant to tool-calling
    invoke_metadata = response.response_metadata
    stream_metadata = full.response_metadata
    for result in [invoke_metadata, stream_metadata]:
        for expected_key in ["RequestId", "HTTPStatusCode", "HTTPHeaders"]:
            assert result["ResponseMetadata"][expected_key]
        assert isinstance(result["ResponseMetadata"]["RetryAttempts"], int)


def test_tool_calling_camel_case() -> None:
    model = ChatBedrockConverse(model="us.anthropic.claude-sonnet-4-5-20250929-v1:0")

    def classifyQuery(queryType: Literal["cat", "dog"]) -> None:
        pass

    chat = model.bind_tools([classifyQuery], tool_choice="any")
    response = chat.invoke("How big are cats?")
    assert isinstance(response, AIMessage)
    assert len(response.tool_calls) == 1
    tool_call = response.tool_calls[0]
    assert tool_call["name"] == "classifyQuery"
    assert tool_call["args"] == {"queryType": "cat"}

    full = None
    for chunk in chat.stream("How big are cats?"):
        full = chunk if full is None else full + chunk  # type: ignore[assignment]
    assert isinstance(full, AIMessageChunk)
    assert len(full.tool_calls) == 1
    tool_call = full.tool_calls[0]
    assert tool_call["name"] == "classifyQuery"
    assert tool_call["args"] == {"queryType": "cat"}
    assert full.tool_calls[0]["args"] == response.tool_calls[0]["args"]


def test_tool_calling_strict() -> None:
    model = ChatBedrockConverse(model="us.anthropic.claude-sonnet-4-5-20250929-v1:0")

    class GetWeather(BaseModel):
        """Get the current weather in a given location."""

        location: str = Field(description="The city and state, e.g. San Francisco, CA")

    chat = model.bind_tools([GetWeather], strict=True, tool_choice="any")
    response = chat.invoke("What is the weather in Paris?")
    assert isinstance(response, AIMessage)
    assert len(response.tool_calls) == 1
    tool_call = response.tool_calls[0]
    assert tool_call["name"] == "GetWeather"
    assert isinstance(tool_call["args"], dict)
    assert "location" in tool_call["args"]


def test_structured_output_streaming() -> None:
    model = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", temperature=0
    )
    query = (
        "What weighs more, a pound of bricks or a pound of feathers? "
        "Limit your response to 20 words."
    )

    # TypedDict
    class AnswerWithJustification(TypedDict):
        """An answer to the user question along with justification for the answer."""

        answer: Annotated[str, ...]
        justification: Annotated[str, ...]

    chat = model.with_structured_output(AnswerWithJustification)
    chunk_count = 0
    for chunk in chat.stream(query):
        chunk_count = chunk_count + 1
        assert isinstance(chunk, dict)
    assert chunk_count > 1

    # Pydantic
    class AnAnswerWithJustification(BaseModel):
        """An answer to the user question along with justification for the answer."""

        answer: Annotated[str, ...]
        justification: Annotated[str, ...]

    chat = model.with_structured_output(AnAnswerWithJustification)
    chunk_count = 0
    for chunk in chat.stream(query):
        chunk_count = chunk_count + 1
        assert isinstance(chunk, AnAnswerWithJustification)
    assert chunk_count > 1


def test_tool_use_with_cache_point() -> None:
    """Test toolUse with cachepoint to verify cache metrics are being reported.

    This test creates tools with a length exceeding 1024 tokens to ensure
    caching is triggered, and verifies the response metrics indicate cache
    activity.

    """
    # Define a large number of tools to exceed 1024 tokens
    tool_classes = []

    # Each tool is simple but we'll define many of them
    for i in range(1, 20):
        # Creating a unique class for each tool
        tool_class_name = f"CalculateTool{i}"

        # Define the class using a closure to properly scope the fields
        def create_tool_class(i: int) -> Type[BaseModel]:
            class ToolClass(BaseModel):
                number1: float = Field(description=f"First number for calculation {i}")
                number2: float = Field(description=f"Second number for calculation {i}")
                operation: Literal["add", "subtract", "multiply", "divide"] = Field(
                    description=f"Operation {i} to perform on the numbers"
                )

            ToolClass.__doc__ = f"A tool to calculate the {i}th mathematical operation"
            return ToolClass

        tool_class = create_tool_class(i)
        tool_class.__name__ = tool_class_name
        tool_classes.append(tool_class)

    # Create the model instance
    model = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", temperature=0
    )

    # Create cache point configuration
    cache_point = ChatBedrockConverse.create_cache_point()

    # Bind tools with cache point
    chat = model.bind_tools(tool_classes + [cache_point], tool_choice="any")

    # Invocation
    response = chat.invoke("What's 5 + 3?")
    assert isinstance(response, AIMessage)
    assert len(response.tool_calls) == 1

    # Verify the response has cache metrics
    assert response.usage_metadata is not None
    input_token_details = response.usage_metadata.get("input_token_details")
    if input_token_details:
        cache_read_input_tokens = input_token_details.get("cache_read", 0)
        cache_write_input_tokens = input_token_details.get("cache_creation", 0)
        assert cache_read_input_tokens + cache_write_input_tokens != 0


@pytest.mark.skip(reason="Needs guardrails setup to run.")
def test_guardrails() -> None:
    params = {
        "region_name": "us-west-2",
        "model": "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
        "temperature": 0,
        "max_tokens": 100,
        "stop": [],
        "guardrail_config": {
            "guardrailIdentifier": "e7esbceow153",
            "guardrailVersion": "1",
            "trace": "enabled",
        },
    }
    chat_model = ChatBedrockConverse(**params)  # type: ignore[arg-type]
    messages = [
        HumanMessage(
            content=[
                "Create a playlist of 2 heavy metal songs.",
                {
                    "guardContent": {
                        "text": {"text": "Only answer with a list of songs."}
                    }
                },
            ]
        )
    ]
    response = chat_model.invoke(messages)

    assert (
        response.content == "Sorry, I can't answer questions about heavy metal music."
    )
    assert response.response_metadata["stopReason"] == "guardrail_intervened"
    assert response.response_metadata["trace"] is not None

    stream = chat_model.stream(messages)
    response = next(stream)
    for chunk in stream:
        response += chunk

    assert (
        response.content[0]["text"]  # type: ignore[index]
        == "Sorry, I can't answer questions about heavy metal music."
    )
    assert response.response_metadata["stopReason"] == "guardrail_intervened"
    assert response.response_metadata["trace"] is not None


@pytest.mark.parametrize(
    "thinking_model",
    [
        "us.anthropic.claude-sonnet-4-20250514-v1:0",
        "us.anthropic.claude-opus-4-20250514-v1:0",
        "us.anthropic.claude-opus-4-1-20250805-v1:0",
        "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
        "us.anthropic.claude-haiku-4-5-20251001-v1:0",
    ],
)
def test_structured_output_tool_choice_not_supported(thinking_model: str) -> None:
    llm = ChatBedrockConverse(model=thinking_model)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        structured_llm = llm.with_structured_output(ClassifyQuery)
        response = structured_llm.invoke("How big are cats? Use the tool.")
    assert len(w) == 0
    assert isinstance(response, ClassifyQuery)

    # Unsupported params
    llm = ChatBedrockConverse(
        model=thinking_model,
        max_tokens=5000,
        additional_model_request_fields={
            "thinking": {"type": "enabled", "budget_tokens": 2000}
        },
    )
    with pytest.warns(match="structured output"):
        structured_llm = llm.with_structured_output(ClassifyQuery)
    response = structured_llm.invoke("How big are cats? Use the tool.")
    assert isinstance(response, ClassifyQuery)

    with pytest.raises(OutputParserException):
        structured_llm.invoke("Hello!")


def test_structured_output_thinking_force_tool_use() -> None:
    # Structured output currently relies on forced tool use, which is not supported
    # when `thinking` is enabled for Claude 3.7. When this test fails, it means that
    # the feature is supported and the workarounds in `with_structured_output` should
    # be removed.

    # Instantiate as convenience for getting client
    llm = ChatBedrockConverse(model="us.anthropic.claude-sonnet-4-5-20250929-v1:0")
    messages = [
        {
            "role": "user",
            "content": [{"text": "Generate a username for Sally with green hair"}],
        }
    ]
    params = {
        "modelId": "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
        "inferenceConfig": {"maxTokens": 5000},
        "toolConfig": {
            "tools": [
                {
                    "toolSpec": {
                        "name": "ClassifyQuery",
                        "description": "Classify a query.",
                        "inputSchema": {
                            "json": {
                                "properties": {
                                    "queryType": {
                                        "description": (
                                            "Classify a query as related to cats or "
                                            "dogs."
                                        ),
                                        "enum": ["cat", "dog"],
                                        "type": "string",
                                    }
                                },
                                "required": ["query_type"],
                                "type": "object",
                            }
                        },
                    }
                }
            ],
            "toolChoice": {"tool": {"name": "ClassifyQuery"}},
        },
        "additionalModelRequestFields": {
            "thinking": {"type": "enabled", "budget_tokens": 2000}
        },
    }
    with pytest.raises(llm.client.exceptions.ValidationException):
        llm.client.converse(messages=messages, **params)


@pytest.mark.default_cassette("test_agent_loop.yaml.gz")
@pytest.mark.vcr
@pytest.mark.parametrize("output_version", ["v0", "v1"])
def test_agent_loop(output_version: Literal["v0", "v1"]) -> None:
    @tool
    def get_weather(location: str) -> str:
        """Get the weather for a location."""
        return "It's sunny."

    llm = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-20250514-v1:0",
        output_version=output_version,
    )
    llm_with_tools = llm.bind_tools([get_weather])
    input_message = HumanMessage("What is the weather in San Francisco, CA?")
    tool_call_message = llm_with_tools.invoke([input_message])
    assert isinstance(tool_call_message, AIMessage)
    tool_calls = tool_call_message.tool_calls
    assert len(tool_calls) == 1
    tool_call = tool_calls[0]
    tool_message = get_weather.invoke(tool_call)
    assert isinstance(tool_message, ToolMessage)
    response = llm_with_tools.invoke(
        [
            input_message,
            tool_call_message,
            tool_message,
        ]
    )
    assert isinstance(response, AIMessage)


@pytest.mark.default_cassette("test_agent_loop_streaming.yaml.gz")
@pytest.mark.vcr
@pytest.mark.parametrize("output_version", ["v0", "v1"])
def test_agent_loop_streaming(output_version: Literal["v0", "v1"]) -> None:
    @tool
    def get_weather(location: str) -> str:
        """Get the weather for a location."""
        return "It's sunny."

    llm = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-20250514-v1:0",
        output_version=output_version,
    )
    llm_with_tools = llm.bind_tools([get_weather])
    input_message = HumanMessage("What is the weather in San Francisco, CA?")

    tool_call_message: Optional[BaseMessageChunk] = None
    for chunk in llm_with_tools.stream([input_message]):
        assert isinstance(chunk, AIMessageChunk)
        tool_call_message = (
            chunk if tool_call_message is None else tool_call_message + chunk
        )
    assert isinstance(tool_call_message, AIMessageChunk)

    tool_calls = tool_call_message.tool_calls
    assert len(tool_calls) == 1
    tool_call = tool_calls[0]
    tool_message = get_weather.invoke(tool_call)
    assert isinstance(tool_message, ToolMessage)
    response = llm_with_tools.invoke(
        [
            input_message,
            tool_call_message,
            tool_message,
        ]
    )
    assert isinstance(response, AIMessage)


def test_streaming_tool_use_round_trip() -> None:
    """Test that streaming tool call messages can be sent back to Bedrock.

    Regression test for https://github.com/langchain-ai/langchain-aws/issues/827.
    After streaming, content[].tool_use.input is a JSON string instead of a
    dict. When a message is reconstructed from content alone (e.g., loaded
    from a checkpoint without tool_calls), _lc_content_to_bedrock must parse
    string input to a dict to avoid Bedrock ValidationException.
    """

    @tool
    def get_weather(city: str) -> str:
        """Get the current weather for a city."""
        return "It's sunny and 72F."

    llm = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
    )
    llm_with_tools = llm.bind_tools([get_weather], tool_choice="any")

    input_message = HumanMessage("What is the weather in Paris?")

    full: Optional[BaseMessageChunk] = None
    for chunk in llm_with_tools.stream([input_message]):
        assert isinstance(chunk, AIMessageChunk)
        full = chunk if full is None else full + chunk
    assert isinstance(full, AIMessageChunk)

    for tc_chunk in full.tool_call_chunks:
        assert tc_chunk["args"] is None or isinstance(tc_chunk["args"], str)

    assert len(full.tool_calls) == 1
    tool_call = full.tool_calls[0]
    assert tool_call["name"] == "get_weather"
    assert isinstance(tool_call["args"], dict)
    assert isinstance(full.content, list)
    tool_block = next(
        b for b in full.content if isinstance(b, dict) and b.get("type") == "tool_use"
    )
    assert isinstance(tool_block["input"], str), (
        "After streaming accumulation, content[].tool_use.input should be a "
        "string. If this assertion fails, the streaming behavior has changed "
        "and this test may need updating."
    )

    restored_msg = AIMessage(content=full.content)
    assert restored_msg.tool_calls == []

    tool_result = ToolMessage(
        content=get_weather.invoke(tool_call).content,
        tool_call_id=tool_call["id"],
    )

    response = llm_with_tools.invoke([input_message, restored_msg, tool_result])
    assert isinstance(response, AIMessage)


@pytest.mark.default_cassette("test_thinking.yaml.gz")
@pytest.mark.vcr
@pytest.mark.parametrize("output_version", ["v0", "v1"])
def test_thinking(output_version: Literal["v0", "v1"]) -> None:
    llm = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-20250514-v1:0",
        max_tokens=4096,
        additional_model_request_fields={
            "thinking": {"type": "enabled", "budget_tokens": 1024},
        },
        output_version=output_version,
    )

    input_message = {"role": "user", "content": "What is 3^3?"}
    full: Optional[BaseMessageChunk] = None
    for chunk in llm.stream([input_message]):
        assert isinstance(chunk, AIMessageChunk)
        full = chunk if full is None else full + chunk
    assert isinstance(full, AIMessageChunk)

    # Raw content
    if output_version == "v0":
        assert [block["type"] for block in full.content] == [  # type: ignore[index]
            "reasoning_content",
            "text",
        ]  # type: ignore[index,union-attr]
        assert "text" in full.content[0]["reasoning_content"]  # type: ignore[index,union-attr]
        assert "signature" in full.content[0]["reasoning_content"]  # type: ignore[index,union-attr]
    else:
        # v1
        assert [block["type"] for block in full.content] == ["reasoning", "text"]  # type: ignore[index,union-attr]
        assert "signature" in full.content[0]["extras"]  # type: ignore[index,union-attr]

    # Parsed
    content_blocks = full.content_blocks
    assert [block["type"] for block in content_blocks] == ["reasoning", "text"]
    assert content_blocks[0]["type"] == "reasoning"
    assert content_blocks[0].get("reasoning")
    assert "signature" in content_blocks[0]["extras"]

    next_message = {"role": "user", "content": "Thanks!"}
    response = llm.invoke([input_message, full, next_message])

    if output_version == "v0":
        assert [block["type"] for block in response.content] == [  # type: ignore[index]
            "reasoning_content",
            "text",
        ]  # type: ignore[index,union-attr]
        assert "text" in response.content[0]["reasoning_content"]  # type: ignore[index,union-attr]
        assert "signature" in response.content[0]["reasoning_content"]  # type: ignore[index,union-attr]
    else:
        # v1
        assert [block["type"] for block in response.content] == ["reasoning", "text"]  # type: ignore[index,union-attr]
        assert "signature" in response.content[0]["extras"]  # type: ignore[index,union-attr]


PLAINTEXT_DOCUMENT = {
    "document": {
        "format": "txt",
        "name": "company_policy",
        "source": {
            "text": (
                "Company leave policy: Employees get 20 days annual leave. "
                "Consult with your manager for details."
            )
        },
        "context": "HR Policy Manual Section 3.2",
        "citations": {"enabled": True},
    },
}

BLOCKS_DOCUMENT = {
    "document": {
        "format": "txt",
        "name": "company_policy",
        "source": {
            "content": [
                {"text": "Company leave policy: Employees get 20 days annual leave."},
                {"text": "Consult with your manager for details."},
            ]
        },
        "context": "HR Policy Manual Section 3.2",
        "citations": {"enabled": True},
    },
}


PDF_URL = "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf"
PDF_DATA = base64.b64encode(httpx.get(PDF_URL).content).decode("utf-8")

STANDARD_PDF_DOCUMENT = {
    "type": "file",
    "mime_type": "application/pdf",
    "base64": PDF_DATA,
    "name": "my-pdf",  # Converse requires a filename
}


@pytest.mark.vcr
@pytest.mark.parametrize("document", [PLAINTEXT_DOCUMENT, BLOCKS_DOCUMENT])
def test_citations(document: dict[str, Any]) -> None:
    llm = ChatBedrockConverse(model="us.anthropic.claude-sonnet-4-20250514-v1:0")

    input_message = {
        "role": "user",
        "content": [
            document,
            {"type": "text", "text": "How many days of annual leave do employees get?"},
        ],
    }

    full: Optional[BaseMessageChunk] = None
    for chunk in llm.stream([input_message]):
        assert isinstance(chunk, AIMessageChunk)
        full = chunk if full is None else full + chunk
    assert isinstance(full, AIMessageChunk)

    # Raw content
    assert any(block.get("citations") for block in full.content)  # type: ignore[union-attr]

    # Parsed
    content_blocks = full.content_blocks
    assert any(block.get("annotations") for block in content_blocks)
    for block in content_blocks:
        if (block["type"] == "text") and "annotations" in block:
            assert isinstance(block["annotations"], list)
            for annotation in block["annotations"]:
                assert "title" in annotation
                assert "cited_text" in annotation

    next_message = {"role": "user", "content": "Who should they consult with?"}
    response = llm.invoke([input_message, full, next_message])
    assert any(block.get("citations") for block in response.content)  # type: ignore[union-attr]


@pytest.mark.default_cassette("test_citations[document0].yaml.gz")
@pytest.mark.vcr
@pytest.mark.parametrize("output_version", ["v0", "v1"])
def test_citations_v1(output_version: Literal["v0", "v1"]) -> None:
    llm = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-20250514-v1:0",
        output_version=output_version,
    )

    input_message = {
        "role": "user",
        "content": [
            PLAINTEXT_DOCUMENT,
            {"type": "text", "text": "How many days of annual leave do employees get?"},
        ],
    }

    full: Optional[BaseMessageChunk] = None
    for chunk in llm.stream([input_message]):
        assert isinstance(chunk, AIMessageChunk)
        full = chunk if full is None else full + chunk
    assert isinstance(full, AIMessageChunk)

    # Raw content
    if output_version == "v0":
        assert any(block.get("citations") for block in full.content)  # type: ignore[union-attr]
    else:
        # v1
        assert any(block.get("annotations") for block in full.content)  # type: ignore[union-attr]

    next_message = {"role": "user", "content": "Who should they consult with?"}
    response = llm.invoke([input_message, full, next_message])
    if output_version == "v0":
        assert any(block.get("citations") for block in response.content)  # type: ignore[union-attr]
    else:
        # v1
        assert any(block.get("annotations") for block in response.content)  # type: ignore[union-attr]


@pytest.mark.vcr
def test_pdf_citations() -> None:
    model = ChatBedrockConverse(model="us.anthropic.claude-sonnet-4-5-20250929-v1:0")

    message = HumanMessage(
        [
            {"type": "text", "text": "What is the title of this document?"},
            {**STANDARD_PDF_DOCUMENT, "citations": {"enabled": True}},
        ]
    )
    response = model.invoke([message])
    assert any(block.get("citations") for block in response.content)  # type: ignore[union-attr]


def test_bedrock_pdf_inputs() -> None:
    model = ChatBedrockConverse(model="us.anthropic.claude-sonnet-4-5-20250929-v1:0")

    message = HumanMessage(
        [
            {"type": "text", "text": "What is the title of this document?"},
            STANDARD_PDF_DOCUMENT,
        ]
    )
    _ = model.invoke([message])

    # Test OpenAI Chat Completions format
    message = HumanMessage(
        [
            {
                "type": "text",
                "text": "What is the title of this document?",
            },
            {
                "type": "file",
                "file": {
                    "filename": "my-pdf",
                    "file_data": f"data:application/pdf;base64,{PDF_DATA}",
                },
            },
        ]
    )
    _ = model.invoke([message])


def test_get_num_tokens_from_messages_integration() -> None:
    chat = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
    )

    base_messages = [
        SystemMessage(content="You are a helpful assistant."),
        HumanMessage(content="Why did the chicken cross the road?"),
    ]

    token_count = chat.get_num_tokens_from_messages(base_messages)

    assert isinstance(token_count, int)
    assert token_count == 38


@pytest.mark.parametrize("streaming", [False, True])
def test_request_headers(tmp_path: Any, streaming: bool) -> None:
    # Test that we can attach headers to requests
    cassette_path = tmp_path / "headers.yaml"

    vcr = VCR(record_mode="all")
    with vcr.use_cassette(str(cassette_path)):
        model = ChatBedrockConverse(
            model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
            default_headers={"X-Foo": "Bar"},
        )
        if streaming:
            _ = list(model.stream("hi"))
        else:
            _ = model.invoke("hi")

    cassette = yaml.safe_load(cassette_path.read_text())
    interactions = cassette["interactions"]
    request = interactions[0]["request"]
    headers = request["headers"]

    assert headers["X-Foo"] == ["Bar"]


# --- Native structured outputs integration tests ---


class JokeSchema(BaseModel):
    """A joke with setup and punchline."""

    setup: str = Field(description="The setup of the joke")
    punchline: str = Field(description="The punchline of the joke")


class GetWeather(BaseModel):
    """Get the current weather in a given location."""

    location: str = Field(..., description="The city and state, e.g. San Francisco, CA")


def test_structured_output_json_schema_pydantic() -> None:
    """Test method='json_schema' with Pydantic model returns validated instance."""
    model = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", temperature=0
    )
    structured = model.with_structured_output(JokeSchema, method="json_schema")
    result = structured.invoke("Tell me a short joke about programming")
    assert isinstance(result, JokeSchema)
    assert result.setup
    assert result.punchline


def test_structured_output_json_schema_dict() -> None:
    """Test method='json_schema' with dict schema returns matching dict."""
    model = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", temperature=0
    )
    schema = {
        "title": "Joke",
        "description": "A joke with setup and punchline.",
        "type": "object",
        "properties": {
            "setup": {"type": "string", "description": "The setup of the joke"},
            "punchline": {
                "type": "string",
                "description": "The punchline of the joke",
            },
        },
        "required": ["setup", "punchline"],
    }
    structured = model.with_structured_output(schema, method="json_schema")
    result = structured.invoke("Tell me a short joke about programming")
    assert isinstance(result, dict)
    assert "setup" in result
    assert "punchline" in result


def test_structured_output_json_schema_streaming() -> None:
    """Test that streaming works with method='json_schema'."""
    model = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", temperature=0
    )
    structured = model.with_structured_output(JokeSchema, method="json_schema")
    result = None
    for chunk in structured.stream("Tell me a short joke about programming"):
        result = chunk
    assert isinstance(result, JokeSchema)
    assert result.setup
    assert result.punchline


def test_structured_output_json_schema_include_raw() -> None:
    """Test include_raw=True returns dict with raw, parsed, parsing_error."""
    model = ChatBedrockConverse(
        model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", temperature=0
    )
    structured = model.with_structured_output(
        JokeSchema, method="json_schema", include_raw=True
    )
    result = structured.invoke("Tell me a short joke about programming")
    assert isinstance(result, dict)
    assert "raw" in result
    assert "parsed" in result
    assert "parsing_error" in result
    assert isinstance(result["raw"], AIMessage)
    assert isinstance(result["parsed"], JokeSchema)
    assert result["parsing_error"] is None
