<?php
/**
 * Local MySQL Vector Store implementation.
 *
 * @package AISupportAssistant\Embeddings
 */

// phpcs:disable WordPress.DB.DirectDatabaseQuery

namespace AISupportAssistant\Embeddings;

use AISupportAssistant\Database\DatabaseManager;
use AISupportAssistant\Helpers\Logger;

if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

/**
 * Class MySQLVectorStore
 */
class MySQLVectorStore implements VectorStoreInterface {

	/**
	 * Upsert a chunk vector embedding into MySQL chunks table.
	 *
	 * @param int    $document_id  Local document database record ID.
	 * @param int    $chunk_number Sequential index of the chunk.
	 * @param array  $vector       Float array representing embedding values.
	 * @param string $text         Chunk raw text.
	 * @param array  $metadata     Optional key/value annotations.
	 * @return bool True on success, false on failure.
	 */
	public function upsert( int $document_id, int $chunk_number, array $vector, string $text, array $metadata = array() ): bool {
		global $wpdb;
		$table = DatabaseManager::get_chunks_table();

		$page_number = isset( $metadata['page_number'] ) ? intval( $metadata['page_number'] ) : null;
		$token_count = isset( $metadata['token_count'] ) ? intval( $metadata['token_count'] ) : 0;

		$data = array(
			'document_id'     => $document_id,
			'chunk_number'    => $chunk_number,
			'page_number'     => $page_number,
			'chunk_text'      => $text,
			'embedding'       => wp_json_encode( $vector ),
			'embedding_model' => $metadata['embedding_model'] ?? EmbeddingService::$last_model,
			'token_count'     => $token_count,
			'created_at'      => current_time( 'mysql' ),
		);

		$result = $wpdb->query(
			$wpdb->prepare(
				'INSERT INTO %i (document_id, chunk_number, page_number, chunk_text, embedding, embedding_model, token_count, created_at)
				 VALUES (%d, %d, %s, %s, %s, %s, %d, %s)
				 ON DUPLICATE KEY UPDATE page_number = VALUES(page_number), chunk_text = VALUES(chunk_text), embedding = VALUES(embedding), embedding_model = VALUES(embedding_model), token_count = VALUES(token_count), created_at = VALUES(created_at)',
				$table, $document_id, $chunk_number, $page_number, $text, wp_json_encode( $vector ), $data['embedding_model'], $token_count, $data['created_at']
			)
		);

		return false !== $result;
	}

	/**
	 * Compute cosine similarity between two vectors.
	 *
	 * @param array $vec1 Vector A.
	 * @param array $vec2 Vector B.
	 * @return float Similarity score.
	 */
	public static function cosine_similarity( array $vec1, array $vec2 ): float {
		$dot_product = 0.0;
		$norm_a      = 0.0;
		$norm_b      = 0.0;

		$count = count( $vec1 );
		if ( $count !== count( $vec2 ) ) {
			return 0.0;
		}

		for ( $i = 0; $i < $count; $i++ ) {
			$a            = (float) ( $vec1[ $i ] ?? 0.0 );
			$b            = (float) ( $vec2[ $i ] ?? 0.0 );
			$dot_product += $a * $b;
			$norm_a      += $a * $a;
			$norm_b      += $b * $b;
		}

		if ( 0.0 === $norm_a || 0.0 === $norm_b ) {
			return 0.0;
		}

		return $dot_product / ( sqrt( $norm_a ) * sqrt( $norm_b ) );
	}

	/**
	 * Search vector database for top-K matching records.
	 *
	 * @param array $vector Query embedding float vector.
	 * @param int   $top_k  Maximum matches limit.
	 * @return array Array of matching chunks with similarity scores.
	 */
	public function search( array $vector, int $top_k ): array {
		global $wpdb;
		$chunks_table = DatabaseManager::get_chunks_table();
		$docs_table   = DatabaseManager::get_documents_table();

		$top_k = max( 1, min( 50, $top_k ) );
		$model = EmbeddingService::$last_model;
		if ( '' === $model ) {
			throw new \RuntimeException( 'Unable to identify the query embedding model.' );
		}
		$results = array();
		$threshold = max( -1.0, min( 1.0, (float) \AISupportAssistant\Admin\Settings::get_instance()->get( 'similarity_threshold', 0.5 ) ) );
		$last_id = 0;
		$scanned = 0;
		do {
			// Keyset pagination bounds memory; retain only the best matches.
			$rows = $wpdb->get_results(
				$wpdb->prepare(
					'SELECT c.id, c.document_id, c.chunk_number, c.page_number, c.chunk_text, c.embedding, d.filename
					FROM %i c JOIN %i d ON c.document_id = d.id
					WHERE c.id > %d AND c.embedding_model = %s
					AND (d.status IN (%s, %s, %s) OR (d.status = %s AND d.error_message LIKE %s))
					ORDER BY c.id ASC LIMIT 100',
					$chunks_table,
					$docs_table,
					$last_id,
					$model,
					'completed',
					'pending',
					'indexing',
					'failed',
					'Indexing failed; existing index remains available:%'
				),
				ARRAY_A
			);
			if ( ! empty( $wpdb->last_error ) ) {
				throw new \RuntimeException( esc_html( 'Vector store search failed: ' . $wpdb->last_error ) );
			}
			foreach ( $rows as $row ) {
				$last_id = (int) $row['id'];
				if ( ++$scanned > 10000 ) {
					throw new \RuntimeException( 'The local vector index exceeds the supported limit of 10,000 matching chunks.' );
				}
				$chunk_vector = json_decode( $row['embedding'], true );
				if ( ! is_array( $chunk_vector ) || count( $chunk_vector ) !== count( $vector ) ) {
					continue;
				}
				$score = self::cosine_similarity( $vector, $chunk_vector );
				if ( $score < $threshold ) { continue; }
				$results[] = array(
					'document_id'  => (int) $row['document_id'],
					'chunk_number' => (int) $row['chunk_number'],
					'page_number'  => null !== $row['page_number'] ? (int) $row['page_number'] : null,
					'chunk_text'   => $row['chunk_text'],
					'filename'     => $row['filename'],
					'score'        => $score,
				);
				usort(
					$results,
					static function ( $a, $b ) {
						return $b['score'] <=> $a['score'];
					}
				);
				$results = array_slice( $results, 0, $top_k );
			}
		} while ( count( $rows ) === 100 );
		return $results;
	}

	/**
	 * Delete all vectors associated with a specific document.
	 *
	 * @param int $document_id Document database ID.
	 * @return bool True on success, false on failure.
	 */
	public function delete_document( int $document_id ): bool {
		global $wpdb;
		$table  = DatabaseManager::get_chunks_table();
		$result = $wpdb->delete( $table, array( 'document_id' => $document_id ) );
		return false !== $result;
	}

	/**
	 * Clear the entire vector index.
	 *
	 * @return bool True on success, false on failure.
	 */
	public function clear(): bool {
		global $wpdb;
		$table  = DatabaseManager::get_chunks_table();
		$result = $wpdb->query( $wpdb->prepare( 'TRUNCATE TABLE %i', $table ) );
		return false !== $result;
	}
}