<?php
class AimogenOllamaAPI 
{
    private $baseUrl;
    public function __construct($baseUrl) 
    {
        $this->baseUrl = rtrim($baseUrl, '/');
    }
    private function makeRequest($endpoint, $data, $stream) 
    {
        $url = $this->baseUrl . $endpoint;
        
        $payload = json_encode($data);
        if ($payload === false) {
            aimogen_log_to_file('Failed to encode payload in request: ' . json_encode($data));
            return false;
        }

        $args = array(
            'body'        => $payload,
            'headers'     => array(
                'Content-Type' => 'application/json',
            ),
            'method'      => 'POST',
            'data_format' => 'body',
            'timeout'     => AIMOGEN_DEFAULT_BIG_TIMEOUT,
        );

        if ($stream === true) {
            add_action('http_api_curl', array($this, 'filterCurlForStream'));
        }

        $response = wp_remote_post($url, $args);

        if ($stream === true) {
            remove_action('http_api_curl', array($this, 'filterCurlForStream'));
        }

        if (is_wp_error($response)) {
            aimogen_log_to_file('Error making request to ' . $url . ': ' . $response->get_error_message());
            return false;
        }

        $body = wp_remote_retrieve_body($response);

        if ($stream === true) {
            return '';
        }
        $res = json_decode($body, true);
        if ($res === null) {
            aimogen_log_to_file('Failed to decode response: ' . $body);
            return false;
        }

        return $res;
    }
    public function filterCurlForStream($handle)
    {
        //curl is required for streaming - using CURLOPT_WRITEFUNCTION
        curl_setopt($handle, CURLOPT_SSL_VERIFYPEER, false);// phpcs:ignore WordPress.WP.AlternativeFunctions.curl_curl_setopt
        curl_setopt($handle, CURLOPT_SSL_VERIFYHOST, false);// phpcs:ignore WordPress.WP.AlternativeFunctions.curl_curl_setopt
        curl_setopt($handle, CURLOPT_WRITEFUNCTION, function ($curl_info, $data) // phpcs:ignore WordPress.WP.AlternativeFunctions.curl_curl_setopt
        {
            $my_copy_data = trim($data);
            $prefix = 'data: ';
            if (substr($my_copy_data, 0, strlen($prefix)) == $prefix) {
                $my_copy_data = substr($my_copy_data, strlen($prefix));
            }
            $suffix = 'data: [DONE]';
            $needle_length = strlen($suffix);
            if (substr($my_copy_data, -$needle_length) === $suffix) {
                $my_copy_data = substr($my_copy_data, 0, -$needle_length);
            }
            $my_copy_data = trim($my_copy_data);
            $response = json_decode($my_copy_data, true);
            if (isset($response['error']) && !empty($response['error'])) 
            {
                $message = isset($response['error']['message']) && !empty($response['error']['message']) ? $response['error']['message'] : '';
                if (empty($message) && isset($response['error']['code']) && $response['error']['code'] == 'invalid_api_key') {
                    $message = "Incorrect API key provided. You can find your API key at https://platform.openai.com/account/api-keys.";
                }
                echo "event: message\n";
                
                // added wp_strip_all_tags now!
                echo 'data: {"error":[{"message":"' . esc_html(wp_strip_all_tags($message)) . '"}]}';
                echo "\n\n";
                $l1 = ob_get_length();
                if($l1 === false)
                {
                    $l1 = 0;
                }
                if (ob_get_length())
                {
                    ob_end_flush();
                }
                flush();
                echo 'data: {"choices":[{"finish_reason":"stop"}]}';
                echo "\n\n";
                $l2 = ob_get_length();
                if($l2 === false)
                {
                    $l2 = 0;
                }
                if (ob_get_length())
                {
                    ob_end_flush();
                }
                flush();
                return $l1 + $l2;
            } 
            else 
            {
                //not escaped because of streaming
                // phpcs:ignore WordPress.Security.EscapeOutput.OutputNotEscaped
                echo "data: {$data}\n\n";
                if (ob_get_length())
                {
                    ob_flush();
                }
                flush();
                return strlen($data);
            }
        });
    }
    public function chatCompletion($model, $messages, $stream = true, $options = [], $functions = false, $toolOutputs = false, $vision_file = false) 
    {
        if(!empty($vision_file) && isset($messages[0]))
        {
            if(is_array($messages[0]['content']) && isset($messages[0]['content'][0]['text']))
            {
                $messages[0]['content'] = $messages[0]['content'][0]['text'];
            }
            $base64_vision = '';
            if(!empty($base64_vision))
            {
                $vision_file = $base64_vision;
            }
            $messages[0]['images'] = array($vision_file);
        }
        if(!empty($toolOutputs) && is_array($toolOutputs))
        {
            foreach($toolOutputs as $tout)
            {
                if(isset($tout['message']))
                {
                    $messages[] = $tout['message'];
                }
                if(isset($tout['output']))
                {
                    $tool_rez = ['role' => 'tool', 'content' => $tout['output']];
                    $messages[] = $tool_rez;
                }
            }
        }
        $data = [
            'model' => $model,
            'messages' => $messages,
            'stream' => $stream
        ];
        if(!empty($functions))
        {
            if($functions !== false && !empty($functions) && isset($functions['functions']) && !empty($functions['functions']))
            {
                $data['tools'] = $functions['functions'];
            }
        }
        if(!empty($options))
        {
            $data['options'] = $options;
        }
        $resp = $this->makeRequest('/api/chat', $data, $stream);
        if($resp === false)
        {
            return false;
        }
        if($stream === true)
        {
            return '';
        }
        if(!isset($resp['message']['content']))
        {
            aimogen_log_to_file('Failed to interpret Ollama chat API response: ' . json_encode($resp));
            return false;
        }
        return $resp;
    }
    public function embeddings($model, $prompt) 
    {
        $data = [
            'model' => $model,
            'prompt' => $prompt
        ];
        $resp = $this->makeRequest('/api/embeddings', $data, false);
        if($resp === false)
        {
            return false;
        }
        if(!isset($resp['embedding']) || empty($resp['embedding']))
        {
            aimogen_log_to_file('Failed to interpret Ollama embedding API response: ' . json_encode($resp));
            return false;
        }
        return $resp['embedding'];
    }
    public function generate($model, $prompt, $stream = false, $options = []) 
    {
        $data = [
            'model' => $model,
            'prompt' => $prompt,
            'stream' => $stream
        ];
        if(!empty($options))
        {
            $data['options'] = $options;
        }
        $resp = $this->makeRequest('/api/generate', $data, $stream);
        if($resp === false)
        {
            return false;
        }
        if($stream === true)
        {
            return '';
        }
        if(!isset($resp['response']))
        {
            aimogen_log_to_file('Failed to interpret Ollama completion API response: ' . json_encode($resp));
            return false;
        }
        return $resp['response'];
    }
}
?>