<?php

namespace CirrusSearch\Maintenance;

use CirrusSearch\Elastica\SearchAfter;
use CirrusSearch\Maintenance\Exception\IndexDumperException;
use CirrusSearch\SearchConfig;
use Elastica;
use Elastica\JSON;
use Elastica\Query;

/**
 * Dump an index to stdout
 *
 * @license GPL-2.0-or-later
 */

// @codeCoverageIgnoreStart
$IP = getenv( 'MW_INSTALL_PATH' );
if ( $IP === false ) {
	$IP = __DIR__ . '/../../..';
}
require_once "$IP/maintenance/Maintenance.php";
require_once __DIR__ . '/../includes/Maintenance/Maintenance.php';
// @codeCoverageIgnoreEnd

/**
 * Dump an index from elasticsearch.
 */
class DumpIndex extends Maintenance {

	/**
	 * @var string
	 */
	private $indexSuffix;

	/**
	 * @var string
	 */
	private $indexBaseName;

	/**
	 * @var string
	 */
	private $indexIdentifier;

	/**
	 * @var int number of docs per shard we export
	 */
	private $inputChunkSize = 500;

	/**
	 * @var bool
	 */
	private $logToStderr = false;

	/**
	 * @var int
	 */
	private $lastProgressPrinted;

	public function __construct() {
		parent::__construct();
		$this->addDescription( "Dump an index into a 'json' based format stdout. " .
			"This format complies to the elasticsearch bulk format and can be directly used " .
			"with a curl command like : " .
			"curl -s -XPOST localhost:9200/{index}/_bulk --data-binary @dump-file\n" .
			"Note that you need to specify the index in the URL because the bulk commands do not " .
			"contain the index name. Beware that the bulk import is not meant to import very large " .
			"files, sweet spot seems to be between 2000 and 5000 documents (see examples below)." .
			"\nThis always operates on a single cluster." .
			"\n\nExamples :\n" .
			" - Dump a general index :" .
			"\n\tdumpIndex --indexSuffix general\n" .
			" - Dump a large content index into compressed chunks of 100000 documents :" .
			"\n\tdumpIndex --indexSuffix content | split -d -a 9 -l 100000  " .
			"--filter 'gzip -c > \$FILE.txt.gz' - \"\" \n" .
			"\nYou can import the data with the following commands :\n" .
			" - Import chunks of 2000 documents :" .
			"\n\tcat dump | split -l 4000 --filter 'curl -s http://elastic:9200/{indexName}/_bulk " .
			"--data-binary @- > /dev/null'\n" .
			" - Import 3 chunks of 2000 documents in parallel :" .
			"\n\tcat dump | parallel --pipe -L 2 -N 2000 -j3 'curl -s http://elastic:9200/{indexName}/_bulk " .
			"--data-binary @- > /dev/null'" );
		$this->addOption( 'indexSuffix', 'Index to dump. Either content or general.', false, true );
		$this->addOption( 'indexType', 'BC form of --indexSuffix.', false, true );
		$this->addOption( 'baseName', 'What basename to use, ' .
			'defaults to wiki id.', false, true );
		$this->addOption( 'filter', 'Dump only the documents that match the filter query ' .
			'(queryString syntax).', false, true );
		$this->addOption( 'limit', 'Maximum number of documents to dump, 0 means no limit. Defaults to 0.',
			false, true );
		$this->addOption( 'indexIdentifier', 'Force the index identifier, use the alias otherwise.', false, true );
		$this->addOption( 'sourceFields', 'List of comma separated source fields to extract.', false, true );
	}

	/**
	 * @return bool True if this script only operates on clusters specified
	 *  in CirrusSearchManagedClusters. Can be set to false for read-only
	 *  scripts that don't care where they read from.
	 */
	protected function requireManagedCluster() {
		return false;
	}

