"""Tests for AWS Bedrock auto-instrumentation."""
from __future__ import annotations

import json
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any
from unittest.mock import MagicMock, patch

import pytest

from fibinaci._testing.mock_client import MockFibinaci


class TestBedrockInstrumentation:
    def test_extract_converse_usage(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.bedrock import _extract_converse_usage

        response = {
            "usage": {"inputTokens": 100, "outputTokens": 50},
            "stopReason": "end_turn",
        }
        inp, out, stop = _extract_converse_usage(response)
        assert inp == 100
        assert out == 50
        assert stop == "end_turn"

    def test_extract_converse_usage_empty(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.bedrock import _extract_converse_usage

        inp, out, stop = _extract_converse_usage({})
        assert inp is None
        assert out is None
        assert stop is None

    def test_extract_invoke_usage_anthropic(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.bedrock import _extract_invoke_usage
        import io

        body_data = json.dumps({
            "usage": {"input_tokens": 80, "output_tokens": 40},
            "stop_reason": "end_turn",
        }).encode()
        response = {"body": io.BytesIO(body_data)}
        inp, out, stop = _extract_invoke_usage(response, "anthropic.claude-3-sonnet")
        assert inp == 80
        assert out == 40
        assert stop == "end_turn"

    def test_extract_invoke_usage_titan(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.bedrock import _extract_invoke_usage
        import io

        body_data = json.dumps({
            "inputTextTokenCount": 60,
            "results": [{"tokenCount": 30}],
            "completionReason": "FINISH",
        }).encode()
        response = {"body": io.BytesIO(body_data)}
        inp, out, stop = _extract_invoke_usage(response, "amazon.titan-text-express-v1")
        assert inp == 60
        assert out == 30
        assert stop == "FINISH"

    def test_extract_invoke_usage_no_body(self, mock_fibinaci: MockFibinaci):
        from fibinaci.instrumentation.bedrock import _extract_invoke_usage

        inp, out, stop = _extract_invoke_usage({}, "some-model")
        assert inp is None
        assert out is None
        assert stop is None

    def test_instrument_hooks_create_client(self, mock_fibinaci: MockFibinaci):
        botocore_mod = MagicMock()
        client_creator_cls = MagicMock()
        original_create = MagicMock()
        client_creator_cls.create_client = original_create
        botocore_mod.client.ClientCreator = client_creator_cls

        with patch.dict("sys.modules", {
            "botocore": botocore_mod,
            "botocore.client": botocore_mod.client,
        }):
            from fibinaci.instrumentation.bedrock import BedrockInstrumentor
            inst = BedrockInstrumentor()
            inst.instrument()
            assert client_creator_cls.create_client is not original_create
            inst.uninstrument()
