/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.solr.cloud;

import java.io.IOException;
import java.lang.invoke.MethodHandles;
import java.net.URI;
import java.net.URL;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.apache.solr.SolrTestCaseJ4.SuppressSSL;
import org.apache.solr.client.solrj.SolrClient;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.cloud.SocketProxy;
import org.apache.solr.client.solrj.request.CollectionAdminRequest;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.SolrInputDocument;
import org.apache.solr.common.cloud.CollectionStatePredicate;
import org.apache.solr.common.cloud.DocCollection;
import org.apache.solr.common.cloud.Replica;
import org.apache.solr.common.cloud.Slice;
import org.apache.solr.common.cloud.ZkStateReader;
import org.apache.solr.common.util.TimeSource;
import org.apache.solr.core.SolrCore;
import org.apache.solr.embedded.JettySolrRunner;
import org.apache.solr.util.TestInjection;
import org.apache.solr.util.TimeOut;
import org.apache.zookeeper.KeeperException;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@SuppressSSL(bugUrl = "https://issues.apache.org/jira/browse/SOLR-5776")
public class TestPullReplicaErrorHandling extends SolrCloudTestCase {

  private static final int REPLICATION_TIMEOUT_SECS = 10;

  private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
  private static Map<URI, SocketProxy> proxies;
  private static Map<URI, JettySolrRunner> jettys;

  private String collectionName = null;

  private String suggestedCollectionName() {
    return (getTestClass().getSimpleName().replace("Test", "")
            + "_"
            + getSaferTestName().split(" ")[0])
        .replaceAll("(.)(\\p{Upper})", "$1_$2")
        .toLowerCase(Locale.ROOT);
  }

  @BeforeClass
  public static void setupCluster() throws Exception {
    System.setProperty("solr.zkclienttimeout", "20000");

    configureCluster(4).addConfig("conf", configset("cloud-minimal")).configure();
    // Add proxies
    proxies = new HashMap<>(cluster.getJettySolrRunners().size());
    jettys = new HashMap<>(cluster.getJettySolrRunners().size());
    for (JettySolrRunner jetty : cluster.getJettySolrRunners()) {
      SocketProxy proxy = new SocketProxy();
      jetty.setProxyPort(proxy.getListenPort());
      cluster.stopJettySolrRunner(jetty); // TODO: Can we avoid this restart
      cluster.startJettySolrRunner(jetty);
      cluster.waitForAllNodes(30);
      proxy.open(jetty.getBaseUrl().toURI());
      if (log.isInfoEnabled()) {
        log.info("Adding proxy for URL: {}. Proxy: {}", jetty.getBaseUrl(), proxy.getUrl());
      }
      proxies.put(proxy.getUrl(), proxy);
      jettys.put(proxy.getUrl(), jetty);
    }
  }

  @AfterClass
  public static void tearDownCluster() {
    if (null != proxies) {
      for (SocketProxy proxy : proxies.values()) {
        proxy.close();
      }
      proxies = null;
    }
    jettys = null;
    TestInjection.reset();
  }

  @Override
  public void setUp() throws Exception {
    super.setUp();
    collectionName = suggestedCollectionName();
    expectThrows(SolrException.class, () -> getCollectionState(collectionName));
    cluster.waitForAllNodes(30);
  }

  @Override
  public void tearDown() throws Exception {
    if (cluster.getSolrClient().getClusterState().getCollectionOrNull(collectionName) != null) {
      log.info("tearDown deleting collection");
      CollectionAdminRequest.deleteCollection(collectionName).process(cluster.getSolrClient());
      log.info("Collection deleted");
      waitForDeletion(collectionName);
    }
    collectionName = null;
    super.tearDown();
  }

  public void testCantConnectToPullReplica() throws Exception {
    int numShards = 2;
    CollectionAdminRequest.createCollection(collectionName, "conf", numShards, 1, 0, 1)
        .process(cluster.getSolrClient());
    cluster.waitForActiveCollection(collectionName, numShards, numShards * 2);
    addDocs(10);
    DocCollection docCollection = assertNumberOfReplicas(numShards, 0, numShards, false, true);
    Slice s = docCollection.getSlices().iterator().next();
    SocketProxy proxy = getProxyForReplica(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0));
    try {
      proxy.close();
      for (int i = 1; i <= 10; i++) {
        addDocs(10 + i);
        try (SolrClient leaderClient = getHttpSolrClient(s.getLeader().getCoreUrl())) {
          assertNumDocs(10 + i, leaderClient);
        }
      }

      SolrServerException e =
          expectThrows(
              SolrServerException.class,
              () -> {
                try (SolrClient pullReplicaClient =
                    getHttpSolrClient(
                        s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0).getCoreUrl())) {
                  pullReplicaClient.query(new SolrQuery("*:*")).getResults().getNumFound();
                }
              });

