<?php
/**
 * Pinecone vector store.
 *
 * @package AI House
 */

defined( 'ABSPATH' ) || exit;

/**
 * Class MYBOT_Vector_Pinecone
 */
class MYBOT_Vector_Pinecone implements MYBOT_Vector_Store {

	/**
	 * Index URL, for example https://my-index-abc123.svc.environment.pinecone.io
	 *
	 * @var string
	 */
	protected $url;

	/**
	 * API key.
	 *
	 * @var string
	 */
	protected $key;

	/**
	 * Namespace used to separate data between the sites.
	 *
	 * @var string
	 */
	protected $namespace;

	/**
	 * Constructor.
	 *
	 * @param string $url       Index URL.
	 * @param string $key       API key.
	 * @param string $namespace Namespace.
	 */
	public function __construct( $url, $key, $namespace = 'mybot' ) {
		$this->url       = rtrim( (string) $url, '/' );
		$this->key       = (string) $key;
		$this->namespace = (string) $namespace;
	}

	/**
	 * Store ID.
	 *
	 * @return string
	 */
	public function get_id() {
		return 'pinecone';
	}

	/**
	 * Display name.
	 *
	 * @return string
	 */
	public function get_label() {
		return 'Pinecone';
	}

	/**
	 * Whether the store is configured.
	 *
	 * @return bool
	 */
	public function is_configured() {
		return '' !== $this->url && '' !== $this->key;
	}

	/**
	 * Authentication headers.
	 *
	 * @return array
	 */
	protected function headers() {
		return array( 'Api-Key' => $this->key );
	}

	/**
	 * Writes vectors.
	 *
	 * @param array $records List of records.
	 * @return true|WP_Error
	 */
	public function upsert( $records ) {
		if ( ! $this->is_configured() ) {
			return new WP_Error( 'mybot_vector_unconfigured', __( 'Pinecone is not configured.', 'aihouse' ) );
		}

		if ( empty( $records ) ) {
			return true;
		}

		$vectors = array();

		foreach ( $records as $record ) {
			$metadata = isset( $record['metadata'] ) && is_array( $record['metadata'] ) ? $record['metadata'] : array();

			// Pinecone only accepts flat metadata with simple types.
			$metadata['content'] = isset( $record['content'] ) ? (string) $record['content'] : '';

			$vectors[] = array(
				'id'       => (string) $record['id'],
				'values'   => isset( $record['values'] ) ? $record['values'] : array(),
				'metadata' => $metadata,
			);
		}

		$response = MYBOT_HTTP::post_json(
			$this->url . '/vectors/upsert',
			array(
				'vectors'   => $vectors,
				'namespace' => $this->namespace,
			),
			$this->headers(),
			20,
			1
		);

		return is_wp_error( $response ) ? $response : true;
	}

	/**
	 * Queries vectors.
	 *
	 * @param array $vector    Query vector.
	 * @param int   $top_k     Number of results.
	 * @param float $min_score Score threshold.
	 * @return array|WP_Error
	 */
	public function query( $vector, $top_k = 4, $min_score = 0.0 ) {
		if ( ! $this->is_configured() ) {
			return new WP_Error( 'mybot_vector_unconfigured', __( 'Pinecone is not configured.', 'aihouse' ) );
		}

		$response = MYBOT_HTTP::post_json(
			$this->url . '/query',
			array(
				'vector'          => $vector,
				'topK'            => max( 1, (int) $top_k ),
				'includeMetadata' => true,
				'namespace'       => $this->namespace,
			),
			$this->headers(),
			15
		);

		if ( is_wp_error( $response ) ) {
			return $response;
		}

		$results = array();

		foreach ( (array) ( isset( $response['matches'] ) ? $response['matches'] : array() ) as $match ) {
			$score = isset( $match['score'] ) ? (float) $match['score'] : 0.0;

			if ( $score < (float) $min_score ) {
				continue;
			}

			$metadata = isset( $match['metadata'] ) && is_array( $match['metadata'] ) ? $match['metadata'] : array();
			$content  = isset( $metadata['content'] ) ? (string) $metadata['content'] : '';
			unset( $metadata['content'] );

			$results[] = array(
				'id'       => isset( $match['id'] ) ? (string) $match['id'] : '',
				'score'    => $score,
				'content'  => $content,
				'metadata' => $metadata,
			);
		}

		return $results;
	}

	/**
	 * Deletes by a list of IDs.
	 *
	 * @param array $ids List of IDs.
	 * @return true|WP_Error
	 */
	public function delete( $ids ) {
		$ids = array_values( array_filter( array_map( 'strval', (array) $ids ), 'strlen' ) );

		return $this->delete_request( array( 'ids' => $ids ) );
	}

	/**
	 * Deletes every vector that belongs to a post.
	 *
	 * @param int $source_id Post ID.
	 * @return true|WP_Error
	 */
	public function delete_by_source( $source_id ) {
		return $this->delete_request( array( 'filter' => array( 'source_id' => (int) absint( $source_id ) ) ) );
	}

	/**
	 * Sends a delete request.
	 *
	 * @param array $body Request body.
	 * @return true|WP_Error
	 */
	protected function delete_request( $body ) {
		if ( ! $this->is_configured() ) {
			return new WP_Error( 'mybot_vector_unconfigured', __( 'Pinecone is not configured.', 'aihouse' ) );
		}

		$body['namespace'] = $this->namespace;

		$response = MYBOT_HTTP::post_json( $this->url . '/vectors/delete', $body, $this->headers(), 15 );

		return is_wp_error( $response ) ? $response : true;
	}
}