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

use boru\ai\Asset\Attachment;
use boru\ai\Schema\ContentPart;
use boru\ai\Schema\Message;
use boru\ai\Schema\Request;
use boru\ai\Schema\Response;
use boru\ai\Schema\ToolCall;
use boru\ai\Schema\ToolResult;

class GeminiMapper
{
    /** @var GeminiFunctionSchemaNormalizer|null */
    protected $functionSchemaNormalizer;

    public function functionSchemaNormalizer(GeminiFunctionSchemaNormalizer $normalizer = null)
    {
        if ($normalizer !== null) {
            $this->functionSchemaNormalizer = $normalizer;
            return $this;
        }

        if (!$this->functionSchemaNormalizer) {
            $this->functionSchemaNormalizer = new GeminiFunctionSchemaNormalizer();
        }

        return $this->functionSchemaNormalizer;
    }

    public function mapRequest(Request $request)
    {
        return $this->buildRequestPayload($request, null, array());
    }

    public function mapContinuationRequest(Request $request, Response $previousResponse, array $toolResults)
    {
        return $this->buildRequestPayload($request, $previousResponse, $toolResults);
    }

    protected function buildRequestPayload(Request $request, Response $previousResponse = null, array $toolResults = array())
    {
        $payload = array(
            'contents' => $this->mapMessages($request),
        );

        $attachmentContents = $this->mapAttachments($request);
        if (!empty($attachmentContents)) {
            $payload['contents'] = array_merge($payload['contents'], $attachmentContents);
        }

        if ($request->instructions()) {
            $payload['systemInstruction'] = array(
                'parts' => array(
                    array(
                        'text' => $request->instructions(),
                    ),
                ),
            );
        }

        $tools = $this->mapTools($request->tools());
        if (!empty($tools)) {
            $payload['tools'] = $tools;
        }

        if ($previousResponse instanceof Response) {
            $previousContent = $this->extractCandidateContent($previousResponse->raw());
            if (!empty($previousContent)) {
                $payload['contents'][] = $previousContent;
            }
        }

        if (!empty($toolResults)) {
            $toolResultContents = $this->mapToolResults($toolResults);
        } else {
            $toolResultContents = $this->mapToolResults($request->toolResults());
        }

        if (!empty($toolResultContents)) {
            $payload['contents'] = array_merge($payload['contents'], $toolResultContents);
        }

        foreach ($request->options() as $key => $value) {
            if ($value === null || $key === 'previous_response_id' || $key === 'stream') {
                continue;
            }
            $payload[$key] = $value;
        }

        return $payload;
    }

    protected function mapMessages(Request $request)
    {
        $contents = array();

        foreach ($request->messages() as $message) {
            if (!$message instanceof Message) {
                continue;
            }

            $parts = array();
            foreach ($message->parts() as $part) {
                if (!$part instanceof ContentPart) {
                    continue;
                }

                if ($part->type() === 'text') {
                    $parts[] = array(
                        'text' => (string)$part->value(),
                    );
                } elseif ($part->type() === 'image') {
                    $parts[] = array(
                        'fileData' => array(
                            'mimeType' => 'image/*',
                            'fileUri' => (string)$part->value(),
                        ),
                    );
                } elseif ($part->type() === 'file') {
                    $parts[] = array(
                        'fileData' => array(
                            'mimeType' => 'application/octet-stream',
                            'fileUri' => (string)$part->value(),
                        ),
                    );
                }
            }

            if (empty($parts)) {
                continue;
            }

            $contents[] = array(
                'role' => $this->mapRole($message->role()),
                'parts' => $parts,
            );
        }

        return $contents;
    }

    protected function mapAttachments(Request $request)
    {
        $contents = array();

        foreach ($request->attachments() as $attachment) {
            if (!$attachment instanceof Attachment) {
                continue;
            }

            $binding = $request->findAssetBinding($attachment->assetId(), 'gemini');
            if (!$binding || !$binding->providerUri()) {
                continue;
            }

            $mimeType = $binding->mimeType() ? $binding->mimeType() : 'application/octet-stream';
            $contents[] = array(
                'role' => 'user',
                'parts' => array(
                    array(
                        'fileData' => array(
                            'mimeType' => $mimeType,
                            'fileUri' => $binding->providerUri(),
                        ),
                    ),
                ),
            );
        }

        return $contents;
    }

    protected function mapTools(array $tools)
    {
        $functionDeclarations = array();

        foreach ($tools as $tool) {
            if (!is_array($tool) || !isset($tool['name'])) {
                continue;
            }

            $schema = isset($tool['schema']) && is_array($tool['schema'])
                ? $tool['schema']
                : array(
                    'type' => 'object',
                    'properties' => array(),
                );

            $functionDeclarations[] = array(
                'name' => $tool['name'],
                'description' => isset($tool['description']) ? $tool['description'] : '',
                'parameters' => $this->functionSchemaNormalizer()->normalize($schema),
            );
        }

        if (empty($functionDeclarations)) {
            return array();
        }

        return array(
            array(
                'functionDeclarations' => $functionDeclarations,
            ),
        );
    }

