"""Test Embedding"""

import os
from unittest import mock

import pytest

from databuilder_logic_sdk.core.embedding import (
    EmbeddingResponse,
    EmbeddingResult,
    EmbeddingService,
    create_embedding_service,
)

from .utils import gen_test_config, reset_global_config


def test_create_embedding_service() -> None:
    """Test creating an embedding service."""
    service = create_embedding_service("bge-large-zh", dimension=0, endpoints=["http://test-endpoint"])
    assert isinstance(service, EmbeddingService)
    assert service.config.model_name == "bge-large-zh"
    assert service.config.dimension == 0
    assert service.config.endpoints == ["http://test-endpoint"]


def test_default_endpoints_from_config() -> None:
    """Test that default endpoints are taken from global config."""
    reset_global_config()

    # Mock the environment variable for ServiceConfig
    mock_config = gen_test_config(
        embedding={
            "endpoint": "http://config-endpoint",
            "modelServingName": "test-serving",
            "baseModelName": "bge-large-zh",
        }
    )
    with mock.patch.dict(os.environ, {"ServiceConfig": str(mock_config).replace("'", '"')}):
        # Reset the global config cache to force re-reading from environment
        service = create_embedding_service("bge-large-zh", dimension=0)
        assert isinstance(service, EmbeddingService)
        assert service.config.model_name == "bge-large-zh"
        assert service.config.dimension == 0
        assert service.config.endpoints == ["http://config-endpoint"]


def test_api_key_from_config_when_arg_missing() -> None:
    """Test fallback to ServiceConfig.embedding.apiKey when api_key argument is missing."""
    reset_global_config()
    mock_config = gen_test_config(
        embedding={
            "endpoint": "http://config-endpoint",
            "modelServingName": "test-serving",
            "baseModelName": "bge-large-zh",
            "apiKey": "config-api-key",
        }
    )

    with mock.patch.dict(os.environ, {"ServiceConfig": str(mock_config).replace("'", '"')}):
        service = create_embedding_service("bge-large-zh", dimension=0)

    mock_response = mock.Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"data": [{"embedding": [0.1, 0.2]}]}
    with mock.patch("databuilder_logic_sdk.core.embedding.requests.post", return_value=mock_response) as mock_post:
        response = service.embed(["test"])

    assert len(response.results) == 1
    headers = mock_post.call_args.kwargs["headers"]
    assert headers["Authorization"] == "Bearer config-api-key"


def test_empty_api_key_from_config_does_not_add_authorization_header() -> None:
    """Test empty ServiceConfig.embedding.apiKey does not add Authorization header."""
    reset_global_config()
    mock_config = gen_test_config(
        embedding={
            "endpoint": "http://config-endpoint",
            "modelServingName": "test-serving",
            "baseModelName": "bge-large-zh",
            "apiKey": "   ",
        }
    )

    with mock.patch.dict(os.environ, {"ServiceConfig": str(mock_config).replace("'", '"')}):
        service = create_embedding_service("bge-large-zh", dimension=0)

    mock_response = mock.Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"data": [{"embedding": [0.1, 0.2]}]}
    with mock.patch("databuilder_logic_sdk.core.embedding.requests.post", return_value=mock_response) as mock_post:
        response = service.embed(["test"])

    assert len(response.results) == 1
    headers = mock_post.call_args.kwargs["headers"]
    assert "Authorization" not in headers


def test_api_key_argument_has_priority_over_service_config() -> None:
    """Test explicit api_key argument has priority over ServiceConfig.embedding.apiKey."""
    reset_global_config()
    mock_config = gen_test_config(
        embedding={
            "endpoint": "http://config-endpoint",
            "modelServingName": "test-serving",
            "baseModelName": "bge-large-zh",
            "apiKey": "config-api-key",
        }
    )

    with mock.patch.dict(os.environ, {"ServiceConfig": str(mock_config).replace("'", '"')}):
        service = create_embedding_service("bge-large-zh", dimension=0, api_key="arg-api-key")

    mock_response = mock.Mock()
    mock_response.status_code = 200
    mock_response.json.return_value = {"data": [{"embedding": [0.1, 0.2]}]}
    with mock.patch("databuilder_logic_sdk.core.embedding.requests.post", return_value=mock_response) as mock_post:
        response = service.embed(["test"])

    assert len(response.results) == 1
    headers = mock_post.call_args.kwargs["headers"]
    assert headers["Authorization"] == "Bearer arg-api-key"


def test_wrong_base_name_in_config() -> None:
    """Test that create_embedding_service raises ValueError when model name does not match config base name."""
    reset_global_config()
    # Mock the environment variable for ServiceConfig
    mock_config = gen_test_config()
    with mock.patch.dict(os.environ, {"ServiceConfig": str(mock_config).replace("'", '"')}):
        with pytest.raises(ValueError):
            create_embedding_service("bge-large-zh", dimension=0)


def test_embed_not_implemented() -> None:
    """Test that embed method raises ValueError for unimplemented service."""
    with pytest.raises(ValueError):
        create_embedding_service("unknown-model", dimension=0, endpoints=["http://test-endpoint"])


def test_mock_text_embed_7b_bf16_method() -> None:
    """Test the embed method of TextEmbed7BBF16Embedding service."""
    service = create_embedding_service("text-embed_7b_bf16", dimension=0, endpoints=["http://test-endpoint"])
    texts = ["测试文本1", "测试文本2"]
    with mock.patch.object(
        service,
        "embed",
        return_value=EmbeddingResponse(
            results=[
                EmbeddingResult(embedding=[0.1, 0.2, 0.3], index=0),
                EmbeddingResult(embedding=[0.4, 0.5, 0.6], index=1),
            ]
        ),
    ):
        response = service.embed(texts)
    assert len(response.results) == len(texts)
    for i, result in enumerate(response.results):
        assert result.index == i
        assert len(result.embedding) > 0


def test_text_embed_7b_bf16() -> None:
    """Test the embed method."""
    service = create_embedding_service("text-embed_7b_bf16", dimension=0, endpoints=["10.174.136.82:8691"])
    texts = ["", "sdf"]
    response = service.embed(texts)
    assert len(response.results) == len(texts)
    for i, result in enumerate(response.results):
        assert result.index == i


def test_bge_large_zh() -> None:
    """Test the embed method of BGE Large Chinese model."""
    service = create_embedding_service(
        "bge-large-zh",
        endpoints=["http://qianfan.baidubce.com"],
        api_key=os.getenv("BCE_API_KEY"),
    )
    texts = ["测试BGE中文模型", "再测试BGE"]
    response = service.embed(texts)
    assert len(response.results) == len(texts)
    for i, result in enumerate(response.results):
        assert result.index == i
        assert len(result.embedding) == 1024


def test_bge_large_en() -> None:
    """Test the embed method of BGE Large English model."""
    service = create_embedding_service(
        "bge-large-en",
        endpoints=["http://qianfan.baidubce.com"],
        api_key=os.getenv("BCE_API_KEY"),
    )
    texts = ["Testing BGE English model", "Testing BGE again"]
    response = service.embed(texts)
    assert len(response.results) == len(texts)
    for i, result in enumerate(response.results):
        assert result.index == i
        assert len(result.embedding) == 1024
