"""Cheap smoke tests for POST /v1/embeddings (always run; total cost < $0.00001)."""
from __future__ import annotations

import base64
import math
import struct

import pytest

PRICE = {"text-embedding-3-small": 0.02, "text-embedding-3-large": 0.13, "text-embedding-ada-002": 0.10}
NATIVE_DIMS = {"text-embedding-3-small": 1536, "text-embedding-3-large": 3072, "text-embedding-ada-002": 1536}


@pytest.mark.parametrize("model", list(NATIVE_DIMS))
def test_embedding_shape_and_norm(openai, model):
    st, body, _ = openai("POST", "/v1/embeddings", {"model": model, "input": "OK"},
                         est_cost_usd=PRICE[model] * 2 / 1e6, note=f"test_embeddings {model}")
    assert st == 200, body
    assert body["object"] == "list" and body["data"][0]["object"] == "embedding" and body["data"][0]["index"] == 0
    vec = body["data"][0]["embedding"]
    assert len(vec) == NATIVE_DIMS[model]
    assert abs(math.sqrt(sum(x * x for x in vec)) - 1.0) < 1e-3, "embeddings are documented as unit-length"
    assert body["usage"]["prompt_tokens"] >= 1 and body["usage"]["total_tokens"] == body["usage"]["prompt_tokens"]


def test_dimensions_and_base64(openai):
    st, body, _ = openai("POST", "/v1/embeddings", {"model": "text-embedding-3-small", "input": "OK", "dimensions": 16,
                                                    "encoding_format": "base64"}, est_cost_usd=1e-7, note="test dims=16 base64")
    assert st == 200, body
    raw = base64.b64decode(body["data"][0]["embedding"])
    assert len(raw) == 16 * 4
    floats = struct.unpack("<16f", raw)
    assert abs(math.sqrt(sum(x * x for x in floats)) - 1.0) < 1e-3, "shortened vectors are re-normalized"


def test_dimensions_rejected_on_ada(openai):
    st, body, _ = openai("POST", "/v1/embeddings", {"model": "text-embedding-ada-002", "input": "OK", "dimensions": 16},
                         note="test dimensions on ada-002 (expect 400)")
    assert st == 400
    assert body["error"]["type"] == "invalid_request_error"
    assert "dimensions" in body["error"]["message"]


def test_array_input_preserves_order(openai):
    st, body, _ = openai("POST", "/v1/embeddings", {"model": "text-embedding-3-small", "input": ["OK", "Reply with OK."], "dimensions": 8},
                         est_cost_usd=2e-7, note="test array input")
    assert st == 200, body
    assert [d["index"] for d in body["data"]] == [0, 1]
    assert all(len(d["embedding"]) == 8 for d in body["data"])