# Copyright Coralogix Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import base64
import os

import boto3
import pytest
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.sampling import ALWAYS_OFF

from llm_tracekit.bedrock.instrumentor import BedrockInstrumentor
from llm_tracekit.instrumentation_utils import (
    OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT,
)

EVENT_STREAM_CT = "application/vnd.amazon.eventstream"


pytest.register_assert_rewrite("tests.bedrock.utils")


@pytest.fixture(autouse=True)
def bedrock_env_vars():
    for env_var_name in [
        "AWS_ACCESS_KEY_ID",
        "AWS_SECRET_ACCESS_KEY",
        "AWS_SESSION_TOKEN",
    ]:
        if not os.getenv(env_var_name):
            os.environ[env_var_name] = f"test_{env_var_name.lower()}"

    if not os.getenv("AWS_DEFAULT_REGION"):
        os.environ["AWS_DEFAULT_REGION"] = "us-east-1"


@pytest.fixture
def bedrock_client_with_content(instrument_with_content):
    return boto3.client("bedrock-runtime")


@pytest.fixture
def bedrock_client_no_content(instrument_no_content):
    return boto3.client("bedrock-runtime")


@pytest.fixture
def bedrock_agent_client_with_content(instrument_with_content):
    return boto3.client("bedrock-agent-runtime")


@pytest.fixture
def bedrock_agent_client_no_content(instrument_no_content):
    return boto3.client("bedrock-agent-runtime")


@pytest.fixture
def claude_model_id() -> str:
    return "anthropic.claude-3-5-sonnet-20240620-v1:0"


@pytest.fixture
def llama_model_id() -> str:
    return "meta.llama3-8b-instruct-v1:0"


@pytest.fixture
def agent_id() -> str:
    return os.environ.get("AWS_BEDROCK_AGENT_ID", "test_agent_id")


@pytest.fixture
def agent_alias_id() -> str:
    return os.environ.get("AWS_BEDROCK_AGENT_ALIAS_ID", "test_agent_alias_id")


def handle_recording_boto_response(response: dict) -> dict:
    """Prepares boto3 reponses for recoring/playback.

    Args:
        response: HTTP response

    Returns:
        Fixed response

    Notes:
        * This function is required to play back streaming and invoke_model output
        * This function was generated by ChatGPT
    """
    # runs both when recording and when loading the cassette
    headers = response["headers"]
    if headers.get("x-vcr-base64") == ["yes"]:
        # playback branch ─ decode back to raw bytes
        response["body"]["string"] = base64.b64decode(response["body"]["string"])
        headers.pop("x-vcr-base64")
        return response

    if EVENT_STREAM_CT in (headers.get("Content-Type") or [""])[0]:
        # record branch ─ protect the bytes from PyYAML
        raw = response["body"]["string"]  # bytes from socket
        response["body"]["string"] = base64.b64encode(raw).decode()
        headers["x-vcr-base64"] = ["yes"]

    # Fix Content-Length header
    response["headers"]["Content-Length"] = [str(len(response["body"]["string"]))]
    return response


@pytest.fixture(scope="module")
def vcr_config():
    return {
        "filter_headers": [
            ("X-Amz-Security-Token", "test-security-token"),
            ("Authorization", "test-auth"),
        ],
        "decode_compressed_response": True,
        "before_record_response": handle_recording_boto_response,
    }


@pytest.fixture(scope="function")
def instrument_no_content(tracer_provider, meter_provider):
    os.environ.update({OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "False"})

    instrumentor = BedrockInstrumentor()
    instrumentor.instrument(
        tracer_provider=tracer_provider,
        meter_provider=meter_provider,
    )

    yield instrumentor
    os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None)
    instrumentor.uninstrument()


@pytest.fixture(scope="function")
def instrument_with_content(tracer_provider, meter_provider):
    os.environ.update({OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "True"})
    instrumentor = BedrockInstrumentor()
    instrumentor.instrument(
        tracer_provider=tracer_provider,
        meter_provider=meter_provider,
    )

    yield instrumentor
    os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None)
    instrumentor.uninstrument()


@pytest.fixture(scope="function")
def instrument_with_content_unsampled(span_exporter, meter_provider):
    os.environ.update({OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "True"})

    tracer_provider = TracerProvider(sampler=ALWAYS_OFF)
    tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))

    instrumentor = BedrockInstrumentor()
    instrumentor.instrument(
        tracer_provider=tracer_provider,
        meter_provider=meter_provider,
    )

    yield instrumentor
    os.environ.pop(OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, None)
    instrumentor.uninstrument()
