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

use boru\ai\AI;
use boru\ai\Config\ConfigResolver;
use boru\ai\Conversation\ConversationReplayService;
use boru\ai\Embedding\EmbeddingProviderInterface;
use boru\ai\Embedding\EmbeddingRequest;
use boru\ai\Embedding\EmbeddingResult;
use boru\ai\Provider\Capabilities;
use boru\ai\Provider\ProviderInterface;
use boru\ai\Provider\ProviderSignal;
use boru\ai\Schema\Request;
use boru\ai\Schema\Response;
use boru\ai\Stream\StreamEvent;

class GeminiProvider implements ProviderInterface, EmbeddingProviderInterface
{
    /** @var Capabilities */
    protected $capabilities;

    /** @var GeminiClient */
    protected $client;

    /** @var GeminiMapper */
    protected $mapper;

    /** @var array */
    protected $modelPrefixes = array('gemini-');

    public function __construct(GeminiClient $client = null, GeminiMapper $mapper = null, Capabilities $capabilities = null)
    {
        $this->client = $client ? $client : new GeminiClient();
        $this->mapper = $mapper ? $mapper : new GeminiMapper();
        $this->capabilities = $capabilities ? $capabilities : new Capabilities(array(
            'streaming' => true,
            'tools' => true,
            'vision' => true,
            'files' => true,
            'reasoning' => true,
            'sessionState' => false,
            'embeddings' => true,
        ));
    }

    public function name()
    {
        return 'gemini';
    }

    public function capabilities()
    {
        return $this->capabilities;
    }

    public function supportsModel($model)
    {
        $model = strtolower((string)$model);
        if ($model === '') {
            return false;
        }

        foreach ($this->modelPrefixes as $prefix) {
            if (strpos($model, strtolower($prefix)) === 0) {
                return true;
            }
        }

        return false;
    }

    public function execute(Request $request, $eventCallback = null)
    {
        $request = $this->applyReplayIfAvailable($request);

        if ($this->shouldStream($request, $eventCallback)) {
            return $this->streamResponse($request, null, array(), $eventCallback);
        }

        $payload = $this->mapper->mapRequest($request);
        $result = $this->client->generateContent($request->model(), $payload);
        return $this->mapper->mapResponse($result);
    }

    public function continueResponse(Request $request, Response $previousResponse, array $toolResults, $eventCallback = null)
    {
        $continuedRequest = $request->copy();
        foreach ($toolResults as $toolResult) {
            $continuedRequest->addToolResult($toolResult);
        }

        if ($this->shouldStream($continuedRequest, $eventCallback)) {
            return $this->streamResponse($continuedRequest, $previousResponse, $toolResults, $eventCallback);
        }

        $payload = $this->mapper->mapContinuationRequest($continuedRequest, $previousResponse, $toolResults);
        $result = $this->client->generateContent($continuedRequest->model(), $payload);
        return $this->mapper->mapResponse($result);
    }

    public function embed(EmbeddingRequest $request)
    {
        $inputs = $request->inputs();
        if (count($inputs) === 1) {
            $response = $this->client->embedContent($request->model(), $this->mapEmbeddingPayload($inputs[0], $request));
        } else {
            $requests = array();
            foreach ($inputs as $input) {
                $requests[] = $this->mapEmbeddingPayload($input, $request);
            }
            $response = $this->client->batchEmbedContents($request->model(), array('requests' => $requests));
        }

        $result = new EmbeddingResult();
        $result->provider($this->name());
        $result->model($request->model());

        if (isset($response['usageMetadata']) && is_array($response['usageMetadata'])) {
            $result->usage($response['usageMetadata']);
        }

        if (isset($response['embedding']) && is_array($response['embedding'])) {
            $result->addEmbedding($this->extractGeminiVector($response['embedding']), $inputs[0], 0);
            return $result;
        }

        $embeddings = isset($response['embeddings']) && is_array($response['embeddings']) ? $response['embeddings'] : array();
        foreach ($embeddings as $index => $embedding) {
            $input = isset($inputs[$index]) ? $inputs[$index] : null;
            $result->addEmbedding($this->extractGeminiVector($embedding), $input, $index);
        }

        return $result;
    }

