<?php
namespace boru\ai\Tests\Provider\Gemini;

use boru\ai\Embedding\EmbeddingRequest;
use boru\ai\Provider\Gemini\GeminiMapper;
use boru\ai\Provider\Gemini\GeminiProvider;
use boru\ai\Tests\Fakes\FakeGeminiClient;
use PHPUnit\Framework\TestCase;

class GeminiEmbeddingProviderTest extends TestCase
{
    public function testEmbedContentMapsSingleInputToNormalizedVector()
    {
        $client = new FakeGeminiClient(array(
            array(
                'embedding' => array(
                    'values' => array(0.4, 0.5),
                ),
                'usageMetadata' => array(
                    'promptTokenCount' => 2,
                ),
            ),
        ));

        $provider = new GeminiProvider($client, new GeminiMapper());
        $request = new EmbeddingRequest();
        $request->model('text-embedding-004');
        $request->input('hello');
        $request->option('task_type', 'RETRIEVAL_DOCUMENT');

        $result = $provider->embed($request);

        $this->assertSame('embedContent', $client->requests[0]['method']);
        $this->assertSame('RETRIEVAL_DOCUMENT', $client->requests[0]['payload']['taskType']);
        $this->assertSame(array(0.4, 0.5), $result->vector());
        $this->assertSame(2, $result->usage()['promptTokenCount']);
    }

    public function testBatchEmbedContentsMapsMultipleInputs()
    {
        $client = new FakeGeminiClient(array(
            array(
                'embeddings' => array(
                    array('values' => array(1.0, 1.1)),
                    array('values' => array(2.0, 2.2)),
                ),
            ),
        ));

        $provider = new GeminiProvider($client, new GeminiMapper());
        $request = new EmbeddingRequest();
        $request->model('text-embedding-004');
        $request->inputs(array('a', 'bb'));

        $result = $provider->embed($request);

        $this->assertSame('batchEmbedContents', $client->requests[0]['method']);
        $this->assertCount(2, $client->requests[0]['payload']['requests']);
        $this->assertSame(array(2.0, 2.2), $result->vector(1));
    }
}
