<?php
/**
 * Vector Search
 * Handles semantic search using cosine similarity
 *
 * @package AskAny
 */

namespace Askany;

if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

class Vector_Search {
	private $database;
	
	public function __construct( $database ) {
		$this->database = $database;
	}
	
	/**
	 * Search for similar content using vector similarity
	 */
	public function search( $query_embedding, $options = array() ) {
		$defaults = array(
			'top_k' => 5,
			'threshold' => 0.7,
			'content_types' => array(),
			'exclude_ids' => array(),
		);
		
		$options = array_merge( $defaults, $options );
		
		// Get all embeddings from database
		$embeddings = $this->get_all_embeddings( $options );
		
		if ( empty( $embeddings ) ) {
			return array();
		}
		
		// Calculate similarity scores
		$results = array();
		foreach ( $embeddings as $embedding ) {
			$stored_vector = json_decode( $embedding['embedding'], true );
			$similarity = $this->calculate_cosine_similarity( $query_embedding, $stored_vector );
			
			if ( $similarity >= $options['threshold'] ) {
				$results[] = array(
					'id' => $embedding['id'],
					'vector_id' => $embedding['vector_id'],
					'content_id' => $embedding['content_id'],
					'content_type' => $embedding['content_type'],
					'source_url' => $embedding['source_url'],
					'title' => $embedding['title'],
					'excerpt' => $embedding['excerpt'],
					'chunk_index' => $embedding['chunk_index'],
					'total_chunks' => $embedding['total_chunks'],
					'similarity' => $similarity,
				);
			}
		}
		
		// Sort by similarity (highest first)
		usort( $results, function( $a, $b ) {
			return $b['similarity'] <=> $a['similarity'];
		} );
		
		// Return top K results
		return array_slice( $results, 0, $options['top_k'] );
	}
	
	/**
	 * Calculate cosine similarity between two vectors
	 */
	public function calculate_cosine_similarity( $vector_a, $vector_b ) {
		if ( count( $vector_a ) !== count( $vector_b ) ) {
			return 0;
		}
		
		$dot_product = 0;
		$magnitude_a = 0;
		$magnitude_b = 0;
		
		for ( $i = 0; $i < count( $vector_a ); $i++ ) {
			$dot_product += $vector_a[ $i ] * $vector_b[ $i ];
			$magnitude_a += $vector_a[ $i ] * $vector_a[ $i ];
			$magnitude_b += $vector_b[ $i ] * $vector_b[ $i ];
		}
		
		$magnitude_a = sqrt( $magnitude_a );
		$magnitude_b = sqrt( $magnitude_b );
		
		if ( $magnitude_a == 0 || $magnitude_b == 0 ) {
			return 0;
		}
		
		return $dot_product / ( $magnitude_a * $magnitude_b );
	}
	
	/**
	 * Get all embeddings from database
	 */
	private function get_all_embeddings( $options ) {
		global $wpdb;
		$table = $wpdb->prefix . 'askany_embeddings';
		
		$where = array( '1=1' );
		
		if ( ! empty( $options['content_types'] ) ) {
			$types = array_map( 'esc_sql', $options['content_types'] );
			$types_str = "'" . implode( "','", $types ) . "'";
			$where[] = "content_type IN ({$types_str})";
		}
		
		if ( ! empty( $options['exclude_ids'] ) ) {
			$ids = array_map( 'intval', $options['exclude_ids'] );
			$ids_str = implode( ',', $ids );
			$where[] = "content_id NOT IN ({$ids_str})";
		}
		
		$where_clause = implode( ' AND ', $where );
		
		$query = "SELECT * FROM {$table} WHERE {$where_clause}";
		
		return $wpdb->get_results( $query, ARRAY_A );
	}
	
	/**
	 * Reassemble chunks from search results
	 */
	public function reassemble_chunks( $search_results ) {
		$reassembled = array();
		$processed_sources = array();
		
		foreach ( $search_results as $result ) {
			// If not chunked, add directly
			if ( $result['total_chunks'] <= 1 ) {
				$reassembled[] = $result;
				continue;
			}
			
			// Extract base hash from vector_id
			$base_hash = $this->get_base_hash( $result['vector_id'] );
			
			// Skip if already processed this source
			if ( in_array( $base_hash, $processed_sources ) ) {
				continue;
			}
			
			// Fetch all chunks for this source
			$all_chunks = $this->get_all_chunks( $base_hash );
			
			if ( empty( $all_chunks ) ) {
				continue;
			}
			
			// Sort by chunk_index
			usort( $all_chunks, function( $a, $b ) {
				return $a['chunk_index'] <=> $b['chunk_index'];
			} );
			
			// Combine into single result
			$combined = array(
				'id' => $result['id'],
				'vector_id' => $base_hash,
				'content_id' => $result['content_id'],
				'content_type' => $result['content_type'],
				'source_url' => $result['source_url'],
				'title' => $result['title'],
				'excerpt' => $result['excerpt'],
				'similarity' => $result['similarity'],
				'is_chunked' => true,
				'chunks' => $all_chunks,
			);
			
			$reassembled[] = $combined;
			$processed_sources[] = $base_hash;
		}
		
		return $reassembled;
	}
	
	/**
	 * Get base hash from vector_id
	 */
	private function get_base_hash( $vector_id ) {
		if ( strpos( $vector_id, '_chunk_' ) !== false ) {
			return substr( $vector_id, 0, strpos( $vector_id, '_chunk_' ) );
		}
		return $vector_id;
	}
	
	/**
	 * Get all chunks for a base hash
	 */
	private function get_all_chunks( $base_hash ) {
		global $wpdb;
		$table = $wpdb->prefix . 'askany_embeddings';
		
		$query = $wpdb->prepare(
			"SELECT * FROM {$table} WHERE vector_id LIKE %s ORDER BY chunk_index ASC",
			$base_hash . '_chunk_%'
		);
		
		return $wpdb->get_results( $query, ARRAY_A );
	}
}