    protected function shouldStream(Request $request, $eventCallback = null)
    {
        $stream = $request->option('stream');
        if ($stream !== null) {
            return $stream ? true : false;
        }
        return is_callable($eventCallback);
    }

    protected function streamResponse(Request $request, Response $previousResponse = null, array $toolResults = array(), $eventCallback = null)
    {
        if ($previousResponse instanceof Response) {
            $payload = $this->mapper->mapContinuationRequest($request, $previousResponse, $toolResults);
        } else {
            $payload = $this->mapper->mapRequest($request);
        }

        $parser = new GeminiStreamParser($this->mapper);

        $httpResponse = $this->client->streamGenerateContent($request->model(), $payload, function ($chunk) use ($parser, $eventCallback) {
            $parser->push($chunk, function (ProviderSignal $signal) use ($eventCallback) {
                $event = $this->mapSignalToStreamEvent($signal);
                if ($event && is_callable($eventCallback)) {
                    call_user_func($eventCallback, $event);
                }
            });
        });

        $error = $parser->error();
        if (is_array($error)) {
            $message = isset($error['message']) ? $error['message'] : 'Gemini streaming request failed';
            $code = isset($error['code']) ? (int)$error['code'] : $httpResponse->statusCode();
            throw new \RuntimeException($message, $code);
        }

        $final = $parser->finalResponse();
        if ($final instanceof Response) {
            return $final;
        }

        $decoded = $httpResponse->json();
        if (is_array($decoded) && isset($decoded['error'])) {
            $message = isset($decoded['error']['message']) ? $decoded['error']['message'] : 'Gemini streaming request failed';
            $code = isset($decoded['error']['code']) ? (int)$decoded['error']['code'] : $httpResponse->statusCode();
            throw new \RuntimeException($message, $code);
        }

        if (is_array($decoded)) {
            return $this->mapper->mapResponse($decoded);
        }

        $text = $parser->currentText();
        if ($httpResponse->isSuccess() && $text !== '') {
            $response = new Response();
            $response->provider('gemini');
            $response->model($request->model());
            $response->text($text);
            $response->status('completed');
            return $response;
        }

        throw new \RuntimeException('Gemini streaming response did not produce a valid final response', $httpResponse->statusCode());
    }

    protected function applyReplayIfAvailable(Request $request)
    {
        $conversation = $request->metadata('conversation');
        if (!$conversation instanceof \boru\ai\Conversation\Conversation) {
            return $request;
        }

        $resolver = new ConfigResolver(AI::config());
        $service = new ConversationReplayService();
        $options = $service->resolvedOptions('gemini', $resolver, array(
            'enabled' => true,
            'strategy' => 'full',
            'max_turns' => 10,
        ));
        if (empty($options['enabled'])) {
            return $request;
        }

        return $service->applyReplay($request, $conversation, $options);
    }

    protected function mapSignalToStreamEvent(ProviderSignal $signal)
    {
        $type = $signal->type();

        if ($type === ProviderSignal::TEXT_DELTA) {
            return new StreamEvent(StreamEvent::TEXT_DELTA, $signal->data(), $signal->response());
        }

        if ($type === ProviderSignal::TEXT_COMPLETED) {
            return new StreamEvent(StreamEvent::TEXT_COMPLETED, $signal->data(), $signal->response());
        }

        if ($type === ProviderSignal::RESPONSE_ERROR) {
            return new StreamEvent(StreamEvent::TURN_ERROR, $signal->data(), $signal->response());
        }

        return null;
    }

    protected function mapEmbeddingPayload($input, EmbeddingRequest $request)
    {
        $payload = array(
            'model' => 'models/' . $request->model(),
            'content' => array(
                'parts' => array(
                    array(
                        'text' => (string)$input,
                    ),
                ),
            ),
        );

        if ($request->option('task_type') !== null) {
            $payload['taskType'] = $request->option('task_type');
        }

        if ($request->option('title') !== null) {
            $payload['title'] = $request->option('title');
        }

        if ($request->option('output_dimensionality') !== null) {
            $payload['outputDimensionality'] = (int)$request->option('output_dimensionality');
        }

        return $payload;
    }

    protected function extractGeminiVector(array $embedding)
    {
        if (isset($embedding['values']) && is_array($embedding['values'])) {
            return $embedding['values'];
        }
        return array();
    }
}