      // Replica should still be active, since it doesn't disconnect from ZooKeeper
      assertNumberOfReplicas(numShards, 0, numShards, true, true);
      {
        long numFound = 0;
        TimeOut t = new TimeOut(REPLICATION_TIMEOUT_SECS, TimeUnit.SECONDS, TimeSource.NANO_TIME);
        while (numFound < 20 && !t.hasTimedOut()) {
          Thread.sleep(200);
          numFound =
              cluster
                  .getSolrClient()
                  .query(collectionName, new SolrQuery("*:*"))
                  .getResults()
                  .getNumFound();
        }
      }
    } finally {
      proxy.reopen();
    }

    try (SolrClient pullReplicaClient =
        getHttpSolrClient(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0).getCoreUrl())) {
      assertNumDocs(20, pullReplicaClient);
    }
  }

  public void testCantConnectToLeader() throws Exception {
    int numShards = 1;
    CollectionAdminRequest.createCollection(collectionName, "conf", numShards, 1, 0, 1)
        .process(cluster.getSolrClient());
    cluster.waitForActiveCollection(collectionName, numShards, numShards * 2);
    addDocs(10);
    DocCollection docCollection = assertNumberOfReplicas(numShards, 0, numShards, false, true);
    Slice s = docCollection.getSlices().iterator().next();
    SocketProxy proxy = getProxyForReplica(s.getLeader());
    try {
      // wait for replication
      try (SolrClient pullReplicaClient =
          getHttpSolrClient(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0).getCoreUrl())) {
        assertNumDocs(10, pullReplicaClient);
      }
      proxy.close();
      expectThrows(SolrException.class, () -> addDocs(1));
      try (SolrClient pullReplicaClient =
          getHttpSolrClient(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0).getCoreUrl())) {
        assertNumDocs(10, pullReplicaClient);
      }
      assertNumDocs(10, cluster.getSolrClient(collectionName));
    } finally {
      log.info("Opening leader node");
      proxy.reopen();
    }
    //     Back to normal
    //    Even if the leader is back to normal, the replica can get broken pipe for some time when
    // trying to connect to it. The commit
    //    can fail if it's sent to the replica, and it forwards it to the leader, and since it uses
    // CUSC the error is hidden! That breaks
    //    the last part of this test.
    //    addDocs(20);
    //    assertNumDocs(20, cluster.getSolrClient(), 300);
    //    try (SolrClient pullReplicaClient =
    // getHttpSolrClient(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0).getCoreUrl())) {
    //      assertNumDocs(20, pullReplicaClient);
    //    }
  }

  public void testPullReplicaDisconnectsFromZooKeeper() throws Exception {
    int numShards = 1;
    CollectionAdminRequest.createCollection(collectionName, "conf", numShards, 1, 0, 1)
        .process(cluster.getSolrClient());
    addDocs(10);
    DocCollection docCollection = assertNumberOfReplicas(numShards, 0, numShards, false, true);
    Slice s = docCollection.getSlices().iterator().next();
    try (SolrClient pullReplicaClient =
        getHttpSolrClient(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0).getCoreUrl())) {
      assertNumDocs(10, pullReplicaClient);
    }
    addDocs(20);
    JettySolrRunner jetty = getJettyForReplica(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0));
    cluster.expireZkSession(jetty);
    addDocs(30);
    waitForState("Expecting node to be disconnected", collectionName, activeReplicaCount(1, 0, 0));
    addDocs(40);
    waitForState("Expecting node to be reconnected", collectionName, activeReplicaCount(1, 0, 1));
    try (SolrClient pullReplicaClient =
        getHttpSolrClient(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0).getCoreUrl())) {
      assertNumDocs(40, pullReplicaClient);
    }
  }

  public void testCloseHooksDeletedOnReconnect() throws Exception {
    CollectionAdminRequest.createCollection(collectionName, "conf", 1, 1, 0, 1)
        .process(cluster.getSolrClient());
    addDocs(10);

    DocCollection docCollection = assertNumberOfReplicas(1, 0, 1, false, true);
    Slice s = docCollection.getSlices().iterator().next();
    JettySolrRunner jetty = getJettyForReplica(s.getReplicas(EnumSet.of(Replica.Type.PULL)).get(0));
    SolrCore core = jetty.getCoreContainer().getCores().iterator().next();

    for (int i = 0; i < (TEST_NIGHTLY ? 5 : 2); i++) {
      cluster.expireZkSession(jetty);
      waitForState(
          "Expecting node to be disconnected", collectionName, activeReplicaCount(1, 0, 0));
      waitForState("Expecting node to reconnect", collectionName, activeReplicaCount(1, 0, 1));
      // We have two active ReplicationHandler with two close hooks each, one for triggering
      // recovery and one for doing interval polling
      assertEquals(5, core.getCloseHooks().size());
    }
  }

  private void assertNumDocs(int numDocs, SolrClient client, int timeoutSecs)
      throws InterruptedException, SolrServerException, IOException {
    TimeOut t = new TimeOut(timeoutSecs, TimeUnit.SECONDS, TimeSource.NANO_TIME);
    long numFound = -1;
    while (!t.hasTimedOut()) {
      Thread.sleep(200);
      // if client is an HttpSolrClient, then the collection is in the path
      // otherwise with a CloudSolrClient then it's the defaultCollection
      numFound = client.query(new SolrQuery("*:*")).getResults().getNumFound();
      if (numFound == numDocs) {
        return;
      }
    }
    fail("Didn't get expected doc count. Expected: " + numDocs + ", Found: " + numFound);
  }

  private void assertNumDocs(int numDocs, SolrClient client)
      throws InterruptedException, SolrServerException, IOException {
    assertNumDocs(numDocs, client, REPLICATION_TIMEOUT_SECS);
  }

  private void addDocs(int numDocs) throws SolrServerException, IOException {
    List<SolrInputDocument> docs = new ArrayList<>(numDocs);
    for (int i = 0; i < numDocs; i++) {
      docs.add(new SolrInputDocument("id", String.valueOf(i), "fieldName_s", String.valueOf(i)));
    }
    cluster.getSolrClient().add(collectionName, docs);
    cluster.getSolrClient().commit(collectionName);
  }

  private DocCollection assertNumberOfReplicas(
      int numWriter, int numActive, int numPassive, boolean updateCollection, boolean activeOnly)
      throws KeeperException, InterruptedException {
    if (updateCollection) {
      cluster.getZkStateReader().forceUpdateCollection(collectionName);
    }
    DocCollection docCollection = getCollectionState(collectionName);
    assertNotNull(docCollection);
    assertEquals(
        "Unexpected number of writer replicas: " + docCollection,
        numWriter,
        docCollection.getReplicas(EnumSet.of(Replica.Type.NRT)).stream()
            .filter(r -> !activeOnly || r.getState() == Replica.State.ACTIVE)
            .count());
    assertEquals(
        "Unexpected number of pull replicas: " + docCollection,
        numPassive,
        docCollection.getReplicas(EnumSet.of(Replica.Type.PULL)).stream()
            .filter(r -> !activeOnly || r.getState() == Replica.State.ACTIVE)
            .count());
    assertEquals(
        "Unexpected number of active replicas: " + docCollection,
        numActive,
        docCollection.getReplicas(EnumSet.of(Replica.Type.TLOG)).stream()
            .filter(r -> !activeOnly || r.getState() == Replica.State.ACTIVE)
            .count());
    return docCollection;
  }

  protected JettySolrRunner getJettyForReplica(Replica replica) throws Exception {
    String replicaBaseUrl = replica.getStr(ZkStateReader.BASE_URL_PROP);
    assertNotNull(replicaBaseUrl);
    URL baseUrl = new URL(replicaBaseUrl);

    JettySolrRunner proxy = jettys.get(baseUrl.toURI());
    assertNotNull("No proxy found for " + baseUrl + "!", proxy);
    return proxy;
  }

  protected SocketProxy getProxyForReplica(Replica replica) throws Exception {
    String replicaBaseUrl = replica.getStr(ZkStateReader.BASE_URL_PROP);
    assertNotNull(replicaBaseUrl);
    URL baseUrl = new URL(replicaBaseUrl);

    SocketProxy proxy = proxies.get(baseUrl.toURI());
    if (proxy == null && !baseUrl.toExternalForm().endsWith("/")) {
      baseUrl = new URL(baseUrl.toExternalForm() + "/");
      proxy = proxies.get(baseUrl.toURI());
    }
    assertNotNull("No proxy found for " + baseUrl + "!", proxy);
    return proxy;
  }

  private void waitForDeletion(String collection) throws InterruptedException, KeeperException {
    TimeOut t = new TimeOut(10, TimeUnit.SECONDS, TimeSource.NANO_TIME);
    while (cluster.getSolrClient().getClusterState().hasCollection(collection)) {
      log.info("Collection not yet deleted");
      try {
        Thread.sleep(100);
        if (t.hasTimedOut()) {
          fail("Timed out waiting for collection " + collection + " to be deleted.");
        }
        cluster.getZkStateReader().forceUpdateCollection(collection);
      } catch (SolrException e) {
        return;
      }
    }
  }

  private CollectionStatePredicate activeReplicaCount(
      int numWriter, int numActive, int numPassive) {
    return (liveNodes, collectionState) -> {
      int writersFound = 0, activesFound = 0, passivesFound = 0;
      if (collectionState == null) return false;
      for (Slice slice : collectionState) {
        for (Replica replica : slice) {
          if (replica.isActive(liveNodes))
            switch (replica.getType()) {
              case TLOG:
                activesFound++;
                break;
              case PULL:
                passivesFound++;
                break;
              case NRT:
                writersFound++;
                break;
              default:
                throw new AssertionError("Unexpected replica type");
            }
        }
      }
      return numWriter == writersFound && numActive == activesFound && numPassive == passivesFound;
    };
  }
}
