"""chat_serialize：多模态消息 token 粗估。"""

from __future__ import annotations

from aicd.agent.chat_serialize import (
    VISION_IMAGE_TOKENS_ESTIMATE,
    estimate_message_tokens,
)


def test_estimate_message_tokens_plain_string() -> None:
    msg = {"role": "user", "content": "abcd"}
    assert estimate_message_tokens(msg) == 2  # 4 // 4 + 1


def test_estimate_message_tokens_multimodal_one_image() -> None:
    msg = {
        "role": "user",
        "content": [
            {"type": "text", "text": "x"},
            {"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}},
        ],
    }
    # text: 1 char -> 1//4+1 = 1; one image -> + VISION_IMAGE_TOKENS_ESTIMATE
    assert estimate_message_tokens(msg) == 1 + VISION_IMAGE_TOKENS_ESTIMATE


def test_estimate_message_tokens_multimodal_two_images() -> None:
    msg = {
        "role": "user",
        "content": [
            {"type": "image_url", "image_url": {"url": "http://a"}},
            {"type": "image_url", "image_url": {"url": "http://b"}},
        ],
    }
    assert estimate_message_tokens(msg) == 2 * VISION_IMAGE_TOKENS_ESTIMATE
