<?php

namespace wpsolr\core\classes\engines\opensearch_vector_php;


use OpenSearch\Namespaces\MlNamespace;
use wpsolr\core\classes\engines\opensearch_php\WPSOLR_OpenSearchClient;
use wpsolr\core\classes\hosting_api\WPSOLR_Hosting_Api_Abstract;

trait WPSOLR_OpenSearch_Vector_Client {
	use WPSOLR_OpenSearchClient {
		_before_client_build as _parent_before_client_build;
	}

	protected function _get_class_for_file_settings(): string {
		// Use same file settings as Opensearch
		return get_parent_class( $this );
	}

	protected string $_FIELD_NAME_TEXT_CONCATENATED_FIELDS_EMBEDDING = 'text_embedding';

	protected function _customize_file_json( array &$index_parameters, array $configuration ): array {
		$model_dimensions                                                                                                 = $this->_get_model_dimensions();
		$model_id                                                                                                         = $this->_register_model();
		$index_parameters[ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_WEAVIATE_HUGGINGFACE_CONFIG_MODEL_QUERY ] = $model_id;

		$custom_configuration = [
			'body' => [
				'settings' => [
					'index' => [
						'knn'              => true,
						"default_pipeline" => $this->get_index_label(),
					]
				],
				'mappings' => [
					'properties' => [
						$this->_FIELD_NAME_TEXT_CONCATENATED_FIELDS_EMBEDDING => [
							'type'      => 'knn_vector',
							'dimension' => $model_dimensions,
							'method'    => [
								'name'       => 'hnsw',
								'engine'     => 'lucene', # FAISS crashes
								'space_type' => 'cosinesimil',
								'parameters' => [
									'm'               => 16,
									'ef_construction' => 100
								]
							]
						]
					]
				]
			]
		];


		return array_merge_recursive( $configuration, $custom_configuration );
	}

