"""测试 Embedding API"""

import pytest
from pydantic import ValidationError
from octen.models.embedding import EmbeddingRequest, EmbeddingResponse


def test_embedding_request_single_input():
    """测试单个输入"""
    request = EmbeddingRequest(input="Hello, world!")
    assert request.input == "Hello, world!"
    assert request.truncation is True  # 默认值


def test_embedding_request_batch_input():
    """测试批量输入"""
    texts = ["text 1", "text 2", "text 3"]
    request = EmbeddingRequest(input=texts)
    assert request.input == texts
    assert len(request.input) == 3


def test_embedding_request_with_options():
    """测试带选项的请求"""
    request = EmbeddingRequest(
        input="test",
        model="octen-embedding-4b",
        dimension=768,
        input_type="query",
        truncation=False,
    )
    assert request.model == "octen-embedding-4b"
    assert request.dimension == 768
    assert request.input_type == "query"
    assert request.truncation is False


def test_embedding_request_validation():
    """测试请求验证"""
    # 无效的模型
    with pytest.raises(ValidationError):
        EmbeddingRequest(input="test", model="invalid-model")

    # 无效的输入类型
    with pytest.raises(ValidationError):
        EmbeddingRequest(input="test", input_type="invalid")

    # 无效的维度
    with pytest.raises(ValidationError):
        EmbeddingRequest(input="test", dimension=0)


def test_embedding_response():
    """测试嵌入响应"""
    response_data = {
        "code": 0,
        "msg": "success",
        "data": {
            "model": "octen-embedding-4b",
            "results": [
                {"index": 0, "embedding": [0.1, 0.2, 0.3, 0.4, 0.5]},
                {"index": 1, "embedding": [0.6, 0.7, 0.8, 0.9, 1.0]},
            ],
        },
        "meta": {
            "usage": {"total_tokens": 50},
        },
    }

    response = EmbeddingResponse(**response_data)
    assert response.code == 0
    assert response.msg == "success"
    assert response.model == "octen-embedding-4b"
    assert len(response.results) == 2


def test_embedding_response_get_embeddings():
    """测试获取嵌入向量"""
    response_data = {
        "code": 0,
        "msg": "success",
        "data": {
            "results": [
                {"index": 0, "embedding": [0.1, 0.2, 0.3]},
                {"index": 1, "embedding": [0.4, 0.5, 0.6]},
            ],
        },
        "meta": {},
    }

    response = EmbeddingResponse(**response_data)
    embeddings = response.get_embeddings()

    assert len(embeddings) == 2
    assert embeddings[0] == [0.1, 0.2, 0.3]
    assert embeddings[1] == [0.4, 0.5, 0.6]


def test_embedding_response_get_first_embedding():
    """测试获取第一个嵌入向量"""
    response_data = {
        "code": 0,
        "msg": "success",
        "data": {
            "results": [
                {"index": 0, "embedding": [0.1, 0.2, 0.3]},
            ],
        },
        "meta": {},
    }

    response = EmbeddingResponse(**response_data)
    first_embedding = response.get_first_embedding()

    assert first_embedding == [0.1, 0.2, 0.3]


def test_embedding_response_empty():
    """测试空响应"""
    response_data = {
        "code": 0,
        "msg": "success",
        "data": {"results": []},
        "meta": {},
    }

    response = EmbeddingResponse(**response_data)
    assert response.get_embeddings() == []
    assert response.get_first_embedding() is None
