<?php

namespace rednaoeasycalculationforms\Managers\AIManagers\APICaller;

use rednaoeasycalculationforms\core\Exception\FriendlyException;
use rednaoeasycalculationforms\core\Managers\LogManager\LogManager;
use rednaoeasycalculationforms\Utilities\Sanitizer;

class Gemini2Caller extends APiCallerBase
{
    public function Chat($message, $apiKey = '')
    {
        $this->MaybeLogMessage($message);

        // Endpoint for Gemini 2.0 Flash with API key as query parameter
        if ($apiKey == '')
            $apiKey = $this->GetAPIKey();

        $url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key=" . urlencode($apiKey);

        // Check if the first message is a system instruction
        $systemInstruction = null;
        $contentsMessages = $message;

        if (!empty($message) && isset($message[0]->role) && $message[0]->role === 'system') {
            $systemInstruction = [
                "parts" => [
                    ["text" => $message[0]->content]
                ]
            ];
            // Remove the system message from the contents array
            $contentsMessages = array_slice($message, 1);
        }

        // Map the remaining role-content array to Gemini's "contents" format
        $contents = array_map(function ($msg) {
            $role = $msg->role;
            if($role == 'assistant')
                $role = 'model';
            return [
                "role" => $role,
                "parts" => [
                    [
                        "text" => $msg->content
                    ]
                ]
            ];
        }, $contentsMessages);

        // Prepare the request body
        $requestBody = [
            "contents" => $contents,
            "generationConfig" => [
                "temperature" => 0.5,
                "maxOutputTokens" => 8000,
                "responseMimeType" => "application/json" // For JSON response format
            ]
        ];

        // Add systemInstruction if it exists
        if ($systemInstruction !== null) {
            $requestBody["systemInstruction"] = $systemInstruction;
        }

        // Make the API request using WordPress HTTP API
        $response = \wp_remote_post($url, [
            'headers' => [
                'Content-Type' => 'application/json', // Only Content-Type needed now
            ],
            'body' => json_encode($requestBody),
            'timeout' => 60
        ]);

        // Check if the request failed
        if (is_wp_error($response)) {
            throw new FriendlyException($response->get_error_message());
        }

        // Get HTTP response code
        $httpCode = wp_remote_retrieve_response_code($response);
        $body = wp_remote_retrieve_body($response);
        $decodedBody = json_decode($body, true); // Convert JSON response to array

        // If the API returns an error message, use it
        $errorMessage = Sanitizer::GetStringValueFromPath($decodedBody, ['error', 'message'], 'Invalid API Key');

        // Check if API Key is valid (Gemini returns 200 for success)
        if ($httpCode !== 200) {
            throw new FriendlyException("The API key does not seem to be a valid Gemini API key. Error: " . $errorMessage);
        }

        // Extract the response content
        $candidates = Sanitizer::GetValueFromPath($decodedBody, ['candidates'], null);
        if ($candidates == null || empty($candidates)) {
            throw new FriendlyException("There was no response, please try again");
        }

        $content = $candidates[0]['content']['parts'][0]['text'];

        $usageMetadata = Sanitizer::GetValueFromPath($decodedBody, ['usageMetadata'], null);
        $inputTokenCount = $usageMetadata->promptTokenCount ?? 0; // Input tokens
        $outputTokenCount = $usageMetadata->candidatesTokenCount ?? 0; // Output tokens
        $totalTokenCount = $usageMetadata->totalTokenCount ?? 0; // Total tokens

        LogManager::LogDebug("Gemini 2.0 Flash Response:\r\n" . $content);
        LogManager::LogDebug($this->GetModelName(). ' Usage: Input Tokens: ' . $inputTokenCount . ' Output Tokens: ' . $outputTokenCount . ' Total Tokens: ' . $totalTokenCount);
        return $content;
    }

    public function GetModelName()
    {
        return 'gemini2';
    }
}