<?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;
use Limb_Chatbot\Includes\Services\Knowledge\Query_Intelligence_Service;
use Limb_Chatbot\Includes\Data_Objects\Dataset_Entry;

/**
 * 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 {
				$search_text = trim( $message->extract_text() ?? '' );

				$vector = ( new Embedding_Utility( Vector::make( [
					'input'       => $search_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() );

				// Dense retrieval: initial candidate set from vector DB.
				$limit       = $this->utility->get_kb_vector_similarity_count() ?? self::DEFAULT_DIMENSION;
				$raw_results = $this->vector_db_utility->query( $vector, $limit )->get() ?? null;

				if ( empty( $raw_results ) || ! is_array( $raw_results ) ) {
					return null;
				}

				// Hybrid retrieval + heuristic reranking (dense + sparse keyword scoring).
				return $this->rerank_hybrid( $raw_results, $search_text );
			} catch ( Exception $e ) {
				if ( $this->utility->chatbot->is_preview() ) {
					throw $e;
				}

				return null;
			}
		}

		return null;
	}

	/**
	 * Combine dense vector similarity with sparse keyword-based scoring and rerank.
	 *
	 * @param  Vector[]  $vectors      Raw vector search results (with similarity score 0–1).
	 * @param  string    $search_text  Normalized search text used for embedding.
	 *
	 * @return Vector[] Reranked vector list with updated hybrid scores.
	 *
	 * @since 1.0.0
	 */
	private function rerank_hybrid( array $vectors, string $search_text ): array {
		if ( empty( $vectors ) ) {
			return [];
		}

		// Extract keyword candidates directly from the user query (sparse signals).
		$query_intel = new Query_Intelligence_Service( $this->utility );
		$keywords    = $query_intel->extract_keywords( $search_text );

		// If no meaningful keywords, fall back to pure vector ranking.
		if ( empty( $keywords ) ) {
			return $vectors;
		}

		$keywordResults = [];
		$vectorResults  = [];

		// Build sparse and dense score maps keyed by dataset_entry_id.
		foreach ( $vectors as $vector ) {
			if ( ! $vector instanceof Vector ) {
				continue;
			}

			$entry_id = $vector->get_dataset_entry_id();
			if ( empty( $entry_id ) ) {
				continue;
			}

			// Dense score from vector similarity (already normalized 0–1).
			$vector_score = isset( $vector->score ) ? (float) $vector->score : 0.0;
			$vectorResults[ $entry_id ] = $vector_score;

			// Sparse keyword score based on overlap between query keywords and chunk text.
			$entry = $vector->dataset_entry();
			if ( ! $entry instanceof Dataset_Entry ) {
				continue;
			}

			$chunk_text = mb_strtolower( $entry->stringify_for_inference() );
			if ( $chunk_text === '' ) {
				continue;
			}

			$matches = 0;
			foreach ( $keywords as $keyword ) {
				if ( mb_stripos( $chunk_text, $keyword ) !== false ) {
					$matches++;
				}
			}

			if ( $matches > 0 && count( $keywords ) > 0 ) {
				// Normalize keyword score to 0–1 by fraction of matched keywords.
				$keywordResults[ $entry_id ] = $matches / count( $keywords );
			}
		}

		// If no sparse signal at all, keep original dense ordering.
		if ( empty( $keywordResults ) ) {
			return $vectors;
		}

		// Merge candidate scores and compute final hybrid score.
		$candidates = [];

		foreach ( $keywordResults as $id => $score ) {
			$candidates[ $id ]['keyword'] = $score;
		}

		foreach ( $vectorResults as $id => $score ) {
			$candidates[ $id ]['vector'] = $score;
		}

		foreach ( $candidates as $id => &$candidate ) {
			$candidate['keyword'] = $candidate['keyword'] ?? 0.0;
			$candidate['vector']  = $candidate['vector'] ?? 0.0;

			// STEP 3 — Heuristic Reranking (FAST)
			$candidate['final'] = ( $candidate['vector'] * 0.6 ) + ( $candidate['keyword'] * 0.4 );
		}
		unset( $candidate );

		// Apply final scores back onto vectors and sort deterministically.
		usort( $vectors, function ( $a, $b ) use ( $candidates ) {
			if ( ! $a instanceof Vector || ! $b instanceof Vector ) {
				return 0;
			}

			$a_id = $a->get_dataset_entry_id();
			$b_id = $b->get_dataset_entry_id();

			$a_score = $a_id && isset( $candidates[ $a_id ]['final'] ) ? $candidates[ $a_id ]['final'] : 0.0;
			$b_score = $b_id && isset( $candidates[ $b_id ]['final'] ) ? $candidates[ $b_id ]['final'] : 0.0;

			if ( $a_score === $b_score ) {
				return 0;
			}

			return ( $a_score < $b_score ) ? 1 : -1;
		} );

		foreach ( $vectors as $vector ) {
			if ( ! $vector instanceof Vector ) {
				continue;
			}
			$entry_id = $vector->get_dataset_entry_id();
			if ( $entry_id && isset( $candidates[ $entry_id ]['final'] ) ) {
				$vector->score = (float) $candidates[ $entry_id ]['final'];
			}
		}

		return $vectors;
	}
}