package router

import (
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestBuildEmbeddingResponse(t *testing.T) {
	response := buildEmbeddingResponse([][]float32{
		{0.5165501, 0.16796638, 0.60452324, 0.2851341, -0.07298655, 0.26138917, 0.06126004, 0.24445628, -0.33593276, -0.09055198},
	}, 123)
	jsonBytes, err := json.Marshal(response)
	require.NoError(t, err)
	assert.JSONEq(t, `
	{
		"embeddings":["oTwEP2H/Kz4Jwho/Gf2RPvl5lb3N1IU+z+t6Pb9Sej5h/6u+UXO5vQ=="],
		"usage":{"prompt_tokens":123,"total_tokens":123}
	}`, string(jsonBytes))

	response = buildEmbeddingResponse([][]float32{
		{0.35766974, -0.11206985, 0.5296162, 0.047424357, 0.0027032215, 0.45516786, 0.2433479, 0.46629536, -0.30070752, -0.050106883},
		{-0.18430692, 0.16802885, 0.6041567, 0.22878253, -0.21028067, -0.08325834, 0.09722368, -0.1493491, -0.3141757, -0.58209676},
		{-0.2690117, -0.31755385, 0.1978754, 0.21958698, 0.28966418, -0.565207, -0.41446158, 0.14880371, 0.1182663, -0.3579762},
	}, 123)
	jsonBytes, err = json.Marshal(response)
	require.NoError(t, err)
	assert.JSONEq(t, `
	{
		"embeddings":[
			"fSC3PuGE5b3tlAc/C0BCPYgoMTvDC+k+MTB5PkS+7j5W9pm+4DxNvQ==",
			"9Lo8vsEPLD4Dqho/+EVqPtFTV75Zg6q9NR3HPfjuGL6j26C+SwQVvw==",
			"57uJvmuWor7Zn0o+adtgPt1OlD5osRC/TzTUvgBgGD6aNfI9qEi3vg=="
		],
		"usage":{"prompt_tokens":123,"total_tokens":123}
	}`, string(jsonBytes))

	response = buildEmbeddingResponse([][]float32{
		{},
	}, 456)
	jsonBytes, err = json.Marshal(response)
	require.NoError(t, err)
	assert.JSONEq(t, `
	{
		"embeddings":[""],
		"usage":{"prompt_tokens":456,"total_tokens":456}
	}`, string(jsonBytes))

	response = buildEmbeddingResponse([][]float32{}, 789)
	jsonBytes, err = json.Marshal(response)
	require.NoError(t, err)
	assert.JSONEq(t, `
	{
		"embeddings":[],
		"usage":{"prompt_tokens":789,"total_tokens":789}
	}`, string(jsonBytes))
}