"""
Unit tests for the Bedrock integration helpers.
No AWS credentials or backend required.
"""
from visibe.integrations.bedrock import (
    _extract_bedrock_prompt,
    _extract_bedrock_output,
    _extract_bedrock_usage,
    _extract_bedrock_tool_calls,
    _extract_bedrock_model,
    _detect_provider_from_model,
    _parse_invoke_anthropic,
    _parse_invoke_meta,
    _parse_invoke_mistral,
    _parse_invoke_titan,
    _parse_invoke_nova,
    _parse_invoke_cohere,
    _parse_invoke_ai21,
    _parse_invoke_unknown,
    _get_invoke_parser,
)
from visibe.client import _is_bedrock_client


# ------------------------------------------------------------------
# Prompt extraction
# ------------------------------------------------------------------

def test_extract_bedrock_prompt_basic():
    kwargs = {
        "messages": [
            {"role": "user", "content": [{"text": "What is 2 + 2?"}]}
        ]
    }
    result = _extract_bedrock_prompt(kwargs)
    assert result == "What is 2 + 2?"


def test_extract_bedrock_prompt_with_system():
    kwargs = {
        "system": [{"text": "You are a math tutor."}],
        "messages": [
            {"role": "user", "content": [{"text": "What is 2 + 2?"}]}
        ],
    }
    result = _extract_bedrock_prompt(kwargs)
    assert "[system] You are a math tutor." in result
    assert "What is 2 + 2?" in result


def test_extract_bedrock_prompt_empty():
    result = _extract_bedrock_prompt({})
    assert result == ""


def test_extract_bedrock_prompt_multi_turn_returns_last():
    kwargs = {
        "messages": [
            {"role": "user", "content": [{"text": "Hello"}]},
            {"role": "assistant", "content": [{"text": "Hi there"}]},
            {"role": "user", "content": [{"text": "What is 3 + 3?"}]},
        ]
    }
    result = _extract_bedrock_prompt(kwargs)
    # Should return the last user message
    assert "3 + 3" in result


# ------------------------------------------------------------------
# Output extraction
# ------------------------------------------------------------------

def test_extract_bedrock_output_text():
    response = {
        "output": {
            "message": {
                "role": "assistant",
                "content": [{"text": "The answer is 4."}],
            }
        },
        "usage": {"inputTokens": 10, "outputTokens": 5},
    }
    result = _extract_bedrock_output(response)
    assert result == "The answer is 4."


def test_extract_bedrock_output_empty_response():
    result = _extract_bedrock_output({})
    assert result == ""


def test_extract_bedrock_output_multiple_content_blocks():
    response = {
        "output": {
            "message": {
                "role": "assistant",
                "content": [
                    {"text": "First part."},
                    {"text": " Second part."},
                ],
            }
        }
    }
    result = _extract_bedrock_output(response)
    assert "First part" in result


# ------------------------------------------------------------------
# Usage extraction
# ------------------------------------------------------------------

def test_extract_bedrock_usage_standard():
    response = {"usage": {"inputTokens": 150, "outputTokens": 42}}
    input_t, output_t = _extract_bedrock_usage(response)
    assert input_t == 150
    assert output_t == 42


def test_extract_bedrock_usage_missing():
    input_t, output_t = _extract_bedrock_usage({})
    assert input_t == 0
    assert output_t == 0


def test_extract_bedrock_usage_zero_tokens():
    response = {"usage": {"inputTokens": 0, "outputTokens": 0}}
    input_t, output_t = _extract_bedrock_usage(response)
    assert input_t == 0
    assert output_t == 0


# ------------------------------------------------------------------
# Tool call extraction
# ------------------------------------------------------------------

def test_extract_bedrock_tool_calls_single():
    response = {
        "output": {
            "message": {
                "role": "assistant",
                "content": [
                    {
                        "toolUse": {
                            "toolUseId": "abc123",
                            "name": "get_weather",
                            "input": {"city": "Paris"},
                        }
                    },
                    {"text": "Let me check the weather."},
                ],
            }
        }
    }
    tools = _extract_bedrock_tool_calls(response)
    assert len(tools) == 1
    assert tools[0]["tool_name"] == "get_weather"
    assert tools[0]["status"] == "success"
    assert "Paris" in tools[0]["input"]