	protected function _register_model() {
		switch ( $model_type = $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_LLM_PROVIDER ] ?? '' ) {
			case WPSOLR_Opensearch_Vector_Constants::MODEL_LOCAL_PRETRAINED_TRANSFORMER:
				return $this->_register_model_pretrained_transformer();

			case WPSOLR_Opensearch_Vector_Constants::MODEL_API_OPENAI:
				return $this->_register_model_openai();

			case WPSOLR_Opensearch_Vector_Constants::MODEL_API_COHERE:
				return $this->_register_model_cohere();

			default:
				throw new \Exception( esc_html( "Model type {$model_type} is not implemented yet." ) );
		}
	}

	protected function _register_model_openai(): string {
		if ( empty( $model_name = $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_LANGUAGE_CODE ] ?? '' ) ) {
			throw new \Exception( 'OpenAI model name is required.' );
		}
		if ( empty( $openai_api_key = trim( $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_API_KEY ] ?? '' ) ) ) {
			throw new \Exception( 'OpenAI API key is required.' );
		}

		return $this->_get_or_create_model(
			[
				'name'          => $this->get_index_label(),
				// One model per index, as api key can be different, hence connector id can be different
				'function_name' => 'remote',
				'description'   => "WPSOLR - OpenAI text embedding model for index {$this->get_index_label()}",
			],
			[
				'name'        => $this->get_index_label(),
				'description' => 'Connector to OpenAI embeddings API for index',
				'version'     => '1.0',
				'protocol'    => 'http',
				'parameters'  => [
					'model' => $model_name
				],
				'credential'  => [
					'openai_key' => $openai_api_key
				],
				'actions'     => [
					[
						'action_type'           => 'predict',
						'method'                => 'POST',
						'url'                   => 'https://api.openai.com/v1/embeddings',
						'headers'               => [
							'Authorization' => 'Bearer ${credential.openai_key}'
						],
						'request_body'          => '{ "input": ${parameters.input}, "model": "${parameters.model}" }',
						"pre_process_function"  => "connector.pre_process.openai.embedding",
						"post_process_function" => "connector.post_process.openai.embedding"
					]
				]
			]
		);

	}

	protected function _register_model_cohere(): string {
		if ( empty( $model_name = $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_LANGUAGE_CODE ] ?? '' ) ) {
			throw new \Exception( 'Cohere model name is required.' );
		}
		if ( empty( $cohere_api_key = trim( $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_API_KEY ] ?? '' ) ) ) {
			throw new \Exception( 'Cohere API key is required.' );
		}

		return $this->_get_or_create_model(
			[
				'name'          => $this->get_index_label(),
				// One model per index, as api key can be different, hence connector id can be different
				'function_name' => 'remote',
				'description'   => "WPSOLR - OpenAI text embedding model for index {$this->get_index_label()}",
			],
			[
				'name'        => $this->get_index_label(),
				'description' => 'Connector to Cohere embeddings API for index',
				'version'     => '1',
				'protocol'    => 'http',
				'parameters'  => [
					'model'      => $model_name,
					'input_type' => 'search_document',
					'truncate'   => 'END'
				],
				'credential'  => [
					'cohere_key' => $cohere_api_key
				],
				'actions'     => [
					[
						'action_type'           => 'predict',
						'method'                => 'POST',
						'url'                   => 'https://api.cohere.ai/v1/embed',
						'headers'               => [
							'Authorization'  => 'Bearer ${credential.cohere_key}',
							'Request-Source' => 'unspecified:opensearch'
						],
						'request_body'          => '{ "texts": ${parameters.texts}, "truncate": "${parameters.truncate}", 
						"model": "${parameters.model}", "input_type": "${parameters.input_type}" }',
						'pre_process_function'  => 'connector.pre_process.cohere.embedding',
						'post_process_function' => 'connector.post_process.cohere.embedding'
					]
				]
			]
		);

	}

	protected function _register_model_pretrained_transformer(): string {
		if ( empty( $model_name = $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_LANGUAGE_CODE ] ?? '' ) ) {
			throw new \Exception( 'Pretrained model name is required.' );
		}
		if ( empty( $model_version = $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_WEAVIATE_HUGGINGFACE_CONFIG_MODEL ] ?? '' ) ) {
			throw new \Exception( 'Model version is required.' );
		}

		return $this->_get_or_create_model( [
			'name'         => $model_name,
			'version'      => $model_version,
			'model_format' => 'TORCH_SCRIPT',
		], [] );
	}

	protected function _get_or_create_search_pipeline( array $pipeline_def ) {
		$this->search_engine_client->searchPipeline()->put( $pipeline_def );
	}

	protected function _delete_search_pipeline( string $pipeline_id ) {
		$this->search_engine_client->searchPipeline()->delete( [ 'id' => $pipeline_id ] );
	}

	protected function _get_or_create_ingest_pipeline( array $pipeline_def ) {
		$this->search_engine_client->ingest()->putPipeline( $pipeline_def );
	}

	protected function _delete_ingest_pipeline( string $pipeline_id ) {
		$this->search_engine_client->ingest()->deletePipeline( [ 'id' => $pipeline_id ] );
	}

	protected function _delete_connectors_by_index_name() {
		// Connectors with same name can accidentally exist. Remove them all by name.
		$connectors = $this->search_engine_client->ml()->getConnectors( [
			'body' => [
				'query' => [
					'bool' => [
						'must' => [
							[
								'term' => [ 'name.keyword' => $this->get_index_label() ],
							]
						]
					]
				]
			]
		] );
		foreach ( $connectors['hits']['hits'] ?? [] as $connector ) {
			try {
				$this->search_engine_client->ml()->deleteConnector( [ 'id' => $connector['_id'] ] );
			} catch ( \Exception $e ) {
				// Continue
				$t = 1;
			}
		}

	}

	protected function _delete_models_by_index_name() {
		// Models with same name can accidentally exist. Remove them all by name.
		$models = $this->_get_models_by_name( $this->get_index_label() );
		foreach ( $models as $model ) {
			try {
				$this->_delete_model( $model['_id'] );
			} catch ( \Exception $e ) {
				// Continue
				$t = 1;
			}
		}

	}

	protected function _get_or_create_model_group( MlNamespace $ml_namespace ): string {
		// Check if the model group exists
		$model_group_name = 'wpsolr';
		$groups           = $ml_namespace->getModelGroups( [
			'body' => [ 'query' => [ 'match' => [ 'name' => $model_group_name ] ] ]
		] )['hits']['hits'] ?? [];

		if ( ! empty( $groups ) ) {
			// Group exists, get its ID
			return $groups[0]['_id'];

		} else {
			$model_group = $ml_namespace->registerModelGroup( [
				'body' => json_encode( [
					'name'        => $model_group_name,
					'description' => 'WPSOLR embedding models'
				] )
			] );

			return $model_group['model_group_id'];
		}
	}

	protected function _poll_task_api( MlNamespace $ml_namespace, string $task_id, string $model_name ): string {
		// Poll the task API
		do {
			sleep( 10 ); // wait
			$task_response = $ml_namespace->getTask( [
				'id' => $task_id
			] );
			$state         = $task_response['state'] ?? null;
		} while ( $state !== 'COMPLETED' && $state !== 'FAILED' );

		if ( $state === 'FAILED' ) {
			$error = $task_response['error'] ?? '';
			throw new \Exception( esc_html( "Error while registering the model {$model_name}: {$error}" ) );
		}

		if ( empty( $task_response['model_id'] ) ) {
			throw new \Exception( esc_html( "Error while registering the model {$model_name}: model id returned is empty" ) );
		}

		return $task_response['model_id'];
	}

	protected function _get_or_create_model( array $model_def, array $connector_def ): string {
		$ml_namespace = $this->search_engine_client->ml();

		// Get or create model group
		$model_group_id = $this->_get_or_create_model_group( $ml_namespace );

		// Ensure model is registered
		$models = $this->_get_models_by_name( $model_def['name'] );

		if ( ! empty( $models ) ) {
			$model_id = $models[0]['_id'];
		} else {

			/**
			 * Register the connector eventually
			 */
			if ( ! empty( $connector_def ) ) {
				$connector_id              = $this->_get_or_create_connector( $connector_def );
				$model_def['connector_id'] = $connector_id;
			}

			/**
			 * Register the model
			 */
			$model_def['model_group_id'] = $model_group_id;
			$register_model_response     = $ml_namespace->registerModel( [ 'body' => $model_def ] );
			$model_id                    = $this->_poll_task_api( $ml_namespace, $register_model_response['task_id'], $model_def['name'] );

			/**
			 * Deploy the registered the model
			 */
			$deploy_model_response = $ml_namespace->deployModel( [
				'id' => $model_id
			] );
			$model_id              = $this->_poll_task_api( $ml_namespace, $deploy_model_response['task_id'], $model_def['name'] );
		}

		// Get or create ingestion pipeline
		$this->_get_or_create_ingest_pipeline( [
			'id'   => $this->get_index_label(),
			'body' => [
				'description' => 'Generate embeddings for k-nn model',
				'processors'  => [
					[
						'text_embedding' => [
							'model_id'  => $model_id,
							'field_map' => [ 'content' => $this->_FIELD_NAME_TEXT_CONCATENATED_FIELDS_EMBEDDING, ]
						]
					]
				]
			]
		] );

		// Get or create RRF pipeline
		$this->_get_or_create_search_pipeline( [
			'id'   => $this->get_index_label(),
			'body' => [
				'description'              => 'WPSOLR - Hybrid search with RRF ranking',
				'phase_results_processors' => [
					[
						'score-ranker-processor' => [
							'combination' => [
								'technique'     => 'rrf',
								'rank_constant' => 60,
								'parameters'    => [
									'weights' => [ 0.7, 0.3 ]  // weights for lexical vs semantic queries
								]
							]
						]
					]
				]
			]
		] );

		return $model_id;
	}

	protected function _get_or_create_connector( array $connector_def ): string {
		return $this->search_engine_client->ml()->createConnector( [ 'body' => $connector_def ] )['connector_id'];
	}

	protected function _get_model_dimensions(): string {
		if ( empty( $model_dimensions = $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_CATALOG_BRANCH ] ?? '' ) ) {
			throw new \Exception( 'Model dimensions are required.' );
		}
		if ( ! ctype_digit( $model_dimensions ) ) {
			throw new \Exception( esc_html( "Pretrained model dimensions '{$model_dimensions}' should be an integer." ) );
		}

		return $model_dimensions;
	}

	/**
	 * @return mixed
	 */
	protected function _get_model_id(): string {
		if ( empty( $model_id = $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_FIELDS_INDEX_WEAVIATE_HUGGINGFACE_CONFIG_MODEL_QUERY ] ?? '' ) ) {
			throw new \Exception( 'No model id found.' );
		}

		return $model_id;
	}

	protected function _clean_before_delete_index() {
		// Delete the index search pipeline
		try {
			$this->_delete_search_pipeline( $this->get_index_label() );
		} catch ( \Exception $e ) {
			// Not found
		}

		// Delete the index ingest pipeline
		try {
			$this->_delete_ingest_pipeline( $this->get_index_label() );
		} catch ( \Exception $e ) {
			// Not found
		}

		$model_id = $this->_get_model_id();
		try {
			$model = $this->search_engine_client->ml()->getModel( [ 'id' => $model_id ] );
		} catch ( \Exception $e ) {
			// Not found
		}

		if ( $model ) {
			// Delete the model if not used anymore in any other index pipeline
			$indices       = $this->search_engine_client->indices()->get( [ 'index' => '*' ] );
			$found_indices = false;
			foreach ( array_keys( $indices ) as $index_name ) {
				if ( $this->get_index_label() === $index_name ) {
					continue;
				}

				if ( $pipeline_id = $indices[ $index_name ]['settings']['index']['default_pipeline'] ?? null ) {
					try {
						$pipeline = $this->search_engine_client->ingest()->getPipeline( [ 'id' => $pipeline_id ] );
					} catch ( \Exception $e ) {
						continue;
					}
					foreach ( $pipeline[ $pipeline_id ]['processors'] as $processor ) {
						if ( isset( $processor['text_embedding'] ) && $processor['text_embedding']['model_id'] === $model_id ) {
							$found_indices = true;
							break;
						}
					}
				}
			}

			if ( ! $found_indices ) {
				// Model will be unused after this index is deleted: remove the model
				try {
					$this->_delete_model( $model_id );
				} catch ( \Exception $e ) {
					// Continue
				}
			}
		}

		// Delete model(s) with index label (orphaned)
		$this->_delete_models_by_index_name();

		// Delete the index model connector(s)
		$this->_delete_connectors_by_index_name();
	}

	protected function _delete_model( string $model_id ): void {
		$this->search_engine_client->ml()->undeployModel( [ 'id' => $model_id ] );
		$this->search_engine_client->ml()->deleteModel( [ 'id' => $model_id ] );
	}

	protected function _is_external_model(): bool {
		return ( WPSOLR_Opensearch_Vector_Constants::MODEL_LOCAL_PRETRAINED_TRANSFORMER !== $this->get_config()['extra_parameters'][ WPSOLR_Hosting_Api_Abstract::FIELD_NAME_LLM_PROVIDER ] ?? '' );
	}

	protected function _get_models_by_name( $model_name ): array {
		$model_group_id = $this->_get_or_create_model_group( $this->search_engine_client->ml() );

		$models = $this->search_engine_client->ml()->searchModels( [
			'body' => [
				'query' => [
					'bool' => [
						'must' => [
							[ 'term' => [ 'model_group_id' => $model_group_id ] ],
							[ 'term' => [ 'name.keyword' => $model_name ] ],
						]
					]
				]
			]
		] );

		return $models['hits']['hits'] ?? [];
	}
}
