<?php
namespace boru\boruai\Openai;

use boru\boruai\BoruAI;
use boru\boruai\Openai\Api\Endpoints\EmbeddingsAPI;
use boru\boruai\Openai\embeddings\Document;
use boru\boruai\Openai\Models\BaseModel;
use boru\boruai\Openai\OpenAI;

class OpenAIEmbedding extends BaseModel {
    private $model = null;
    private $embedding = [];
    private $text;

    private static $batchSize = 50;

    /**
     * Generate an embedding for a given text using a given model. Retrieve the embedding using the embedding() method.
     * @param string|array|null $text A string or an array of strings to embed
     * @param string|null $model The model to use for the embedding, defaults to text-embedding-ada-002
     */
    public function __construct($text=null,$model=null) {
        if($model !== null) {
            $this->model($model);
        } else {
            $this->model(BoruAI::defaultModel("embedding","text-embedding-3-small"));
        }
        if($text !== null) {
            $this->text($text);
        }
    }

    /**
     * Set the model and text from an array
     * @param array $array
     * @param bool $embed Whether to embed the text after setting the model and text
     */
    public function fromArray($array,$embed=false) {
        if(isset($array["model"])) {
            $this->model = $array["model"];
        }
        if(isset($array["text"])) {
            $this->text = $array["text"];
        }
        if(isset($array["embedding"])) {
            $this->embedding = $array["embedding"];
        }
        if($embed) {
            $this->embed();
        }
    }

    public function toArray() {
        return [
            "model" => $this->model,
            "text" => $this->text,
            "embedding" => $this->embedding
        ];
    }

    /**
     * @param string|null $model
     * @return string
     */
    public function model($model=null) {
        if ($model!==null) {
            $this->model = $model;
        }
        return $this->model;
    }
    /**
     * @param array|null $embedding
     * @return array
     */
    public function embedding($embedding=null) {
        if ($embedding!==null) {
            $this->embedding = $embedding;
        }
        if($this->embedding === null && $this->text !== null) {
            $this->embed();
        }
        return $this->embedding;
    }
    /**
     * @param string|array|null $text
     * @return string|array
     */
    public function text($text=null) {
        if ($text!==null) {
            $this->text = OpenAI::toUtf8($text);
            $this->embedding = null;
            $this->embed();
        }
        return $this->text;
    }

    public function setEmbedding($embedding) {
        $this->embedding = $embedding;
    }
    
    private function embed() {
        if($this->text === null) {
            return false;
        }
        if(is_array($this->text)) {
            return $this->embedMulti();
        } else {
            return $this->embedSingle();
        }
    }
    private function embedSingle() {
        $parameters = [
            "model" => $this->model(),
            "input" => $this->text()
        ];
        try {
            $response = EmbeddingsAPI::create($parameters);
            //print_r($response);
            if(($data = $response->get("data",false)) !== false) {
                if(empty($data)) {
                    throw new \Exception("No data returned from OpenAI");
                }
                return $this->embedding($data[0]["embedding"]);
            }
        } catch (\Exception $e) {
            throw $e;
            return false;
        }
    }
    private function embedMulti() {
        $texts = $this->text();
        $embeddings = [];
        $batchSize = self::batchSize();
        $batches = array_chunk($texts,$batchSize);
        foreach($batches as $chunkIndex=>$batch) {
            $parameters = [
                "model" => $this->model(),
                "input" => $batch
            ];
            try {
                $response = EmbeddingsAPI::create($parameters);
                if(($data = $response->get("data",false)) !== false) {
                    if(empty($data)) {
                        throw new \Exception("No data returned from OpenAI");
                    }
                    foreach($data as $index=>$item) {
                        $idx = $chunkIndex * $batchSize + $index;
                        $embeddings[$idx] = $item["embedding"];
                    }
                }
            } catch (\Exception $e) {
                throw $e;
                return false;
            }
        }
        return $this->embedding($embeddings);
    }

    public static function batchSize($batchSize=null) {
        if ($batchSize!==null) {
            self::$batchSize = $batchSize;
        }
        return self::$batchSize;
    }

    public static function create($text,$returnObject=false) {
        $embedding = new OpenAIEmbedding($text);
        if($returnObject) {
            return $embedding;
        }
        return $embedding->embedding();
    }

    public static function embedString($string,$model=null) {
        $embedding = new OpenAIEmbedding($string,$model);
        return $embedding->embedding();
    }

    /**
     * Embed a document and store the embedding in the document object
     * @param Document $document
     * @param bool $format Whether to use the formatted content for the embedding
     * @param string|null $model The model to use for the embedding, defaults to text-embedding-ada-002
     * @return static
     */
    public static function embedDocument($document,$format=true,$model=null) {
        $document->embed($format,$model);
        return $document->embedding();
    }

    /**
     * Embed a list of documents and store the embeddings in the document objects
     * @param Document[] $documents
     * @param bool $format Whether to use the formatted content for the embedding
     * @param string|null $model The model to use for the embedding, defaults to text-embedding-ada-002
     * @return static[]
     */
    public static function embedDocuments($documents,$format=true,$model=null) {
        $embeddings = [];
        foreach($documents as $document) {
            $embeddings[] = self::embedDocument($document,$format,$model);
        }
        return $embeddings;
    }
}