	/** @inheritDoc */
	public function execute() {
		$this->disablePoolCountersAndLogging();

		$this->indexSuffix = $this->getBackCompatOption( 'indexSuffix', 'indexType' );
		$this->indexBaseName = $this->getOption( 'baseName',
			$this->getSearchConfig()->get( SearchConfig::INDEX_BASE_NAME ) );

		$indexSuffixes = $this->getConnection()->getAllIndexSuffixes();
		if ( !in_array( $this->indexSuffix, $indexSuffixes ) ) {
			$this->fatalError( 'indexSuffix option must be one of ' .
				implode( ', ', $indexSuffixes ) );
		}

		$this->indexIdentifier = $this->getOption( 'indexIdentifier' );

		$filter = null;
		if ( $this->hasOption( 'filter' ) ) {
			$filter = new Elastica\Query\QueryString( $this->getOption( 'filter' ) );
		}

		$limit = (int)$this->getOption( 'limit', 0 );

		$query = new Query();
		$query->setStoredFields( [ '_id', '_source' ] );
		$query->setSize( $this->inputChunkSize );
		// Elasticsearch docs (mapping-id-field.html) say that sorting by _id is restricted
		// from use, but at least on 7.10.2 this works as one would expect. The _id comes
		// from fielddata, this would be more efficient with a doc-values enabled field.
		// Additionally the cluster setting 'indices.id_field_data.enabled=false' breaks
		// this.
		$query->setSort( [
			[ 'page_id' => 'asc' ],
		] );
		$query->setTrackTotalHits( true );
		if ( $this->hasOption( 'sourceFields' ) ) {
			$sourceFields = explode( ',', $this->getOption( 'sourceFields' ) );
			$query->setSource( [ 'include' => $sourceFields ] );
		}
		if ( $filter ) {
			$bool = new \Elastica\Query\BoolQuery();
			$bool->addFilter( $filter );
			$query->setQuery( $bool );
		}

		$search = new \Elastica\Search( $this->getClient() );
		$search->setQuery( $query );
		$search->addIndex( $this->getIndex() );
		$searchAfter = new SearchAfter( $search );

		$totalDocsInIndex = -1;
		$totalDocsToDump = -1;
		$docsDumped = 0;

		$this->logToStderr = true;

		foreach ( $searchAfter as $results ) {
			if ( $totalDocsInIndex === -1 ) {
				$totalDocsInIndex = $results->getTotalHits();
				$totalDocsToDump = $limit > 0 ? $limit : $totalDocsInIndex;
				$this->output( "Dumping $totalDocsToDump documents ($totalDocsInIndex in the index)\n" );
			}

			foreach ( $results as $result ) {
				$document = [
					'_id' => $result->getId(),
					'_source' => $result->getSource()
				];
				$this->write( $document );
				$docsDumped++;
				if ( $docsDumped >= $totalDocsToDump ) {
					break;
				}
			}
			if ( $limit > 0 && $docsDumped > $limit ) {
				break;
			}
			$this->outputProgress( $docsDumped, $totalDocsToDump );
		}
		$this->output( "Dump done ($docsDumped docs).\n" );

		return true;
	}

	/**
	 * @param array $document Valid elasticsearch document to write to stdout
	 */
	public function write( array $document ) {
		$indexOp = [
			'index' => [
				'_id' => $document['_id']
			] ];

		// We use Elastica wrapper around json_encode.
		// Depending on PHP version JSON_ESCAPE_UNICODE will be used
		$this->writeLine( JSON::stringify( $indexOp ) );
		$this->writeLine( JSON::stringify( $document['_source'] ) );
	}

	/**
	 * @param string $data
	 * @throws IndexDumperException
	 */
	private function writeLine( $data ) {
		if ( !fwrite( STDOUT, $data . "\n" ) ) {
			throw new IndexDumperException( "Cannot write to standard output" );
		}
	}

	/**
	 * @return Elastica\Index being updated
	 */
	private function getIndex() {
		if ( $this->indexIdentifier ) {
			return $this->getConnection()->getIndex( $this->indexBaseName, $this->indexSuffix, $this->indexIdentifier );
		} else {
			return $this->getConnection()->getIndex( $this->indexBaseName, $this->indexSuffix );
		}
	}

	/**
	 * @param string $message
	 */
	public function outputIndented( $message ) {
		$this->output( "\t$message" );
	}

	/**
	 * @param string $message
	 * @param string|null $channel
	 */
	public function output( $message, $channel = null ) {
		if ( $this->mQuiet ) {
			return;
		}
		if ( $this->logToStderr ) {
			// We must log to stderr
			fwrite( STDERR, $message );
		} else {
			parent::output( $message );
		}
	}

	/**
	 * public because php 5.3 does not support accessing private
	 * methods in a closure.
	 * @param int $docsDumped
	 * @param int $limit
	 */
	public function outputProgress( $docsDumped, $limit ) {
		if ( $docsDumped <= 0 ) {
			return;
		}
		$pctDone = (int)( ( $docsDumped / $limit ) * 100 );
		if ( $this->lastProgressPrinted == $pctDone ) {
			return;
		}
		$this->lastProgressPrinted = $pctDone;
		if ( ( $pctDone % 2 ) == 0 ) {
			$this->outputIndented( "$pctDone% done...\n" );
		}
	}

	/**
	 * @return Elastica\Client
	 */
	protected function getClient() {
		return $this->getConnection()->getClient();
	}

}

// @codeCoverageIgnoreStart
$maintClass = DumpIndex::class;
require_once RUN_MAINTENANCE_IF_MAIN;
// @codeCoverageIgnoreEnd