    protected function mapToolResults(array $toolResults)
    {
        $contents = array();

        foreach ($toolResults as $toolResult) {
            if (!$toolResult instanceof ToolResult) {
                continue;
            }

            $part = array(
                'functionResponse' => array(
                    'name' => $toolResult->name(),
                    'response' => $this->normalizeToolResultOutput($toolResult->output()),
                ),
            );

            if ($toolResult->toolCallId()) {
                $part['functionResponse']['id'] = $toolResult->toolCallId();
            }

            $contents[] = array(
                'role' => 'user',
                'parts' => array($part),
            );
        }

        return $contents;
    }

    public function mapResponse(array $payload)
    {
        $response = new Response();
        $response->provider('gemini');

        if (isset($payload['responseId'])) {
            $response->id($payload['responseId']);
        }
        if (isset($payload['modelVersion'])) {
            $response->model($payload['modelVersion']);
        }
        if (isset($payload['usageMetadata']) && is_array($payload['usageMetadata'])) {
            $response->usage($this->normalizeUsage($payload['usageMetadata']));
        }
        if (isset($payload['error'])) {
            $response->error($payload['error']);
        }

        $response->raw($payload);
        $response->text($this->extractText($payload));
        $this->extractToolCalls($response, $payload);

        if (!$response->status()) {
            $response->status($response->hasToolCalls() ? 'tool_calling' : 'completed');
        }

        return $response;
    }

    protected function extractText(array $payload)
    {
        $parts = $this->extractCandidateParts($payload);
        $texts = array();

        foreach ($parts as $part) {
            if (is_array($part) && isset($part['text'])) {
                $texts[] = $part['text'];
            }
        }

        return implode("\n", $texts);
    }

    protected function extractToolCalls(Response $response, array $payload)
    {
        $parts = $this->extractCandidateParts($payload);

        foreach ($parts as $part) {
            if (!is_array($part) || !isset($part['functionCall']) || !is_array($part['functionCall'])) {
                continue;
            }

            $functionCall = $part['functionCall'];
            $toolCall = new ToolCall();

            if (isset($functionCall['id'])) {
                $toolCall->id($functionCall['id']);
            } elseif (isset($functionCall['name'])) {
                $toolCall->id($functionCall['name']);
            }

            if (isset($functionCall['name'])) {
                $toolCall->name($functionCall['name']);
            }

            $arguments = isset($functionCall['args']) && is_array($functionCall['args']) ? $functionCall['args'] : array();
            $toolCall->arguments($arguments);
            $toolCall->argumentsJson(json_encode($arguments));
            $toolCall->status('requested');
            $response->addToolCall($toolCall);
        }
    }

    protected function extractCandidateParts(array $payload)
    {
        if (!isset($payload['candidates']) || !is_array($payload['candidates'])) {
            return array();
        }

        $parts = array();
        foreach ($payload['candidates'] as $candidate) {
            if (!is_array($candidate) || !isset($candidate['content']['parts']) || !is_array($candidate['content']['parts'])) {
                continue;
            }

            foreach ($candidate['content']['parts'] as $part) {
                $parts[] = $part;
            }
        }

        return $parts;
    }

    protected function extractCandidateContent(array $payload)
    {
        if (!isset($payload['candidates'][0]['content']) || !is_array($payload['candidates'][0]['content'])) {
            return array();
        }

        $content = $payload['candidates'][0]['content'];
        if (!isset($content['role'])) {
            $content['role'] = 'model';
        }
        return $content;
    }

    protected function normalizeToolResultOutput($output)
    {
        if (is_array($output)) {
            return $output;
        }

        $decoded = json_decode((string)$output, true);
        if (is_array($decoded)) {
            return $decoded;
        }

        if (is_scalar($output)) {
            return array('result' => (string)$output);
        }

        return array('result' => json_encode($output));
    }

    protected function mapRole($role)
    {
        $role = strtolower((string)$role);
        if ($role === 'assistant') {
            return 'model';
        }
        return 'user';
    }

    protected function normalizeUsage(array $usage)
    {
        $normalized = array(
            'raw' => $usage,
        );

        if (isset($usage['promptTokenCount'])) {
            $normalized['input_tokens'] = (int)$usage['promptTokenCount'];
        }
        if (isset($usage['candidatesTokenCount'])) {
            $normalized['output_tokens'] = (int)$usage['candidatesTokenCount'];
        }
        if (isset($usage['totalTokenCount'])) {
            $normalized['total_tokens'] = (int)$usage['totalTokenCount'];
        }

        return $normalized;
    }
}
