<?php

namespace CirrusSearch\Maintenance;

use CirrusSearch\BuildDocument\Completion\SuggestBuilder;
use CirrusSearch\CirrusTestCase;
use CirrusSearch\Connection;
use CirrusSearch\MetaStore\MetaVersionStore;
use Elastica\Client;
use Elastica\Document;
use Elastica\Exception\ResponseException;
use Elastica\Index;
use Elastica\Request;
use Elastica\Response;

/**
 * @covers \CirrusSearch\Maintenance\CompletionSuggesterIndexer
 */
class CompletionSuggesterIndexerTest extends CirrusTestCase {

	private const BASE_NAME = "mytestwiki";
	private const MASTER_TIMEOUT = "30s";
	private const REPLICA_COUNT = 2;
	private const INDEX_RETRY_ATTEMPTS = 1;
	private const FAKE_ANALYSIS_CONFIG = [ "my_analysis_config" => [] ];
	/**
	 * @var Client|\PHPUnit\Framework\MockObject\MockObject
	 */
	private $client;
	/**
	 * @var Printer|\PHPUnit\Framework\MockObject\MockObject
	 */
	private $printer;
	/**
	 * @var Connection|\PHPUnit\Framework\MockObject\MockObject
	 */
	private $connection;
	private ConfigUtils $utils;
	/**
	 * @var MetaVersionStore|\PHPUnit\Framework\MockObject\MockObject
	 */
	private $versionStore;
	/**
	 * @var SuggesterAnalysisConfigBuilder|\PHPUnit\Framework\MockObject\MockObject
	 */
	private $analysisConfigBuilder;

	/**
	 * @var SuggestBuilder|\PHPUnit\Framework\MockObject\MockObject
	 */
	private $builder;
	private int $batchId;

	private function newFailedRefreshResponse(): Response {
		return new Response(
			[ "_shards" => [ "total" => 1, "failed" => 1, "successful" => 0 ] ],
			200
		);
	}

	private function newSuccessRefreshResponse(): Response {
		return new Response(
			[ "_shards" => [ "total" => 1, "failed" => 0, "successful" => 1 ] ],
			200
		);
	}

	private function assertQueryOldBatch( string $indexName, int $chunkSize, string $path, string $method, array $body ): void {
		$this->assertEquals( "$indexName/_search", $path );
		$this->assertEquals( 'POST', $method );
		$this->assertTrue( isset( $body['query'] ) );
		$this->assertTrue( isset( $body['sort'] ) );
		$this->assertTrue( isset( $body['track_total_hits'] ) );
		$this->assertTrue( isset( $body['size'] ) );
		$this->assertEquals( [ 'bool' => [ 'must_not' => [ [ "term" => [ "batch_id" => $this->batchId ] ] ] ] ],
			$body['query'] );
		$this->assertEquals( [ [ '_id' => 'asc' ] ], $body['sort'] );
		$this->assertTrue( $body['track_total_hits'] );
		$this->assertEquals( $body['size'], $chunkSize );
	}

	private function assertRefresh( string $indexName, string $path, string $method, array $body, array $param ): void {
		$this->assertEquals( "$indexName/_refresh", $path );
		$this->assertEquals( 'POST', $method );
		$this->assertEquals( [], $body );
		$this->assertEquals( [], $param );
	}

	private function assertCountIndexDocs( string $indexName, string $path, string $method, array $body, array $param ): void {
		$this->assertEquals( "$indexName/_search", $path );
		$this->assertEquals( 'POST', $method );
		$this->assertEquals( [
				"query" => [ "match_all" => new \stdClass() ],
				"track_total_hits" => true,
				"size" => 0,
			],
			$body
		);
		$this->assertEquals( [ "search_type" => "query_then_fetch" ], $param );
	}

	private function assertOptimize( string $indexName, string $path, string $method, array $body, array $param ): void {
		$this->assertEquals( "$indexName/_forcemerge", $path );
		$this->assertEquals( 'POST', $method );
		$this->assertEquals( [], $body );
		$this->assertEquals( [ 'max_num_segments' => 1 ], $param );
	}

	private function assertCheckIndexExistence( string $indexName, string $path, string $method, array $body, array $param ): void {
		$this->assertEquals( $indexName, $path );
		$this->assertEquals( 'HEAD', $method );
		$this->assertEquals( [], $body );
		$this->assertEquals( [], $param );
	}

