/*
 * SPDX-License-Identifier: Apache-2.0
 *
 * The OpenSearch Contributors require contributions made to
 * this file be licensed under the Apache-2.0 license or a
 * compatible open source license.
 */

package org.opensearch.action.admin.indices.streamingingestion.state;

import org.apache.lucene.store.AlreadyClosedException;
import org.opensearch.Version;
import org.opensearch.action.support.ActionFilters;
import org.opensearch.cluster.ClusterState;
import org.opensearch.cluster.block.ClusterBlockException;
import org.opensearch.cluster.block.ClusterBlockLevel;
import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
import org.opensearch.cluster.metadata.IngestionSource;
import org.opensearch.cluster.metadata.Metadata;
import org.opensearch.cluster.node.DiscoveryNode;
import org.opensearch.cluster.routing.ShardRouting;
import org.opensearch.cluster.routing.ShardsIterator;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.settings.Settings;
import org.opensearch.core.action.support.DefaultShardOperationFailedException;
import org.opensearch.core.index.Index;
import org.opensearch.core.index.shard.ShardId;
import org.opensearch.index.IndexService;
import org.opensearch.index.shard.IndexShard;
import org.opensearch.index.shard.ShardNotFoundException;
import org.opensearch.indices.IndicesService;
import org.opensearch.indices.replication.common.ReplicationType;
import org.opensearch.telemetry.tracing.noop.NoopTracer;
import org.opensearch.test.OpenSearchTestCase;
import org.opensearch.test.transport.MockTransportService;
import org.opensearch.threadpool.TestThreadPool;
import org.opensearch.threadpool.ThreadPool;
import org.opensearch.transport.TransportService;
import org.opensearch.transport.client.node.NodeClient;

import java.util.Collections;
import java.util.List;

import static org.hamcrest.Matchers.equalTo;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

public class TransportGetIngestionStateActionTests extends OpenSearchTestCase {

    private TestThreadPool threadPool;
    private ClusterService clusterService;
    private TransportService transportService;
    private IndicesService indicesService;
    private ActionFilters actionFilters;
    private IndexNameExpressionResolver indexNameExpressionResolver;
    private TransportGetIngestionStateAction action;
    private NodeClient client;

    @Override
    public void setUp() throws Exception {
        super.setUp();
        threadPool = new TestThreadPool(getClass().getName());
        clusterService = mock(ClusterService.class);
        transportService = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool, NoopTracer.INSTANCE);
        indicesService = mock(IndicesService.class);
        actionFilters = mock(ActionFilters.class);
        indexNameExpressionResolver = mock(IndexNameExpressionResolver.class);
        client = mock(NodeClient.class);

