<?php
namespace boru\boruai\Tools;

use boru\boruai\Openai\OpenAIAssistant;
use boru\boruai\Openai\OpenAIChat;
use boru\boruai\Tiktoken\EncoderProvider;
use boru\boruai\Tiktoken\Encoder;

class Encode {
    /** @var EncoderProvider */
    private $provider;

    /** @var Encoder */
    private $encoder;

    private $text = "";
    private $tokens = [];

    public function __construct($chatOrAssistant=null) {
        $this->provider = new EncoderProvider();
        $this->loadEncoder($chatOrAssistant);
    }

    private function loadEncoder($input) {
        if($input instanceof OpenAIAssistant) {
            $this->encoder = $this->provider->getForModel($input->model());
         } elseif($input instanceof OpenAIChat) {
             $this->encoder = $this->provider->getForModel($input->model());
         } elseif($input instanceof Encoder) {
             $this->encoder = $input;
         } elseif(is_string($input)) {
             $assistant = OpenAIAssistant::fromId($input);
             if($assistant) {
                 $this->encoder = $this->provider->getForModel($assistant->model());
             } else {
                 $this->encoder = $this->provider->getForModel($input);
             }
         } else {
             $this->encoder = $this->provider->getForModel("gpt-3.5-turbo");
         }
    }
    public function assistant($assistant) {
        $this->loadEncoder($assistant);
        return $this;
    }
    public function chat($chat) {
        $this->loadEncoder($chat);
        return $this;
    }
    public function model($model) {
        $this->loadEncoder($model);
        return $this;
    }

    public function text($text=null) {
        if($text!==null) {
            $this->text = $text;
        }
        return $this->text;
    }
    public function tokens($tokens=null) {
        if($tokens!==null) {
            $this->tokens = $tokens;
        }
        return $this->tokens;
    }

    public function encode($text=null) {
        if($text===null) {
            $text = $this->text;
        }
        $this->tokens = $this->encoder->encode($text);
        return $this->tokens;
    }

    public function decode($tokens=null) {
        if($tokens===null) {
            $tokens = $this->text;
        }
        $this->text = $this->encoder->decode($tokens);
        return $this->text;
    }

    public function count($tokens=null) {
        if($tokens===null) {
            $tokens = $this->tokens;
        }
        return count($tokens);
    }
}