	private function assertWaitForGreen( string $indexName, string $path, string $method, array $body, array $param ): void {
		$this->assertEquals( "_cluster/health/$indexName", $path );
		$this->assertEquals( "GET", $method );
		$this->assertEquals( [], $body );
		$this->assertEquals( [], $param );
	}

	protected function setUp(): void {
		parent::setUp(); // TODO: Change the autogenerated stub
		$this->client = $this->getMockBuilder( Client::class )
			->setProxyTarget( new Client( [ 'connections' => [] ] ) )
			->onlyMethods( [ 'request' ] )
			->getMock();
		$this->connection = $this->getMockBuilder( Connection::class )
			->disableOriginalConstructor()
			->onlyMethods( [ 'getClient', 'getConfig' ] )
			->getMock();
		$this->connection->method( 'getClient' )->willReturn( $this->client );
		$this->printer = $this->createMock( Printer::class );
		$this->utils = new ConfigUtils( $this->client, $this->printer );
		$this->versionStore = $this->createMock( MetaVersionStore::class );
		$this->analysisConfigBuilder = $this->createMock( SuggesterAnalysisConfigBuilder::class );
		$this->analysisConfigBuilder->method( 'buildConfig' )->willReturn( self::FAKE_ANALYSIS_CONFIG );
		$this->builder = $this->createMock( SuggestBuilder::class );
		$this->batchId = random_int( 1, 100 );
		$this->builder->method( 'getBatchId' )->willReturn( $this->batchId );
	}

	public function testPrepareRecycle(): void {
		$this->versionStore->expects( $this->never() )->method( $this->anything() );
		$this->client->expects( $this->never() )->method( 'request' );
		$config = $this->indexerConfig( true, 1, false, false, 0 );
		$indexer = new CompletionSuggesterIndexer(
			$this->connection,
			new Index( $this->client, "unused" ),
			null,
			$this->builder,
			$this->printer,
			$this->utils,
			$this->versionStore,
			$this->analysisConfigBuilder,
			$config
		);

		$indexer->prepare();
	}

	public function testPrepareRebuild(): void {
		$config = $this->indexerConfig( false, 1, false, false, 0 );
		$searchConfig = $this->newHashSearchConfig( [] );
		$this->connection->method( 'getConfig' )->willReturn( $searchConfig );
		$state = "check_index";
		$indexName = self::BASE_NAME . "_titlesuggest_first";
		$indexer = new CompletionSuggesterIndexer(
			$this->connection,
			new Index( $this->client, $indexName ),
			null,
			$this->builder,
			$this->printer,
			$this->utils,
			$this->versionStore,
			$this->analysisConfigBuilder,
			$config
		);

		$this->versionStore->expects( $this->never() )->method( $this->anything() );
		$this->client
			->method( 'request' )
			->willReturnCallback(
				function ( string $path, string $method, array $body, array $param ) use ( &$state, $indexName, $config ) {
					switch ( $state ) {
						case "check_index":
							$this->assertCheckIndexExistence( $indexName, $path, $method, $body, $param );
							$state = "create_index";
							return new Response( [], 404 );
						case "create_index":
							$this->assertEquals( $indexName, $path );
							$this->assertEquals( 'PUT', $method );
							$this->assertTrue( isset( $body["settings"]["index"] ) );
							$this->assertEquals(
								[
									"number_of_shards" => $config->getShardCount(),
									"auto_expand_replicas" => "0-0",
									"refresh_interval" => -1,
									"routing.allocation.total_shards_per_node" => $config->getMaxShardPerNode(),
									"analysis" => self::FAKE_ANALYSIS_CONFIG
								],
								$body["settings"]["index"]
							);
							$this->assertEquals( [ "master_timeout" => "30s" ], $param );
							$state = "wait_for_green_but_yellow";
							return new Response( [], 200 );
						case "wait_for_green_but_yellow":
						case "wait_for_green":
							$yellow = $state === "wait_for_green_but_yellow";
							$this->assertWaitForGreen( $indexName, $path, $method, $body, $param );
							$state = $yellow ? "wait_for_green" : "end";
							return new Response( [ "status" => $yellow ? 'yellow' : 'green' ], 200 );
						default:
							$this->fail( "unknown state $state" );
					}
				} );

		$indexer->prepare();
	}

