<?php
if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

/**
 * Embedding storage and similarity search over a dedicated custom table.
 *
 * phpcs:disable WordPress.DB.DirectDatabaseQuery -- the whole class exists to
 * manage the plugin's own table, which core APIs and object caching do not
 * cover; every dynamic value goes through $wpdb->prepare() placeholders.
 */
class Arcai_Vector_Store {

	public function replace_post_chunks( $post_id, $post_date, array $chunks, array $embeddings ) {
		global $wpdb;
		$this->delete_post_chunks( $post_id );

		foreach ( $chunks as $i => $chunk ) {
			$wpdb->insert(
				$wpdb->prefix . 'arcai_chunks',
				array(
					'post_id'     => $post_id,
					'chunk_index' => $i,
					'content'     => $chunk['content'],
					'embedding'   => isset( $embeddings[ $i ] ) ? pack( 'f*', ...$embeddings[ $i ] ) : null,
					'token_count' => $chunk['token_estimate'],
					'post_date'   => $post_date,
				),
				array( '%d', '%d', '%s', '%s', '%d', '%s' )
			);
		}
	}

	public function delete_post_chunks( $post_id ) {
		global $wpdb;
		$wpdb->delete( $wpdb->prefix . 'arcai_chunks', array( 'post_id' => $post_id ), array( '%d' ) );
		update_option( 'arcai_index_generation', (string) microtime( true ), false );
	}

	public function indexed_post_count() {
		global $wpdb;
		return (int) $wpdb->get_var( "SELECT COUNT(DISTINCT post_id) FROM {$wpdb->prefix}arcai_chunks" );
	}

	/**
	 * Cosine-similarity search over stored chunks.
	 *
	 * OpenAI embeddings are unit-normalized, so the dot product equals
	 * cosine similarity.
	 *
	 * @param float[]     $query_embedding
	 * @param int         $limit
	 * @param string|null $date_from Y-m-d or null.
	 * @param string|null $date_to   Y-m-d or null.
	 * @return array[] Each: [ 'post_id', 'content', 'score' ]
	 */
	public function search( array $query_embedding, $limit = 8, $date_from = null, $date_to = null ) {
		global $wpdb;

		$where  = array( 'embedding IS NOT NULL' );
		$params = array();
		if ( $date_from ) {
			$where[]  = 'post_date >= %s';
			$params[] = $date_from . ' 00:00:00';
		}
		if ( $date_to ) {
			$where[]  = 'post_date <= %s';
			$params[] = $date_to . ' 23:59:59';
		}
		$where_sql = implode( ' AND ', $where );

		$top    = array();
		$offset = 0;
		$batch  = 2000;

		do {
			// $where_sql is assembled exclusively from the fixed clause strings
			// above; all user-influenced values are bound via placeholders.
			// phpcs:disable WordPress.DB.PreparedSQL, WordPress.DB.PreparedSQLPlaceholders.ReplacementsWrongNumber, PluginCheck.Security.DirectDB.UnescapedDBParameter
			$rows = $wpdb->get_results(
				$wpdb->prepare(
					"SELECT id, embedding FROM {$wpdb->prefix}arcai_chunks WHERE $where_sql LIMIT %d OFFSET %d",
					array_merge( $params, array( $batch, $offset ) )
				)
			);
			// phpcs:enable WordPress.DB.PreparedSQL, WordPress.DB.PreparedSQLPlaceholders.ReplacementsWrongNumber, PluginCheck.Security.DirectDB.UnescapedDBParameter

			foreach ( $rows as $row ) {
				$vector = unpack( 'f*', $row->embedding );
				$dot    = 0.0;
				$j      = 1;
				foreach ( $query_embedding as $q ) {
					$dot += $q * $vector[ $j ];
					$j++;
				}
				$top[ $row->id ] = $dot;
			}

			if ( count( $top ) > $limit * 4 ) {
				arsort( $top );
				$top = array_slice( $top, 0, $limit * 4, true );
			}

			$offset += $batch;
		} while ( count( $rows ) === $batch );

		arsort( $top );
		$top = array_slice( $top, 0, $limit, true );

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

		$ids          = array_map( 'intval', array_keys( $top ) );
		$placeholders = implode( ',', array_fill( 0, count( $ids ), '%d' ) );
		// $placeholders is a generated list of %d markers, one per ID.
		// phpcs:disable WordPress.DB.PreparedSQL, WordPress.DB.PreparedSQLPlaceholders
		$rows = $wpdb->get_results(
			$wpdb->prepare( "SELECT id, post_id, content FROM {$wpdb->prefix}arcai_chunks WHERE id IN ($placeholders)", $ids )
		);
		// phpcs:enable WordPress.DB.PreparedSQL, WordPress.DB.PreparedSQLPlaceholders

		$results = array();
		foreach ( $rows as $row ) {
			$results[] = array(
				'post_id' => (int) $row->post_id,
				'content' => $row->content,
				'score'   => $top[ $row->id ],
			);
		}
		usort(
			$results,
			function ( $a, $b ) {
				return $b['score'] <=> $a['score'];
			}
		);

		return $results;
	}
}
