<?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;

		$existing_id = $wpdb->get_var(
			$wpdb->prepare(
				"SELECT id FROM %i WHERE document_id = %d AND chunk_number = %d",
				$table,
				$document_id,
				$chunk_number
			)
		);

		$data = array(
			'document_id'  => $document_id,
			'chunk_number' => $chunk_number,
			'page_number'  => $page_number,
			'chunk_text'   => $text,
			'embedding'    => wp_json_encode( $vector ),
			'token_count'  => $token_count,
			'created_at'   => current_time( 'mysql' ),
		);

		if ( $existing_id ) {
			$result = $wpdb->update( $table, $data, array( 'id' => $existing_id ) );
		} else {
			$result = $wpdb->insert( $table, $data );
		}

		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 ) ) {
			// Fallback if dimensions differ, trim or return zero.
			$count = min( count( $vec1 ), count( $vec2 ) );
		}

		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();

		// Fetch all chunks from completed documents.
		$rows = $wpdb->get_results(
			$wpdb->prepare(
				"SELECT 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 d.status = %s",
				$chunks_table,
				$docs_table,
				'completed'
			),
			ARRAY_A
		);

		if ( empty( $rows ) ) {
			return array();
		}

		$results = array();
		foreach ( $rows as $row ) {
			$chunk_vector = json_decode( $row['embedding'], true );
			if ( ! is_array( $chunk_vector ) ) {
				continue;
			}

			// Calculate similarity score.
			$score = self::cosine_similarity( $vector, $chunk_vector );

			$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,
			);
		}

		// Sort by score descending.
		usort(
			$results,
			function( $a, $b ) {
				return $b['score'] <=> $a['score'];
			}
		);

		// Return top K.
		return array_slice( $results, 0, $top_k );
	}

	/**
	 * 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;
	}
}