        action = new TransportGetIngestionStateAction(
            clusterService,
            transportService,
            indicesService,
            actionFilters,
            indexNameExpressionResolver,
            client
        );
    }

    @Override
    public void tearDown() throws Exception {
        super.tearDown();
        ThreadPool.terminate(threadPool, 30, java.util.concurrent.TimeUnit.SECONDS);
    }

    public void testShards() {
        GetIngestionStateRequest request = new GetIngestionStateRequest(new String[] { "test-index" });
        request.setShards(new int[] { 0, 1 });
        ClusterState clusterState = mock(ClusterState.class);
        ShardsIterator shardsIterator = mock(ShardsIterator.class);
        when(clusterState.routingTable()).thenReturn(mock(org.opensearch.cluster.routing.RoutingTable.class));
        when(clusterState.routingTable().allShardsSatisfyingPredicate(any(), any())).thenReturn(shardsIterator);

        Metadata metadata = mock(Metadata.class);
        IndexMetadata indexMetadata = mock(IndexMetadata.class);
        Settings settings = Settings.builder().put(IndexMetadata.SETTING_REPLICATION_TYPE, ReplicationType.DOCUMENT.toString()).build();

        // Set up the mocks
        when(clusterState.metadata()).thenReturn(metadata);
        when(metadata.index("test-index")).thenReturn(indexMetadata);
        when(indexMetadata.getSettings()).thenReturn(settings);

        ShardsIterator result = action.shards(clusterState, request, new String[] { "test-index" });
        assertThat(result, equalTo(shardsIterator));
    }

    public void testCheckGlobalBlock() {
        GetIngestionStateRequest request = new GetIngestionStateRequest(new String[] { "test-index" });
        ClusterState clusterState = mock(ClusterState.class);
        when(clusterState.blocks()).thenReturn(mock(org.opensearch.cluster.block.ClusterBlocks.class));
        when(clusterState.blocks().globalBlockedException(ClusterBlockLevel.METADATA_READ)).thenReturn(null);

        ClusterBlockException result = action.checkGlobalBlock(clusterState, request);
        assertThat(result, equalTo(null));
    }

    public void testCheckRequestBlock() {
        GetIngestionStateRequest request = new GetIngestionStateRequest(new String[] { "test-index" });
        ClusterState clusterState = mock(ClusterState.class);
        when(clusterState.blocks()).thenReturn(mock(org.opensearch.cluster.block.ClusterBlocks.class));
        when(clusterState.blocks().indicesBlockedException(ClusterBlockLevel.METADATA_READ, request.indices())).thenReturn(null);

        ClusterBlockException result = action.checkRequestBlock(clusterState, request, new String[] { "test-index" });
        assertThat(result, equalTo(null));
    }

    public void testShardOperation() {
        GetIngestionStateRequest request = new GetIngestionStateRequest(new String[] { "test-index" });
        ShardRouting shardRouting = mock(ShardRouting.class);
        IndexService indexService = mock(IndexService.class);
        IndexShard indexShard = mock(IndexShard.class);
        ShardIngestionState expectedState = new ShardIngestionState("test-index", 0, "POLLING", "DROP", true, false, "", true, "node_name");

        when(shardRouting.shardId()).thenReturn(mock(ShardId.class));
        when(shardRouting.shardId().getIndex()).thenReturn(mock(Index.class));
        when(shardRouting.shardId().id()).thenReturn(0);
        when(indicesService.indexServiceSafe(any())).thenReturn(indexService);
        when(indexService.getShard(0)).thenReturn(indexShard);
        when(indexShard.routingEntry()).thenReturn(mock(org.opensearch.cluster.routing.ShardRouting.class));
        when(indexShard.getIngestionState()).thenReturn(expectedState);

        DiscoveryNode localNode = mock(DiscoveryNode.class);
        when(localNode.getName()).thenReturn("node_name");
        when(clusterService.localNode()).thenReturn(localNode);

        ShardIngestionState result = action.shardOperation(request, shardRouting);
        assertThat(result.getIndex(), equalTo(expectedState.getIndex()));
        assertThat(result.getShardId(), equalTo(expectedState.getShardId()));
        assertThat(result.getPollerState(), equalTo(expectedState.getPollerState()));
        assertThat(result.getErrorPolicy(), equalTo(expectedState.getErrorPolicy()));
        assertThat(result.getNodeName(), equalTo(expectedState.getNodeName()));

    }

    public void testShardOperationWithShardNotFoundException() {
        GetIngestionStateRequest request = new GetIngestionStateRequest(new String[] { "test-index" });
        ShardRouting shardRouting = mock(ShardRouting.class);
        IndexService indexService = mock(IndexService.class);
        IndexShard indexShard = mock(IndexShard.class);

        when(shardRouting.shardId()).thenReturn(mock(ShardId.class));
        when(shardRouting.shardId().getIndex()).thenReturn(mock(Index.class));
        when(shardRouting.shardId().id()).thenReturn(0);
        when(indicesService.indexServiceSafe(any())).thenReturn(indexService);
        when(indexService.getShard(0)).thenReturn(indexShard);
        when(indexShard.routingEntry()).thenReturn(null);

        expectThrows(ShardNotFoundException.class, () -> action.shardOperation(request, shardRouting));
    }

    public void testShardOperationWithAlreadyClosedException() {
        GetIngestionStateRequest request = new GetIngestionStateRequest(new String[] { "test-index" });
        ShardRouting shardRouting = mock(ShardRouting.class);
        IndexService indexService = mock(IndexService.class);
        IndexShard indexShard = mock(IndexShard.class);

        when(shardRouting.shardId()).thenReturn(mock(ShardId.class));
        when(shardRouting.shardId().getIndex()).thenReturn(mock(Index.class));
        when(shardRouting.shardId().id()).thenReturn(0);
        when(indicesService.indexServiceSafe(any())).thenReturn(indexService);
        when(indexService.getShard(0)).thenReturn(indexShard);
        when(indexShard.routingEntry()).thenReturn(mock(org.opensearch.cluster.routing.ShardRouting.class));
        when(indexShard.getIngestionState()).thenThrow(new AlreadyClosedException("shard is closed"));

        expectThrows(ShardNotFoundException.class, () -> action.shardOperation(request, shardRouting));
    }

    public void testNewResponse() {
        GetIngestionStateRequest request = new GetIngestionStateRequest(new String[] { "test-index" });
        List<ShardIngestionState> responses = Collections.singletonList(
            new ShardIngestionState("test-index", 0, "POLLING", "DROP", true, false, "", true, "node")
        );
        List<DefaultShardOperationFailedException> shardFailures = Collections.emptyList();
        ClusterState clusterState = mock(ClusterState.class);

        GetIngestionStateResponse response = action.newResponse(request, 1, 1, 0, responses, shardFailures, clusterState);

        assertThat(response.getTotalShards(), equalTo(1));
        assertThat(response.getSuccessfulShards(), equalTo(1));
        assertThat(response.getFailedShards(), equalTo(0));
        assertThat(response.getShardStates().length, equalTo(1));
        assertThat(response.getShardStates()[0].getIndex(), equalTo("test-index"));
        assertThat(response.getShardStates()[0].getShardId(), equalTo(0));
    }

    public void testShardsWithAllActiveIngestionEnabled() {
        GetIngestionStateRequest request = new GetIngestionStateRequest(new String[] { "test-index" });
        request.setShards(new int[] { 0, 1 });
        ClusterState clusterState = mock(ClusterState.class);
        ShardsIterator shardsIterator = mock(ShardsIterator.class);

        when(clusterState.routingTable()).thenReturn(mock(org.opensearch.cluster.routing.RoutingTable.class));
        when(clusterState.routingTable().allShardsSatisfyingPredicate(any(), any())).thenReturn(shardsIterator);

        Metadata metadata = mock(Metadata.class);
        IndexMetadata indexMetadata = mock(IndexMetadata.class);
        IngestionSource ingestionSource = mock(IngestionSource.class);

        // Set up mocks for all-active ingestion enabled
        when(clusterState.metadata()).thenReturn(metadata);
        when(metadata.index("test-index")).thenReturn(indexMetadata);
        when(indexMetadata.useIngestionSource()).thenReturn(true);
        when(indexMetadata.getIngestionSource()).thenReturn(ingestionSource);
        when(ingestionSource.isAllActiveIngestionEnabled()).thenReturn(true);

        ShardsIterator result = action.shards(clusterState, request, new String[] { "test-index" });
        assertThat(result, equalTo(shardsIterator));
    }
}