	public function testFinishRecycle(): void {
		$config = $this->indexerConfig( true, 1, false, false, 0 );
		$searchConfig = $this->newHashSearchConfig( [] );
		$this->connection->method( 'getConfig' )->willReturn( $searchConfig );
		$state = "expunge_deletes_before";
		$produceError = true;
		$indexName = self::BASE_NAME . "_titlesuggest_first";
		$indexer = new CompletionSuggesterIndexer(
			$this->connection,
			new Index( $this->client, $indexName ),
			null,
			$this->builder,
			$this->printer,
			$this->utils,
			$this->versionStore,
			$this->analysisConfigBuilder,
			$config
		);
		$this->versionStore->expects( $this->never() )->method( $this->anything() );
		$this->client
			->expects( $this->atLeast( 10 ) )
			->method( 'request' )
			->willReturnCallback(
				function ( string $path, string $method, array $body, array $param ) use ( &$state, $indexName, $config, &$produceError ) {
					switch ( $state ) {
						case "expunge_deletes_before":
							$this->assertEquals( "$indexName/_forcemerge", $path );
							$this->assertEquals( 'POST', $method );
							$this->assertEquals( [], $body );
							$this->assertEquals( [ 'only_expunge_deletes' => 'true', 'flush' => 'false' ], $param );
							$state = "refresh_before";
							return new Response( [], 200 );
						case "refresh_before":
							$this->assertRefresh( $indexName, $path, $method, $body, $param );
							if ( $produceError ) {
								$produceError = false;
								return $this->newFailedRefreshResponse();
							} else {
								$produceError = true;
								$state = "query_old_batch_page_1";
								return $this->newSuccessRefreshResponse();
							}
						case "query_old_batch_page_1":
						case "query_old_batch_page_2":
							$firstPage = $state === "query_old_batch_page_1";
							$this->assertQueryOldBatch( $indexName, $config->getIndexChunkSize(), $path, $method, $body );
							$this->assertEquals( [], $param );
							$state = $firstPage ? "delete_by_query_page_1" : "delete_by_query_page_2";
							return new Response( [
								'_shards' => [ 'total' => 1, 'successful' => 1, 'failed' => 0 ],
								"hits" => [
									"total" => [ "value" => 2, "relation" => "eq" ],
									"hits" => [ [ "_id" => "1", "sort" => [ "1" ] ] ]
								]
							] );
						case "delete_by_query_page_1":
						case "delete_by_query_page_2":
							$this->assertEquals( "$indexName/_delete_by_query", $path );
							$this->assertEquals( "POST", $method );
							$this->assertEquals(
								[ "query" => [ "ids" => [ "values" => [ "1" ] ] ] ],
								$body
							);
							$this->assertEquals( [], $param );
							$lastPage = $state === "delete_by_query_page_2";
							if ( $produceError ) {
								$produceError = false;
								throw new ResponseException( new Request( "unused" ), new Response( [], 500 ) );
							} else {
								$state = $lastPage ? "last_query_old_batch_page" : "query_old_batch_page_2";
								$produceError = true;
								return new Response( [], 200 );
							}
						case "last_query_old_batch_page":
							$state = "optimize";
							return new Response( [
								'_shards' => [ 'total' => 1, 'successful' => 1, 'failed' => 0 ],
								"hits" => [
									"total" => [ "value" => 2, "relation" => "eq" ],
									"hits" => []
								]
							] );
						case "optimize":
							$this->assertOptimize( $indexName, $path, $method, $body, $param );
							$state = "refresh_after";
							return new Response( [], 200 );
						case "refresh_after":
							$this->assertRefresh( $indexName, $path, $method, $body, $param );
							if ( $produceError ) {
								$produceError = false;
								return $this->newFailedRefreshResponse();
							} else {
								$produceError = true;
								$state = "count";
								return $this->newSuccessRefreshResponse();
							}
						case "count":
							$this->assertCountIndexDocs( $indexName, $path, $method, $body, $param );
							if ( $produceError ) {
								$produceError = false;
								return new Response( [ "broken" ], 200 );
							} else {
								$produceError = true;
								$state = "end";
								return new Response( [ "hits" => [ "total" => [ "value" => 2 ] ] ], 200 );
							}
						default:
							$this->fail( "unknown state $state for $method $path" );
					}
				} );
		$indexer->finish();
	}

