<?php

namespace Limb_Chatbot\Includes\Vector_Dbs\Local\Utilities;

use Limb_Chatbot\Includes\Data_Objects\Vector;
use Limb_Chatbot\Includes\Exceptions\Exception;
use Limb_Chatbot\Includes\Services\Collection;
use Limb_Chatbot\Includes\Services\Helper;
use Limb_Chatbot\Includes\Utilities\Vector_Db_Utility as Global_Vector_Db_Utility;

/**
 * Class Vector_Db_Utility
 *
 * Local vector database utility wrapper.
 *
 * @since 1.0.0
 */
class Vector_Db_Utility {
	/**
	 * Special role allowing universal access.
	 */
	private const ROLE_EVERYONE = 'everyone';

	/**
	 * Distance metric used for vectors.
	 *
	 * @var string
	 * @since 1.0.0
	 */
	public string $metric = 'cosine';

	/**
	 * Vector dimensionality.
	 *
	 * @var int
	 * @since 1.0.0
	 */
	public int $dimension = 1536;

	/**
	 * Identifier for this vector DB.
	 *
	 * @var string
	 * @since 1.0.0
	 */
	public static string $id = 'local';

	/**
	 * Global vector db utility instance.
	 *
	 * @var Global_Vector_Db_Utility
	 * @since 1.0.0
	 */
	public Global_Vector_Db_Utility $global_utility;

	/**
	 * Constructor.
	 *
	 * @param  Global_Vector_Db_Utility  $global_utility  Global vector DB utility instance.
	 *
	 * @since 1.0.0
	 */
	public function __construct( Global_Vector_Db_Utility $global_utility ) {
		$this->global_utility = $global_utility;
	}

	/**
	 * Query vectors from the vector DB.
	 *
	 * @param  Vector  $vector  Vector instance to query.
	 * @param  int  $count  Number of results to return.
	 *
	 * @return Collection Collection of matching vectors.
	 *
	 * @throws Exception
	 * @since 1.0.0
	 * @since 1.0.15 added dimension support
	 */
	public function query( Vector $vector, int $count = 5 ): Collection {
		$query_dimension = $vector->get_dimension();
		$query_roles     = $this->get_global_utility()->get_query_access_roles();

		$vectors = $this->get_vectors(
			$this->get_global_utility()->get_vector_index()->get_id(),
			$vector->get_ai_model_id(),
			$query_dimension,
			$query_roles
		);

		if ( $vectors->is_empty() ) {
			return $vectors;
		}

		$similar_vectors = [];
		foreach ( $vectors->get() as $item ) {
			// Single-pass filtering + scoring for better performance on large sets.
			if ( ! $item->get_values() || ! $this->vector_is_accessible_for_roles( $item, $query_roles ) ) {
				continue;
			}

			$item->score      = $this->cosine_similarity( $item->get_values(), $vector->get_values() );
			$similar_vectors[] = $item;
		}
		if ( empty( $similar_vectors ) ) {
			return new Collection( [] );
		}

		usort( $similar_vectors, function ( $a, $b ) {
			return $b->score <=> $a->score;
		} );
		if ( $count > 0 ) {
			$filtered = array_slice( $similar_vectors, 0, $count );
		}

		return new Collection(  $filtered ?? $similar_vectors );
	}

