/*
 * 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
 *
 *   https://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.accumulo.server.util;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.stream.Collectors;

import org.apache.accumulo.core.Constants;
import org.apache.accumulo.core.clientImpl.ClientContext;
import org.apache.accumulo.core.data.TableId;
import org.apache.accumulo.core.dataImpl.KeyExtent;
import org.apache.accumulo.core.metadata.StoredTabletFile;
import org.apache.accumulo.core.metadata.SystemTables;
import org.apache.accumulo.core.metadata.schema.Ample;
import org.apache.accumulo.core.metadata.schema.DataFileValue;
import org.apache.accumulo.core.metadata.schema.TabletMetadata;
import org.apache.accumulo.core.metadata.schema.TabletsMetadata;
import org.easymock.EasyMock;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;

public class TableDiskUsageTest {

  private static final String volume1 = "hdfs://nn1/acc";
  private static final String volume2 = "hdfs://nn2/acc";

  private static final TableId tableId1 = TableId.of("1");
  private static final TableId tableId2 = TableId.of("2");
  private static final TableId tableId3 = TableId.of("3");
  private static final String tabletName1 = "t-0001";
  private static final String tabletName2 = "t-0002";
  private static final String tabletName3 = "t-0003";
  private static final String tabletName4 = "t-0004";

  private static final TreeMap<TableId,String> tableIdToNameMap = new TreeMap<>();

  @BeforeAll
  public static void beforeClass() {
    tableIdToNameMap.put(SystemTables.ROOT.tableId(), SystemTables.METADATA.tableName());
    tableIdToNameMap.put(SystemTables.METADATA.tableId(), SystemTables.METADATA.tableName());
    tableIdToNameMap.put(tableId1, "table1");
    tableIdToNameMap.put(tableId2, "table2");
    tableIdToNameMap.put(tableId3, "table3");
  }

  @Test
  public void testSingleTableMultipleTablets() throws Exception {
    final ClientContext client = EasyMock.createMock(ClientContext.class);
    EasyMock.expect(client.createTableIdToQualifiedNameMap()).andReturn(tableIdToNameMap);
    final TabletsMetadata mockTabletsMetadata = mockTabletsMetadata(client, tableId1);

    List<TabletMetadata> realTabletsMetadata = new ArrayList<>();
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume1, tableId1, tabletName1, "C0001.rf"), 1024);
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume1, tableId1, tabletName1, "C0002.rf"), 1024);
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume1, tableId1, tabletName2, "C0003.rf"), 2048);
    mockTabletsMetadataIter(mockTabletsMetadata, realTabletsMetadata.iterator());

    EasyMock.replay(client, mockTabletsMetadata);

    Map<SortedSet<String>,Long> result = TableDiskUsage.getDiskUsage(tableSet(tableId1), client);

    assertEquals(4096, getTotalUsage(result, tableId1));
    assertEquals(1, result.size());
    Map.Entry<SortedSet<String>,Long> firstResult =
        result.entrySet().stream().findFirst().orElseThrow();
    assertEquals(1, firstResult.getKey().size());
    assertTrue(firstResult.getKey().contains(getTableName(tableId1)));
    assertEquals(4096, firstResult.getValue());

    EasyMock.verify(client, mockTabletsMetadata);
  }

  @Test
  public void testMultipleVolumes() throws Exception {
    final ClientContext client = EasyMock.createMock(ClientContext.class);
    EasyMock.expect(client.createTableIdToQualifiedNameMap()).andReturn(tableIdToNameMap);
    final TabletsMetadata mockTabletsMetadata = mockTabletsMetadata(client, tableId1);

    List<TabletMetadata> realTabletsMetadata = new ArrayList<>();
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume1, tableId1, tabletName1, "C0001.rf"), 1024);
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume1, tableId1, tabletName1, "C0002.rf"), 1024);
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume2, tableId1, tabletName2, "C0003.rf"), 2048);
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume2, tableId1, tabletName2, "C0004.rf"), 10000);
    mockTabletsMetadataIter(mockTabletsMetadata, realTabletsMetadata.iterator());

    EasyMock.replay(client, mockTabletsMetadata);

    Map<SortedSet<String>,Long> result = TableDiskUsage.getDiskUsage(tableSet(tableId1), client);

    assertEquals(14096, getTotalUsage(result, tableId1));
    assertEquals(1, result.size());
    Map.Entry<SortedSet<String>,Long> firstResult =
        result.entrySet().stream().findFirst().orElseThrow();
    assertEquals(1, firstResult.getKey().size());
    assertEquals(14096, firstResult.getValue());

    EasyMock.verify(client, mockTabletsMetadata);
  }

  @Test
  public void testMetadataTable() throws Exception {
    final ClientContext client = EasyMock.createMock(ClientContext.class);
    EasyMock.expect(client.createTableIdToQualifiedNameMap()).andReturn(tableIdToNameMap);
    final TabletsMetadata mockTabletsMetadata =
        mockTabletsMetadata(client, SystemTables.METADATA.tableId());

    List<TabletMetadata> realTabletsMetadata = new ArrayList<>();
    appendFileMetadata(realTabletsMetadata, getTabletFile(volume1, SystemTables.METADATA.tableId(),
        SystemTables.METADATA.tableName(), "C0001.rf"), 1024);
    mockTabletsMetadataIter(mockTabletsMetadata, realTabletsMetadata.iterator());

    EasyMock.replay(client, mockTabletsMetadata);

    Map<SortedSet<String>,Long> result =
        TableDiskUsage.getDiskUsage(tableSet(SystemTables.METADATA.tableId()), client);

    assertEquals(1024, getTotalUsage(result, SystemTables.METADATA.tableId()));
    assertEquals(1, result.size());
    Map.Entry<SortedSet<String>,Long> firstResult =
        result.entrySet().stream().findFirst().orElseThrow();
    assertEquals(1024, firstResult.getValue());

    EasyMock.verify(client, mockTabletsMetadata);
  }

  @Test
  public void testDuplicateFile() throws Exception {
    final ClientContext client = EasyMock.createMock(ClientContext.class);
    EasyMock.expect(client.createTableIdToQualifiedNameMap()).andReturn(tableIdToNameMap);
    final TabletsMetadata mockTabletsMetadata = mockTabletsMetadata(client, tableId1);

    List<TabletMetadata> realTabletsMetadata = new ArrayList<>();
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume1, tableId1, tabletName1, "C0001.rf"), 1024);
    appendFileMetadata(realTabletsMetadata,
        getTabletFile(volume1, tableId1, tabletName1, "C0001.rf"), 1024);
    mockTabletsMetadataIter(mockTabletsMetadata, realTabletsMetadata.iterator());

    EasyMock.replay(client, mockTabletsMetadata);

    Map<SortedSet<String>,Long> result = TableDiskUsage.getDiskUsage(tableSet(tableId1), client);

    assertEquals(1024, getTotalUsage(result, tableId1));
    assertEquals(1, result.size());
    Map.Entry<SortedSet<String>,Long> firstResult =
        result.entrySet().stream().findFirst().orElseThrow();
    assertEquals(1, firstResult.getKey().size());
    assertTrue(firstResult.getKey().contains(getTableName(tableId1)));
    assertEquals(1024, firstResult.getValue());

    EasyMock.verify(client, mockTabletsMetadata);
  }

  @Test
  public void testEmptyTable() throws Exception {
    final ClientContext client = EasyMock.createMock(ClientContext.class);
    EasyMock.expect(client.createTableIdToQualifiedNameMap()).andReturn(tableIdToNameMap);
    final TabletsMetadata mockTabletsMetadata = mockTabletsMetadata(client, tableId1);

    List<TabletMetadata> realTabletsMetadata = new ArrayList<>();
    mockTabletsMetadataIter(mockTabletsMetadata, realTabletsMetadata.iterator());

    EasyMock.replay(client, mockTabletsMetadata);

    Map<SortedSet<String>,Long> result = TableDiskUsage.getDiskUsage(tableSet(tableId1), client);

    assertEquals(0, getTotalUsage(result, tableId1));
    assertEquals(1, result.size());
    Map.Entry<SortedSet<String>,Long> firstResult =
        result.entrySet().stream().findFirst().orElseThrow();
    assertEquals(1, firstResult.getKey().size());
    assertEquals(0, firstResult.getValue());

    EasyMock.verify(client, mockTabletsMetadata);
  }

  @Test
  public void testMultipleTables() throws Exception {
    final ClientContext client = EasyMock.createMock(ClientContext.class);
    EasyMock.expect(client.createTableIdToQualifiedNameMap()).andReturn(tableIdToNameMap);

    final TabletsMetadata mockTabletsMetadata1 = mockTabletsMetadata(client, tableId1);
    List<TabletMetadata> realTabletsMetadata1 = new ArrayList<>();
    appendFileMetadata(realTabletsMetadata1,
        getTabletFile(volume1, tableId1, tabletName1, "C0001.rf"), 1024);
    appendFileMetadata(realTabletsMetadata1,
        getTabletFile(volume1, tableId1, tabletName1, "C0002.rf"), 4096);
    mockTabletsMetadataIter(mockTabletsMetadata1, realTabletsMetadata1.iterator());

    final TabletsMetadata mockTabletsMetadata2 = mockTabletsMetadata(client, tableId2);
    List<TabletMetadata> realTabletsMetadata2 = new ArrayList<>();
    appendFileMetadata(realTabletsMetadata2,
        getTabletFile(volume1, tableId2, tabletName2, "C0003.rf"), 2048);
    appendFileMetadata(realTabletsMetadata2,
        getTabletFile(volume1, tableId2, tabletName2, "C0004.rf"), 3000);
    mockTabletsMetadataIter(mockTabletsMetadata2, realTabletsMetadata2.iterator());

    final TabletsMetadata mockTabletsMetadata3 = mockTabletsMetadata(client, tableId3);
    List<TabletMetadata> realTabletsMetadata3 = new ArrayList<>();
    // shared file
    appendFileMetadata(realTabletsMetadata3,
        getTabletFile(volume1, tableId2, tabletName2, "C0003.rf"), 2048);
    appendFileMetadata(realTabletsMetadata3,
        getTabletFile(volume1, tableId3, tabletName3, "C0005.rf"), 84520);
    appendFileMetadata(realTabletsMetadata3,
        getTabletFile(volume1, tableId3, tabletName3, "C0006.rf"), 3000);
    appendFileMetadata(realTabletsMetadata3,
        getTabletFile(volume1, tableId3, tabletName4, "C0007.rf"), 98456);
    mockTabletsMetadataIter(mockTabletsMetadata3, realTabletsMetadata3.iterator());

    EasyMock.replay(client, mockTabletsMetadata1, mockTabletsMetadata2, mockTabletsMetadata3);

    Map<SortedSet<String>,Long> result =
        TableDiskUsage.getDiskUsage(tableSet(tableId1, tableId2, tableId3), client);

    assertEquals(5120, getTotalUsage(result, tableId1));
    assertEquals(5048, getTotalUsage(result, tableId2));
    assertEquals(188024, getTotalUsage(result, tableId3));

    // Make sure all shared tables exist in map
    assertEquals(4, result.size());
    assertTrue(result.containsKey(tableNameSet(tableId1)));
    assertTrue(result.containsKey(tableNameSet(tableId2)));
    assertTrue(result.containsKey(tableNameSet(tableId2, tableId3)));
    assertTrue(result.containsKey(tableNameSet(tableId3)));

    // Make sure all the shared disk usage computations are correct
    assertEquals(5120, result.get(tableNameSet(tableId1)));
    assertEquals(3000, result.get(tableNameSet(tableId2)));
    assertEquals(2048, result.get(tableNameSet(tableId2, tableId3)));
    assertEquals(185976, result.get(tableNameSet(tableId3)));

    EasyMock.verify(client, mockTabletsMetadata1, mockTabletsMetadata2, mockTabletsMetadata3);
  }

  private static TreeSet<String> tableNameSet(TableId... tableIds) {
    return Set.of(tableIds).stream().map(TableDiskUsageTest::getTableName)
        .collect(Collectors.toCollection(TreeSet::new));
  }

  // Need to use a LinkedHashSet for predictable order due to the fact that
  // we are using mock scanners that always return results in the same order
  private static Set<TableId> tableSet(TableId... tableIds) {
    return new LinkedHashSet<>(List.of(tableIds));
  }

  private static Long getTotalUsage(Map<SortedSet<String>,Long> result, TableId tableId) {
    return result.entrySet().stream()
        .filter(entry -> entry.getKey().contains(getTableName(tableId))).mapToLong(Entry::getValue)
        .sum();
  }

  private static String getTableName(TableId tableId) {
    return tableIdToNameMap.get(tableId);
  }

  private static void appendFileMetadata(List<TabletMetadata> realTabletsMetadata,
      StoredTabletFile file, long size) {
    TabletMetadata tm = TabletMetadata.builder(new KeyExtent(file.getTableId(), null, null))
        .putFile(file, new DataFileValue(size, 1)).build();
    realTabletsMetadata.add(tm);
  }

  private static StoredTabletFile getTabletFile(String volume, TableId tableId, String tablet,
      String fileName) {
    return new StoredTabletFile(StoredTabletFile.serialize(
        volume + Constants.HDFS_TABLES_DIR + "/" + tableId + "/" + tablet + "/" + fileName));
  }

  private TabletsMetadata mockTabletsMetadata(ClientContext client, TableId tableId) {
    final Ample ample = EasyMock.createMock(Ample.class);
    final TabletsMetadata.TableOptions tableOptions =
        EasyMock.createMock(TabletsMetadata.TableOptions.class);
    final TabletsMetadata.TableRangeOptions tableRangeOptions =
        EasyMock.createMock(TabletsMetadata.TableRangeOptions.class);
    final TabletsMetadata.Options options = EasyMock.createMock(TabletsMetadata.Options.class);
    final TabletsMetadata tabletsMetadata = EasyMock.createMock(TabletsMetadata.class);
    EasyMock.expect(client.getAmple()).andReturn(ample);
    EasyMock.expect(ample.readTablets()).andReturn(tableOptions);
    EasyMock.expect(tableOptions.forTable(tableId)).andReturn(tableRangeOptions);
    EasyMock.expect(tableRangeOptions.fetch(TabletMetadata.ColumnType.FILES)).andReturn(options);
    EasyMock.expect(options.build()).andReturn(tabletsMetadata);
    EasyMock.replay(ample, tableOptions, tableRangeOptions, options);
    return tabletsMetadata;
  }

  private void mockTabletsMetadataIter(TabletsMetadata tabletsMetadata,
      Iterator<TabletMetadata> realTabletsMetadata) {
    EasyMock.expect(tabletsMetadata.iterator()).andReturn(realTabletsMetadata);
    tabletsMetadata.close();
    EasyMock.expectLastCall().andAnswer(() -> null);
  }
}