def test_extract_bedrock_tool_calls_none():
    response = {
        "output": {
            "message": {
                "role": "assistant",
                "content": [{"text": "No tool needed."}],
            }
        }
    }
    tools = _extract_bedrock_tool_calls(response)
    assert tools == []


def test_extract_bedrock_tool_calls_empty_response():
    tools = _extract_bedrock_tool_calls({})
    assert tools == []


# ------------------------------------------------------------------
# Model ID extraction
# ------------------------------------------------------------------

def test_extract_bedrock_model_simple():
    kwargs = {"modelId": "anthropic.claude-3-haiku-20240307-v1:0"}
    result = _extract_bedrock_model(kwargs)
    assert result == "anthropic.claude-3-haiku-20240307-v1:0"


def test_extract_bedrock_model_arn():
    kwargs = {
        "modelId": (
            "arn:aws:bedrock:us-east-1:123456789:"
            "inference-profile/anthropic.claude-3-haiku-20240307-v1:0"
        )
    }
    result = _extract_bedrock_model(kwargs)
    assert result == "anthropic.claude-3-haiku-20240307-v1:0"


def test_extract_bedrock_model_missing():
    result = _extract_bedrock_model({})
    assert result == "unknown"


# ------------------------------------------------------------------
# Provider detection
# ------------------------------------------------------------------

def test_detect_provider_anthropic():
    assert _detect_provider_from_model("anthropic.claude-3-haiku-20240307-v1:0") == "anthropic"


def test_detect_provider_meta():
    assert _detect_provider_from_model("meta.llama3-70b-instruct-v1:0") == "meta"


def test_detect_provider_mistral():
    assert _detect_provider_from_model("mistral.mistral-large-2402-v1:0") == "mistral"


def test_detect_provider_amazon():
    assert _detect_provider_from_model("amazon.nova-pro-v1:0") == "amazon"


def test_detect_provider_cohere():
    assert _detect_provider_from_model("cohere.command-r-plus-v1:0") == "cohere"


def test_detect_provider_ai21():
    assert _detect_provider_from_model("ai21.jamba-1.5-large-v1:0") == "ai21"


def test_detect_provider_unknown():
    assert _detect_provider_from_model("unknown-model") == "bedrock"


# ------------------------------------------------------------------
# Client type detection
# ------------------------------------------------------------------

def test_is_bedrock_client_rejects_string():
    assert not _is_bedrock_client("a string")


def test_is_bedrock_client_rejects_int():
    assert not _is_bedrock_client(42)


def test_is_bedrock_client_rejects_none():
    assert not _is_bedrock_client(None)


def test_is_bedrock_client_rejects_object():
    assert not _is_bedrock_client(object())


# ------------------------------------------------------------------
# invoke_model parsers — Anthropic Claude
# ------------------------------------------------------------------

def test_parse_invoke_anthropic_basic():
    request = {"messages": [{"role": "user", "content": "Hello Claude"}]}
    response = {
        "content": [{"type": "text", "text": "Hi there!"}],
        "usage": {"input_tokens": 10, "output_tokens": 5},
    }
    prompt, output, input_t, output_t = _parse_invoke_anthropic(response, request)
    assert prompt == "Hello Claude"
    assert output == "Hi there!"
    assert input_t == 10
    assert output_t == 5


def test_parse_invoke_anthropic_with_system():
    request = {
        "system": "You are helpful.",
        "messages": [{"role": "user", "content": "What time is it?"}],
    }
    response = {
        "content": [{"type": "text", "text": "I don't have a clock."}],
        "usage": {"input_tokens": 20, "output_tokens": 8},
    }
    prompt, output, input_t, output_t = _parse_invoke_anthropic(response, request)
    assert "[system]" in prompt
    assert "What time is it?" in prompt
    assert output == "I don't have a clock."


def test_parse_invoke_anthropic_content_as_list():
    request = {
        "messages": [
            {"role": "user", "content": [{"type": "text", "text": "List content"}]}
        ]
    }
    response = {
        "content": [{"type": "text", "text": "Got it."}],
        "usage": {"input_tokens": 5, "output_tokens": 3},
    }
    prompt, output, _, _ = _parse_invoke_anthropic(response, request)
    assert prompt == "List content"
    assert output == "Got it."