	/**
	 * Retrieves vectors from the vector index filtered by AI model and dimension.
	 *
	 * @param  int  $vector_index_id  Vector Index ID
	 * @param  int  $ai_model_id  AI Model ID
	 * @param  int|null  $dimension  Expected dimension (optional filter)
	 *
	 * @return Collection Collection of Vector objects.
	 * @since 1.0.0
	 * @since 1.0.15 Added dimension support
	 */
	protected function get_vectors( int $vector_index_id, $ai_model_id, ?int $dimension = null, array $query_roles = [] ): Collection {
		$db_strategy = Vector::get_db_strategy();
		$prefix      = $db_strategy->get_wpdb()->prefix;

		// Build SQL query with dimension filter if provided
		$sql = "SELECT `v`.* FROM `{$prefix}lbaic_vectors` as `v` 
				INNER JOIN `{$prefix}lbaic_vector_indexes_vectors` as `vx` 
				ON `v`.id=`vx`.vector_id 
				WHERE `vx`.vector_index_id = %d
				AND `v`.status = %d 
				AND `v`.values IS NOT NULL 
				AND `v`.ai_model_id = %d";

		$query_params = [ $vector_index_id, Vector::STATUS_ACTIVE, $ai_model_id ];

		// Add dimension filter if specified
		if ( $dimension !== null ) {
			$sql .= " AND `v`.dimension = %d";
			$query_params[] = $dimension;
		}

		$sql .= " ORDER BY `vx`.vector_id DESC, `v`.created_at DESC";

		$results = $db_strategy->get_wpdb()->get_results(
			$db_strategy->get_wpdb()->prepare( $sql, ...$query_params ),
			ARRAY_A
		);

		$objects = [];
		foreach ( $results as $result ) {
			$item      = Vector::make( $result );
			$objects[] = $item;
		}

		return new Collection( $objects ?? [] );
	}

	/**
	 * Check whether vector is accessible for current role set.
	 *
	 * @param  Vector  $vector
	 * @param  array  $query_roles
	 *
	 * @return bool
	 */
	private function vector_is_accessible_for_roles( Vector $vector, array $query_roles ): bool {
		$allowed_roles = $vector->get_access_roles();
		if ( ! is_array( $allowed_roles ) || empty( $allowed_roles ) ) {
			return true;
		}
		$allowed_roles = array_values( array_filter( array_map( static function ( $role ) {
			return is_string( $role ) ? sanitize_key( $role ) : '';
		}, $allowed_roles ) ) );
		if ( empty( $allowed_roles ) ) {
			return true;
		}
		if ( in_array( self::ROLE_EVERYONE, $allowed_roles, true ) ) {
			return true;
		}
		if ( empty( $query_roles ) ) {
			return false;
		}

		return ! empty( array_intersect( $allowed_roles, $query_roles ) );
	}

	/**
	 * Calculates the cosine similarity between two vectors.
	 *
	 * @param  array  $vector_a  First vector.
	 * @param  array  $vector_b  Second vector.
	 *
	 * @return float Cosine similarity value between 0 and 1.
	 */
	private function cosine_similarity( array $vector_a, array $vector_b ): float {
		$dot_product = 0;
		$norm_a      = 0;
		$norm_b      = 0;
		foreach ( $vector_a as $i => $value ) {
			$dot_product += $value * $vector_b[ $i ];
			$norm_a      += $value ** 2;
			$norm_b      += $vector_b[ $i ] ** 2;
		}
		$norm_a = sqrt( $norm_a );
		$norm_b = sqrt( $norm_b );

		return ( $norm_a * $norm_b ) ? ( $dot_product / ( $norm_a * $norm_b ) ) : 0;
	}

	/**
	 * Get the global vector DB utility.
	 *
	 * @return Global_Vector_Db_Utility
	 *
	 * @since 1.0.0
	 */
	public function get_global_utility(): Global_Vector_Db_Utility {
		return $this->global_utility;
	}

	/**
	 * Get metric.
	 *
	 * @return string Metric name.
	 *
	 * @since 1.0.0
	 */
	public function get_metric(): string {
		return $this->metric;
	}

	/**
	 * Set metric.
	 *
	 * @param  string  $metric  Metric name.
	 *
	 * @return void
	 *
	 * @since 1.0.0
	 */
	public function set_metric( string $metric ): void {
		$this->metric = $metric;
	}

	/**
	 * Get dimension.
	 *
	 * @return int Vector dimension.
	 *
	 * @since 1.0.0
	 */
	public function get_dimension(): int {
		return $this->dimension;
	}

	/**
	 * Set dimension.
	 *
	 * @param  int  $dimension  Vector dimension.
	 *
	 * @return void
	 *
	 * @since 1.0.0
	 */
	public function set_dimension( int $dimension ): void {
		$this->dimension = $dimension;
	}
}