#!/usr/bin/env python3
"""Test script to specifically reproduce the space + number parsing issue."""

import boto3

# REMOVE THESE LINES TO FIX!!!
from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
BotocoreInstrumentor().instrument()


def test_space_number_issue():
    """Test that reproduces the issue where detlas get parsed as number instead of str:
    IE Chunk 5: {'input': 23} instead of {'input': ' 23 '}"""

    client = boto3.client("bedrock-runtime", region_name="us-east-2")

    # Tool that expects a calculation query
    tool_config = {
        "tools": [
            {
                "toolSpec": {
                    "name": "response_tool",
                    "description": "Respond to the user's request",
                    "inputSchema": {
                        "json": {
                            "type": "object",
                            "properties": {
                                "response": {
                                    "type": "string",
                                    "description": "The response to the user's request",
                                }
                            },
                            "required": ["response"],
                        }
                    },
                }
            }
        ],
        "toolChoice": {"auto": {}},
    }

    # Message that should trigger the specific issue - long equation with many numbers
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "text": "Give me the first 20 characters of pi. Respond using the `response_tool`"
                }
            ],
        }
    ]

    print("=== TESTING SPACE + NUMBER PARSING ISSUE ===")
    print(
        "Looking for numbers in the complex equation that might get split/parsed incorrectly during streaming..."
    )

    response = client.converse_stream(
        modelId="us.anthropic.claude-sonnet-4-20250514-v1:0",
        messages=messages,
        toolConfig=tool_config,
        inferenceConfig={"maxTokens": 1000, "temperature": 0.1},
    )

    tool_chunks = []
    chunk_count = 0

    for chunk in response["stream"]:
        chunk_count += 1

        if "contentBlockDelta" in chunk:
            delta = chunk["contentBlockDelta"]["delta"]

            if "toolUse" in delta:
                tool_use = delta["toolUse"]
                tool_chunks.append({"chunk_num": chunk_count, "tool_use": tool_use})

                print(f"Chunk {chunk_count}: {tool_use}")

                # Check specifically for numeric inputs
                if "input" in tool_use:
                    input_val = tool_use["input"]
                    if isinstance(input_val, (int, float)):
                        print(
                            f"  ⚠️  NUMERIC INPUT DETECTED: {input_val} (type: {type(input_val)})"
                        )
                    elif isinstance(input_val, str) and input_val.strip().isdigit():
                        print(f"  📝 String numeric input: '{input_val}'")

        # Stop after reasonable number of chunks
        if chunk_count > 100:
            break

    print(f"\n=== SUMMARY ===")
    print(f"Total chunks processed: {chunk_count}")
    print(f"Tool use chunks found: {len(tool_chunks)}")

    # Reconstruct what the full tool arguments should look like
    full_input = ""
    for chunk in tool_chunks:
        tool_use = chunk["tool_use"]
        if "input" in tool_use:
            input_val = tool_use["input"]
            if isinstance(input_val, str):
                full_input += input_val
            else:
                full_input += str(input_val)

    print(f"Reconstructed input: '{full_input}'")

    # Check for the specific issue patterns
    numeric_chunks = [
        chunk
        for chunk in tool_chunks
        if "input" in chunk["tool_use"]
        and isinstance(chunk["tool_use"]["input"], (int, float))
    ]
    if numeric_chunks:
        print(
            f"\n⚠️  ISSUE DETECTED: Found {len(numeric_chunks)} chunks with raw numeric values:"
        )
        for chunk in numeric_chunks:
            print(f"  Chunk {chunk['chunk_num']}: {chunk['tool_use']['input']}")


if __name__ == "__main__":
    try:
        test_space_number_issue()

        # asyncio.run(test_agent_run())

    except Exception as e:
        print(f"Error: {e}")
        import traceback

        traceback.print_exc()