def test_parse_invoke_anthropic_empty():
    prompt, output, input_t, output_t = _parse_invoke_anthropic({}, {})
    assert prompt == ""
    assert output == ""
    assert input_t == 0
    assert output_t == 0


# ------------------------------------------------------------------
# invoke_model parsers — Meta Llama
# ------------------------------------------------------------------

def test_parse_invoke_meta_basic():
    request = {"prompt": "<s>[INST] Tell me a joke [/INST]"}
    response = {
        "generation": "Why did the chicken cross the road?",
        "prompt_token_count": 15,
        "generation_token_count": 12,
    }
    prompt, output, input_t, output_t = _parse_invoke_meta(response, request)
    assert "Tell me a joke" in prompt
    assert output == "Why did the chicken cross the road?"
    assert input_t == 15
    assert output_t == 12


def test_parse_invoke_meta_empty():
    prompt, output, input_t, output_t = _parse_invoke_meta({}, {})
    assert prompt == ""
    assert output == ""
    assert input_t == 0
    assert output_t == 0


# ------------------------------------------------------------------
# invoke_model parsers — Mistral (two formats)
# ------------------------------------------------------------------

def test_parse_invoke_mistral_text_completion_format():
    request = {"prompt": "<s>[INST] Hello [/INST]"}
    response = {"outputs": [{"text": "Hi from Mistral!", "stop_reason": "stop"}]}
    prompt, output, input_t, output_t = _parse_invoke_mistral(response, request)
    assert "Hello" in prompt
    assert output == "Hi from Mistral!"
    assert input_t == 0   # Mistral returns no token counts
    assert output_t == 0


def test_parse_invoke_mistral_chat_format():
    request = {"messages": [{"role": "user", "content": "Hello chat"}]}
    response = {"choices": [{"message": {"content": "Chat reply"}, "finish_reason": "stop"}]}
    prompt, output, input_t, output_t = _parse_invoke_mistral(response, request)
    assert prompt == "Hello chat"
    assert output == "Chat reply"
    assert input_t == 0
    assert output_t == 0


def test_parse_invoke_mistral_empty():
    prompt, output, input_t, output_t = _parse_invoke_mistral({}, {})
    assert prompt == ""
    assert input_t == 0
    assert output_t == 0


# ------------------------------------------------------------------
# invoke_model parsers — Amazon Titan
# ------------------------------------------------------------------

def test_parse_invoke_titan_basic():
    request = {"inputText": "Write a haiku"}
    response = {
        "results": [{"outputText": "Petals fall gently", "tokenCount": 6}],
        "inputTextTokenCount": 4,
    }
    prompt, output, input_t, output_t = _parse_invoke_titan(response, request)
    assert prompt == "Write a haiku"
    assert output == "Petals fall gently"
    assert input_t == 4
    assert output_t == 6


def test_parse_invoke_titan_empty():
    prompt, output, input_t, output_t = _parse_invoke_titan({}, {})
    assert prompt == ""
    assert output == ""
    assert input_t == 0
    assert output_t == 0


# ------------------------------------------------------------------
# invoke_model parsers — Amazon Nova
# ------------------------------------------------------------------

def test_parse_invoke_nova_basic():
    request = {
        "messages": [{"role": "user", "content": [{"text": "What is Nova?"}]}]
    }
    response = {
        "output": {
            "message": {"content": [{"text": "Nova is an Amazon model."}]}
        },
        "usage": {"inputTokens": 8, "outputTokens": 7},
    }
    prompt, output, input_t, output_t = _parse_invoke_nova(response, request)
    assert prompt == "What is Nova?"
    assert output == "Nova is an Amazon model."
    assert input_t == 8
    assert output_t == 7


def test_parse_invoke_nova_with_system():
    request = {
        "system": [{"text": "Be concise."}],
        "messages": [{"role": "user", "content": [{"text": "Hi"}]}],
    }
    response = {
        "output": {"message": {"content": [{"text": "Hello!"}]}},
        "usage": {"inputTokens": 5, "outputTokens": 2},
    }
    prompt, output, _, _ = _parse_invoke_nova(response, request)
    assert "[system]" in prompt
    assert "Hi" in prompt
    assert output == "Hello!"


