<?php

namespace Limb_Chatbot\Includes\Services;

use Limb_Chatbot\Includes\Ai_Providers\AI_Providers;
use Limb_Chatbot\Includes\Data_Objects\AI_Model;
use Limb_Chatbot\Includes\Data_Objects\Config;
use Limb_Chatbot\Includes\Data_Objects\Message;
use Limb_Chatbot\Includes\Data_Objects\Vector;
use Limb_Chatbot\Includes\Data_Objects\Vector_Index;
use Limb_Chatbot\Includes\Exceptions\Exception;
use Limb_Chatbot\Includes\Interfaces\Vector_Searcher_Interface;
use Limb_Chatbot\Includes\Repositories\Vector_Repository;
use Limb_Chatbot\Includes\Utilities\Chatbot_Utility;
use Limb_Chatbot\Includes\Utilities\Embedding_Utility;
use Limb_Chatbot\Includes\Utilities\Vector_Db_Utility;

/**
 * Class Vector_Db_Searcher
 *
 * Searches for vectors in a database-backed vector index, generating embeddings as needed.
 *
 * @since 1.0.0
 */
class Vector_Db_Searcher implements Vector_Searcher_Interface {
	const DEFAULT_DIMENSION = 1536;

	/**
	 * Utility object for chatbot-related functionality.
	 *
	 * @var Chatbot_Utility
	 * @since 1.0.0
	 */
	protected Chatbot_Utility $utility;

	/**
	 * Currently loaded vector index object.
	 *
	 * @var Vector_Index|null
	 * @since 1.0.0
	 */
	protected ?Vector_Index $vector_index = null;

	/**
	 * Utility for querying the vector database.
	 *
	 * @var Vector_Db_Utility|null
	 * @since 1.0.0
	 */
	protected ?Vector_Db_Utility $vector_db_utility = null;

	/**
	 * Repository for vector data access.
	 *
	 * @var Vector_Repository
	 * @since 1.0.0
	 */
	protected Vector_Repository $vector_repository;

	/**
	 * Vector_Db_Searcher constructor.
	 *
	 * @param  Chatbot_Utility  $utility  Chatbot utility.
	 * @param  int|null  $vector_index_id  ID of the vector index to use.
	 *
	 * @since 1.0.0
	 */
	public function __construct( Chatbot_Utility $utility, ?int $vector_index_id ) {
		$this->utility      = $utility;
		$this->vector_index = Vector_Index::find( $vector_index_id );
		if ( $this->vector_index ) {
			$this->vector_db_utility = new Vector_Db_Utility( $this->vector_index );
		}
		$this->vector_repository = new Vector_Repository();
	}

	/**
	 * Performs a vector search for the given message.
	 *
	 * @param  Message  $message  Message object containing input text.
	 *
	 * @return array|null Array of search results or null if none.
	 * @throws Exception Throws exception in preview mode on errors.
	 * @since 1.0.0
	 */
	public function search( Message $message ): ?array {
		if ( ! $this->vector_index || ! $this->vector_db_utility ) {
			return [];
		}
		$config_id = $this->utility->chatbot->get_parameter( 'kb_config_id' ) ?? null;
		if ( empty( $config_id ) ) {
			$ai_provider = AI_Providers::instance()->get_ai_provider( $this->utility->get_ai_provider_id() );
			$config_id   = $this->utility->get_config_id();
		} else {
			$config      = Config::find( $config_id );
			$ai_provider = $config->get_related_to_instance();
			$ai_model_id = $this->utility->chatbot->get_parameter( 'kb_ai_model_id' ) ?? null;
		}
		if ( empty( $ai_model_id ) ) {
			$ai_model_id = $ai_provider->get_default_embedding_model();
			if ( $ai_model_id instanceof AI_Model ) {
				$ai_model_id = $ai_model_id->get_id();
			} else {
				return null;
			}
		}
		$existing_vector = $this->vector_repository->get_index_vector( $this->vector_index->get_id(), $ai_model_id );
		if ( $existing_vector ) {
			try {
				$vector = ( new Embedding_Utility( Vector::make( [
					'input'       => trim( $message->extract_text() ?? '' ),
					'config_id'   => $config_id,
					'ai_model_id' => $ai_model_id,
					'dimension'   => $existing_vector->get_dimension() ?? self::DEFAULT_DIMENSION
				] ) ) )->generate();
				if ( $usage = $vector->get_usage() ) {
					( new Usage_Service( $this->utility->get_limits() ) )->update_usage( $usage,
						$existing_vector->get_ai_model() );
				}
				if ( ! $vector->get_values() ) {
					return null;
				}
				$vector->set_model_id( $existing_vector->get_ai_model_id() );

				return $this->vector_db_utility->query( $vector, 3 )->get() ?? null;
			} catch ( Exception $e ) {
				if ( $this->utility->chatbot->is_preview() ) {
					throw $e;
				}

				return null;
			}
		}

		return null;
	}
}