import asyncio
import json
import os
import uuid
from pathlib import Path
from typing import Final

import pytest
import yaml

from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server

MODEL: Final = "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps({
    "output": {"message": {"role": "assistant", "content": [{"text": "bedrock wire control"}]}},
    "stopReason": "end_turn", "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
    "metrics": {"latencyMs": 1},
}).encode()


def bearer_peer(request: Request) -> Reply:
    assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse"
    assert request.headers["authorization"] == f"Bearer {TOKEN}"
    assert "x-amz-security-token" not in request.headers
    body: Final = json.loads(request.body)
    assert body["messages"] == [{"role": "user", "content": [{"text": "synthetic bearer request"}]}]
    assert body["system"] == [{"text": "synthetic system"}]
    assert body["inferenceConfig"]["maxTokens"] == 16
    assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "api_key"}.intersection(body)
    return Reply(body=RESPONSE)


@pytest.mark.covers("other.provider_wire.bedrock.bearer_sdk_skips_credential_chain")
async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
    import litellm

    empty: Final = tmp_path / "empty-aws-config"
    empty.write_text("")
    for name in tuple(name for name in os.environ if name.startswith("AWS_")):
        monkeypatch.delenv(name, raising=False)
    for name, value in {"AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}.items():
        monkeypatch.setenv(name, value)
    with wire_server(bearer_peer) as wire:
        with pytest.raises(litellm.APIConnectionError, match=r"config profile .* could not be found"):
            await asyncio.to_thread(litellm.completion, model=MODEL, aws_profile_name="integration-profile-must-not-be-read", aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url, messages=[{"role": "user", "content": "synthetic credential control"}], timeout=5, num_retries=0)
        assert wire.drain() == ()
        for source in ("argument", "environment"):
            if source == "environment":
                monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", TOKEN)
            parameters: Final = {
                "model": MODEL, "api_key": TOKEN if source == "argument" else None,
                "aws_region_name": "us-east-1", "aws_profile_name": "integration-profile-must-not-be-read",
                "aws_bedrock_runtime_endpoint": wire.url, "timeout": 5, "num_retries": 0,
                "messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
                "max_tokens": 16,
            }
            for asynchronous in (False, True):
                result: Final = await litellm.acompletion(**parameters) if asynchronous else await asyncio.to_thread(litellm.completion, **parameters)
                assert result.choices[0].message.content == "bedrock wire control"
                assert result.choices[0].finish_reason == "stop"
                assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4
                assert len(wire.drain()) == 1


@pytest.mark.covers("other.provider_wire.bedrock.bearer_db_yaml_survives_reload")
def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload(gateway: Gateway, tmp_path: Path) -> None:
    empty: Final = tmp_path / "empty-aws-config"
    empty.write_text("")
    with wire_server(bearer_peer) as wire:
        parameters: Final = {
            "model": MODEL, "api_key": "os.environ/INTEGRATION_BEARER_TOKEN", "aws_region_name": "us-east-1",
            "aws_profile_name": "integration-profile-must-not-be-read", "aws_bedrock_runtime_endpoint": wire.url,
        }
        alias: Final = f"integration-yaml-{uuid.uuid4().hex}"
        configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
        configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}]
        path: Final = tmp_path / "bedrock.yaml"
        path.write_text(yaml.safe_dump(configuration))
        overrides: Final = {"INTEGRATION_BEARER_TOKEN": TOKEN, "AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}
        with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario:
            database_model: Final = scenario.model(**parameters)
            for generation in range(2):
                for model in (alias, database_model):
                    response: Final = candidate.request("POST", "/v1/chat/completions", {
                        "model": model, "messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
                        "max_tokens": 16, "cache": {"no-cache": True},
                    })
                    assert response.status_code == 200, response.text
                    assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
                    assert response.json()["usage"]["total_tokens"] == 15
                    assert len(wire.drain()) == 1, f"Expected actual provider call after reload {generation}"
                if generation == 0:
                    entries: Final = candidate.get("/model/info")["data"]
                    target: Final = next(entry for entry in entries if entry["model_name"] == database_model)
                    response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "bearer reload"}})
                    assert response.status_code == 200, response.text