	public function testRebuildFinish() {
		$config = $this->indexerConfig( false, 1, false, false, 0 );
		$searchConfig = $this->newHashSearchConfig( [] );
		$this->connection->method( 'getConfig' )->willReturn( $searchConfig );
		$state = "optimize";
		$produceError = true;
		$oldIndexName = self::BASE_NAME . "_titlesuggest_old";
		$newIndexName = self::BASE_NAME . "_titlesuggest_new";
		$indexer = new CompletionSuggesterIndexer(
			$this->connection,
			new Index( $this->client, $newIndexName ),
			new Index( $this->client, $oldIndexName ),
			$this->builder,
			$this->printer,
			$this->utils,
			$this->versionStore,
			$this->analysisConfigBuilder,
			$config
		);

		$this->client
			->method( 'request' )
			->willReturnCallback(
				function ( string $path, string $method, array $body, array $param )
				use ( &$state, $newIndexName, $oldIndexName, $config, &$produceError ) {
					switch ( $state ) {
						case "optimize":
							$this->assertOptimize( $newIndexName, $path, $method, $body, $param );
							$state = "refresh";
							return new Response( [], 200 );
						case "refresh":
							$this->assertRefresh( $newIndexName, $path, $method, $body, $param );
							if ( $produceError ) {
								$produceError = false;
								return $this->newFailedRefreshResponse();
							} else {
								$state = "check_old_index_existence";
								$produceError = true;
								return $this->newSuccessRefreshResponse();
							}
						case "check_old_index_existence":
							$this->assertCheckIndexExistence( $oldIndexName, $path, $method, $body, $param );
							$state = "count_old";
							return new Response( [], 200 );
						case "count_old":
							$this->assertCountIndexDocs( $oldIndexName, $path, $method, $body, $param );
							$state = "count_too_low";
							return new Response( [ "hits" => [ "total" => [ "value" => 200000 ] ] ], 200 );
						case "count_too_low":
							$this->assertCountIndexDocs( $newIndexName, $path, $method, $body, $param );
							$state = "count_ok";
							return new Response( [ "hits" => [ "total" => [ "value" => 100000 ] ] ], 200 );
						case "count_ok":
							$this->assertCountIndexDocs( $newIndexName, $path, $method, $body, $param );
							$state = "enable_replicas";
							return new Response( [ "hits" => [ "total" => [ "value" => 200000 ] ] ], 200 );
						case "enable_replicas":
							$this->assertEquals( "PUT", $method );
							$this->assertEquals( "$newIndexName/_settings", $path );
							$this->assertEquals(
								[
									"index" => [
										"auto_expand_replicas" => $config->getReplicaCount()
									]
								],
								$body
							);
							$this->assertEquals( [ "master_timeout" => "30s" ], $param );
							$state = "wait_for_green_but_yellow";
							return new Response( [], 200 );
						case "wait_for_green_but_yellow":
						case "wait_for_green":
							$yellow = $state === "wait_for_green_but_yellow";
							$this->assertWaitForGreen( $newIndexName, $path, $method, $body, $param );
							$state = $yellow ? "wait_for_green" : "validate_alias_get_current_aliases";
							return new Response( [ "status" => $yellow ? 'yellow' : 'green' ], 200 );
						case "validate_alias_get_current_aliases":
							$this->assertEquals( "GET", $method );
							$this->assertEquals( "_alias/" . self::BASE_NAME . "_" . Connection::TITLE_SUGGEST_INDEX_SUFFIX, $path );
							$this->assertEquals( [], $body );
							$this->assertEquals( [], $param );
							$state = "update_aliases";
							return new Response( [
								$oldIndexName => []
							], 200 );
						case "update_aliases":
							$this->assertEquals( "POST", $method );
							$this->assertEquals( "_aliases", $path );
							$aliasName = self::BASE_NAME . "_" . Connection::TITLE_SUGGEST_INDEX_SUFFIX;
							$this->assertEquals( [
								"actions" => [
									[
										"remove" => [
											"index" => $oldIndexName,
											"alias" => $aliasName
										],
									],
									[
										"add" => [
											"index" => $newIndexName,
											"alias" => $aliasName
										],
									]
								]
							], $body );
							$this->assertEquals( [ 'master_timeout' => self::MASTER_TIMEOUT ], $param );
							$state = "update_version_index";
							return new Response( [
								$oldIndexName => []
							], 200 );
						case "check_old_index_existence_before_deletion":
							$this->assertCheckIndexExistence( $oldIndexName, $path, $method, $body, $param );
							$state = "delete_old_index";
							return new Response( [], 200 );
						case "delete_old_index":
							$this->assertEquals( "DELETE", $method );
							$this->assertEquals( $oldIndexName, $path );
							$this->assertEquals( [], $body );
							$this->assertEquals( [ 'master_timeout' => self::MASTER_TIMEOUT ], $param );
							$state = "end";
							return new Response( [], 200 );

						default:
							$this->fail( "unknown state $state for $method $path" );
					}
				}
			);
		$this->versionStore->expects( $this->once() )->method( 'update' )->willReturnCallback(
			function ( string $baseName, string $typeName, bool $altIndex = false, int $altIndexId = 0 ) use ( &$state, $config ) {
				$this->assertEquals( "update_version_index", $state );
				$this->assertEquals( self::BASE_NAME, $baseName );
				$this->assertEquals( Connection::TITLE_SUGGEST_INDEX_SUFFIX, $typeName );
				$this->assertEquals( $config->isAltIndex(), $altIndex );
				$this->assertEquals( $config->getAltIndexId(), $altIndexId );
				$state = "check_old_index_existence_before_deletion";
			}
		);
		$indexer->finish();
	}