def test_parse_invoke_nova_empty():
    prompt, output, input_t, output_t = _parse_invoke_nova({}, {})
    assert prompt == ""
    assert output == ""
    assert input_t == 0
    assert output_t == 0


# ------------------------------------------------------------------
# invoke_model parsers — Cohere (two formats)
# ------------------------------------------------------------------

def test_parse_invoke_cohere_command_r_format():
    request = {"message": "Summarize this"}
    response = {
        "text": "Here is a summary.",
        "meta": {"billed_units": {"input_tokens": 12, "output_tokens": 6}},
    }
    prompt, output, input_t, output_t = _parse_invoke_cohere(response, request)
    assert prompt == "Summarize this"
    assert output == "Here is a summary."
    assert input_t == 12
    assert output_t == 6


def test_parse_invoke_cohere_original_format():
    request = {"prompt": "Tell me something"}
    response = {"generations": [{"text": "Something interesting."}]}
    prompt, output, input_t, output_t = _parse_invoke_cohere(response, request)
    assert prompt == "Tell me something"
    assert output == "Something interesting."
    assert input_t == 0   # Original Command returns no token counts
    assert output_t == 0


def test_parse_invoke_cohere_empty():
    prompt, output, input_t, output_t = _parse_invoke_cohere({}, {})
    assert prompt == ""
    assert input_t == 0
    assert output_t == 0


# ------------------------------------------------------------------
# invoke_model parsers — AI21 Jamba
# ------------------------------------------------------------------

def test_parse_invoke_ai21_basic():
    request = {"messages": [{"role": "user", "content": "What is Jamba?"}]}
    response = {
        "choices": [{"message": {"content": "Jamba is an AI21 model."}}],
        "usage": {"prompt_tokens": 9, "completion_tokens": 8},
    }
    prompt, output, input_t, output_t = _parse_invoke_ai21(response, request)
    assert prompt == "What is Jamba?"
    assert output == "Jamba is an AI21 model."
    assert input_t == 9
    assert output_t == 8


def test_parse_invoke_ai21_empty():
    prompt, output, input_t, output_t = _parse_invoke_ai21({}, {})
    assert prompt == ""
    assert output == ""
    assert input_t == 0
    assert output_t == 0


# ------------------------------------------------------------------
# invoke_model parsers — Unknown fallback
# ------------------------------------------------------------------

def test_parse_invoke_unknown_returns_empty_prompt():
    _, _, input_t, output_t = _parse_invoke_unknown({"some": "data"}, {})
    assert input_t == 0
    assert output_t == 0


def test_parse_invoke_unknown_stringifies_response():
    prompt, output, _, _ = _parse_invoke_unknown({"key": "value"}, {})
    assert prompt == ""
    assert "value" in output


# ------------------------------------------------------------------
# _get_invoke_parser — router
# ------------------------------------------------------------------

def test_get_invoke_parser_anthropic():
    assert _get_invoke_parser("anthropic.claude-3-haiku-20240307-v1:0") is _parse_invoke_anthropic


def test_get_invoke_parser_meta():
    assert _get_invoke_parser("meta.llama3-70b-instruct-v1:0") is _parse_invoke_meta


def test_get_invoke_parser_mistral():
    assert _get_invoke_parser("mistral.mistral-large-2402-v1:0") is _parse_invoke_mistral


def test_get_invoke_parser_titan():
    assert _get_invoke_parser("amazon.titan-text-express-v1") is _parse_invoke_titan


def test_get_invoke_parser_nova():
    assert _get_invoke_parser("amazon.nova-micro-v1:0") is _parse_invoke_nova


def test_get_invoke_parser_cohere():
    assert _get_invoke_parser("cohere.command-r-plus-v1:0") is _parse_invoke_cohere


def test_get_invoke_parser_ai21():
    assert _get_invoke_parser("ai21.jamba-1.5-large-v1:0") is _parse_invoke_ai21


def test_get_invoke_parser_unknown_falls_back():
    assert _get_invoke_parser("some-unknown-model") is _parse_invoke_unknown


def test_get_invoke_parser_case_insensitive():
    assert _get_invoke_parser("ANTHROPIC.Claude-3-Haiku") is _parse_invoke_anthropic
