import json
from datetime import datetime
from typing import AsyncIterator, Dict, Any
import asyncio
import unittest.mock
from unittest.mock import AsyncMock, MagicMock

import litellm
import pytest
from dotenv import load_dotenv
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
    anthropic_messages,
)

from typing import Optional
from litellm.types.utils import StandardLoggingPayload
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.router import Router
import importlib
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM


class TestCustomLogger(CustomLogger):
    def __init__(self):
        super().__init__()
        self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None

    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
        print("inside async_log_success_event")
        self.logged_standard_logging_payload = kwargs.get("standard_logging_object")

        pass


class BaseAnthropicMessagesTest:
    """Base class for anthropic messages tests to reduce code duplication"""

    @property
    def model_config(self) -> Dict[str, Any]:
        """Override in subclasses to provide model-specific configuration"""
        raise NotImplementedError("Subclasses must implement model_config")

    @property
    def expected_model_name_in_logging(self) -> str:
        """
        This is the model name that is expected to be in the logging payload
        """
        raise NotImplementedError(
            "Subclasses must implement expected_model_name_in_logging"
        )

    def _validate_response(self, response: Any):
        """Validate non-streaming response structure"""
        # Handle type checking - response should be a dict for non-streaming
        if isinstance(response, AsyncIterator):
            pytest.fail("Expected non-streaming response but got AsyncIterator")

        assert isinstance(
            response, dict
        ), f"Expected dict response, got {type(response)}"
        assert "id" in response
        assert "content" in response
        assert "model" in response
        assert response.get("role") == "assistant"

    @pytest.mark.asyncio
    async def test_non_streaming_base(self):
        """Base test for non-streaming requests"""
        litellm._turn_on_debug()

        request_params = self.model_config

        # Set up test parameters
        messages = [{"role": "user", "content": "Hello, can you tell me a short joke?"}]

        # Prepare call arguments
        call_args = {
            "messages": messages,
            "max_tokens": 100,
        }

        # Add any additional config from subclass
        call_args.update(request_params)

        # Call the handler
        response = await litellm.anthropic.messages.acreate(**call_args)

        print(f"Non-streaming {request_params['model']} response: ", response)

        # Verify response
        self._validate_response(response)

        print(f"Non-streaming response: {json.dumps(response, indent=2, default=str)}")
        return response

    @pytest.mark.asyncio
    async def test_streaming_base(self):
        """Base test for streaming requests"""
        request_params = self.model_config
        # Set up test parameters
        messages = [{"role": "user", "content": "Hello, can you tell me a short joke?"}]

        # Prepare call arguments
        call_args = {
            "messages": messages,
            "max_tokens": 100,
            "stream": True,
            "client": AsyncHTTPHandler(),
        }

        # Add any additional config from subclass
        call_args.update(request_params)

        # Call the handler
        response = await litellm.anthropic.messages.acreate(**call_args)

        collected_chunks = []
        if isinstance(response, AsyncIterator):
            async for chunk in response:
                print("chunk=", chunk)
                collected_chunks.append(chunk)

        print("collected_chunks=", collected_chunks)
        return collected_chunks

    @pytest.mark.asyncio
    async def test_response_format_consistency(self):
        """
        Test that response content blocks are consistently dicts (not Pydantic objects).

        This ensures that code like response["content"][0]["type"] works
        regardless of the target provider.

        Issue: https://github.com/BerriAI/litellm/issues/20342
        """
        litellm._turn_on_debug()

        request_params = self.model_config

        # Set up test parameters
        messages = [{"role": "user", "content": "Say hi"}]

        # Prepare call arguments
        call_args = {
            "messages": messages,
            "max_tokens": 100,
        }

        # Add any additional config from subclass
        call_args.update(request_params)

        # Call the handler
        response = await litellm.anthropic.messages.acreate(**call_args)

        print(
            f"Response for {request_params['model']}: {json.dumps(response, indent=2, default=str)}"
        )

        # Verify response structure
        assert "content" in response, "Response should have 'content' field"
        assert len(response["content"]) > 0, "Response content should not be empty"

        # Get the first content block
        block = response["content"][0]

        # Check that the block is a dict, not a Pydantic object
        assert isinstance(block, dict), (
            f"Content block should be a dict, but got {type(block)}. "
            f"This means response format is inconsistent across providers."
        )

        # Verify we can access fields using dict syntax (not object attributes)
        try:
            block_type = block["type"]
            print(f"✓ Successfully accessed block['type']: {block_type}")
        except TypeError as e:
            pytest.fail(
                f"Cannot access content block using dict syntax: {e}. "
                f"Block type: {type(block)}"
            )

        # Verify the block has expected structure
        assert "type" in block, "Content block should have 'type' field"
        if block["type"] == "text":
            assert "text" in block, "Text content block should have 'text' field"

        print(
            f"✓ Response format consistency test passed for {request_params['model']}"
        )

    @pytest.mark.asyncio
    async def test_anthropic_messages_litellm_router_streaming_with_logging(self):
        """
        Test that logging and cost tracking works for anthropic_messages with streaming request
        """
        test_custom_logger = TestCustomLogger()
        litellm.callbacks = [test_custom_logger]
        litellm._turn_on_debug()
        router = Router(
            model_list=[
                {
                    "model_name": "claude-special-alias",
                    "litellm_params": {**self.model_config},
                }
            ]
        )

        # Set up test parameters
        messages = [{"role": "user", "content": "Hello, can you tell me a short joke?"}]

        # Call the handler
        response = await router.aanthropic_messages(
            messages=messages,
            model="claude-special-alias",
            max_tokens=100,
            stream=True,
        )

        response_prompt_tokens = 0
        response_completion_tokens = 0
        all_anthropic_usage_chunks = []
        buffer = ""

        async for chunk in response:
            # Decode chunk if it's bytes
            print("chunk=", chunk)

            # Handle SSE format chunks
            if isinstance(chunk, bytes):
                chunk_str = chunk.decode("utf-8")
                buffer += chunk_str
                # Extract the JSON data part from SSE format
                for line in buffer.split("\n"):
                    if line.startswith("data: "):
                        try:
                            json_data = json.loads(line[6:])  # Skip the 'data: ' prefix
                            print(
                                "\n\nJSON data:",
                                json.dumps(json_data, indent=4, default=str),
                            )

                            # Extract usage information
                            if (
                                json_data.get("type") == "message_start"
                                and "message" in json_data
                            ):
                                if "usage" in json_data["message"]:
                                    usage = json_data["message"]["usage"]
                                    all_anthropic_usage_chunks.append(usage)
                                    print(
                                        "USAGE BLOCK",
                                        json.dumps(usage, indent=4, default=str),
                                    )
                            elif "usage" in json_data:
                                usage = json_data["usage"]
                                all_anthropic_usage_chunks.append(usage)
                                print(
                                    "USAGE BLOCK",
                                    json.dumps(usage, indent=4, default=str),
                                )
                        except json.JSONDecodeError:
                            print(f"Failed to parse JSON from: {line[6:]}")
            elif hasattr(chunk, "message"):
                if chunk.message.usage:
                    print(
                        "USAGE BLOCK",
                        json.dumps(chunk.message.usage, indent=4, default=str),
                    )
                    all_anthropic_usage_chunks.append(chunk.message.usage)
            elif hasattr(chunk, "usage"):
                print("USAGE BLOCK", json.dumps(chunk.usage, indent=4, default=str))
                all_anthropic_usage_chunks.append(chunk.usage)

        print(
            "all_anthropic_usage_chunks",
            json.dumps(all_anthropic_usage_chunks, indent=4, default=str),
        )

        # Extract token counts from usage data
        if all_anthropic_usage_chunks:
            response_prompt_tokens = max(
                [usage.get("input_tokens", 0) for usage in all_anthropic_usage_chunks]
            )
            response_completion_tokens = max(
                [usage.get("output_tokens", 0) for usage in all_anthropic_usage_chunks]
            )

        print("input_tokens_anthropic_api", response_prompt_tokens)
        print("output_tokens_anthropic_api", response_completion_tokens)

        await asyncio.sleep(4)

        print(
            "logged_standard_logging_payload",
            json.dumps(
                test_custom_logger.logged_standard_logging_payload,
                indent=4,
                default=str,
            ),
        )

        assert (
            test_custom_logger.logged_standard_logging_payload is not None
        ), "Logging payload should not be None"
        assert (
            test_custom_logger.logged_standard_logging_payload["messages"] == messages
        )
        assert (
            test_custom_logger.logged_standard_logging_payload["response"] is not None
        )
        assert (
            test_custom_logger.logged_standard_logging_payload["model"]
            == self.expected_model_name_in_logging
        )

        # check logged usage + spend
        assert test_custom_logger.logged_standard_logging_payload["response_cost"] > 0
        assert (
            test_custom_logger.logged_standard_logging_payload["prompt_tokens"]
            == response_prompt_tokens
        )
        assert (
            test_custom_logger.logged_standard_logging_payload["completion_tokens"]
            == response_completion_tokens
        )