	public function testIndexData() {
		$sourceDoc = 16;
		$docMultiplier = 2;
		$totalSuggestions = $sourceDoc * $docMultiplier;
		$chunkSize = 10;

		$this->builder->method( 'build' )
			->willReturnCallback(
				static fn ( array $docs ) => array_fill( 0, count( $docs ) * 2, new Document() )
			);
		$indexerConfig = $this->indexerConfig( false, $chunkSize, false, false, 0 );

		$index = new Index( $this->client, self::BASE_NAME . "_titlesuggest_first" );
		$indexer = new CompletionSuggesterIndexer(
			$this->connection,
			$index,
			null,
			$this->builder,
			$this->printer,
			$this->utils,
			$this->versionStore,
			$this->analysisConfigBuilder,
			$indexerConfig
		);

		$this->client->expects( $this->exactly( ceil( $totalSuggestions / $chunkSize ) ) )
			->method( 'request' )
			->willReturnCallback(
				function ( string $path, string $method, string $bulkData, array $params, $contentType ) use ( $index, $chunkSize ) {
					$this->assertEquals( "_bulk", $path );
					$this->assertEquals( "POST", $method );
					$this->assertEquals( Request::NDJSON_CONTENT_TYPE, $contentType );
					$this->assertEquals( [], $params );
					$lines = explode( "\n", $bulkData );
					array_pop( $lines );
					$this->assertTrue( count( $lines ) <= $chunkSize );
					foreach ( $lines as $line ) {
						$doc = json_decode( $line, true );
						$this->assertEquals( JSON_ERROR_NONE, json_last_error() );
						$this->assertArrayHasKey( 'index', $doc );
						$this->assertArrayHasKey( '_index', $doc['index'] );
						$this->assertEquals( $index->getName(), $doc['index']['_index'], );
					}
					return new Response(
						[
							'items' => array_fill( 0, count( $lines ), [ "index" => [ 'result' => 'created' ] ] )
						],
						200 );
				}
			);

		$indexer->addDocument( array_fill( 0, $sourceDoc, [] ) );
		$indexer->flushSuggestDocs();
		$this->assertEquals(
			[
				'doc_built' => $totalSuggestions,
				'doc_indexed' => 0,
				'bulk_requests' => (int)ceil( $totalSuggestions / $chunkSize ),
				'retried_bulk_requests' => 0,
				'doc_sent' => $totalSuggestions,
				'index_results' => [
					'created' => $totalSuggestions,
					'updated' => 0,
					'noop' => 0,
					'unknown' => 0,
					'error' => 0,
				],
			],
			$indexer->getIndexingStats()
		);
	}

	private function indexerConfig(
		bool $recycle,
		int $chunkSize,
		bool $force,
		bool $altIndex,
		int $altIndexId
	): CompletionSuggesterIndexerConfig {
		return new CompletionSuggesterIndexerConfig(
			self::BASE_NAME,
			$altIndex,
			$altIndexId,
			self::REPLICA_COUNT,
			"0-2",
			1,
			$recycle,
			$chunkSize, self::MASTER_TIMEOUT,
			5,
			self::INDEX_RETRY_ATTEMPTS,
			null,
			null,
			true,
			$force
		);
	}

}
