/*
 * 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.hadoop.ozone.debug.om;

import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertEquals;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import org.apache.hadoop.hdds.client.BlockID;
import org.apache.hadoop.hdds.client.StandaloneReplicationConfig;
import org.apache.hadoop.hdds.conf.OzoneConfiguration;
import org.apache.hadoop.hdds.protocol.proto.HddsProtos;
import org.apache.hadoop.ozone.debug.OzoneDebug;
import org.apache.hadoop.ozone.om.OMMetadataManager;
import org.apache.hadoop.ozone.om.OmMetadataManagerImpl;
import org.apache.hadoop.ozone.om.helpers.BucketLayout;
import org.apache.hadoop.ozone.om.helpers.OmBucketInfo;
import org.apache.hadoop.ozone.om.helpers.OmDirectoryInfo;
import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfo;
import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfoGroup;
import org.apache.hadoop.ozone.om.helpers.OmMultipartKeyInfo;
import org.apache.hadoop.ozone.om.helpers.OmVolumeArgs;
import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.KeyInfo;
import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.PartKeyInfo;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import picocli.CommandLine;

/**
 * Unit test for ContainerToKeyMapping tool.
 */
public class TestContainerToKeyMapping {

  @TempDir
  private Path tempDir;

  private OMMetadataManager omMetadataManager;
  private String dbPath;
  private CommandLine cmd;
  private StringWriter outWriter;

  private static final String VOLUME_NAME = "vol1";
  private static final String FSO_BUCKET_NAME = "fso-bucket";
  private static final String OBS_BUCKET_NAME = "obs-bucket";
  private static final long VOLUME_ID = 100L;
  private static final long FSO_BUCKET_ID = 200L;
  private static final long OBS_BUCKET_ID = 250L;
  private static final long DIR_ID = 300L;
  private static final long FILE_ID = 400L;
  private static final long KEY_ID = 450L;
  private static final long CONTAINER_ID_1 = 1L;
  private static final long CONTAINER_ID_2 = 2L;
  private static final long CONTAINER_ID_3 = 3L;
  private static final long CONTAINER_ID_4 = 4L;
  private static final long UNREFERENCED_FILE_ID = 500L;
  private static final long MISSING_DIR_ID = 999L;  // Non-existent parent
  private static final long OPEN_FILE_ID = 600L;
  private static final long OPEN_KEY_ID = 610L;
  private static final long MPU_KEY_ID = 700L;
  private static final long MPU_PART1_ID = 710L;
  private static final long MPU_PART2_ID = 720L;

  @BeforeEach
  public void setup() throws Exception {
    Path dbFile = tempDir.resolve("om.db");
    dbPath = dbFile.toString();
    
    cmd = new OzoneDebug().getCmd();
    outWriter = new StringWriter();
    StringWriter errWriter = new StringWriter();
    cmd.setOut(new PrintWriter(outWriter));
    cmd.setErr(new PrintWriter(errWriter));
    
    OzoneConfiguration conf = new OzoneConfiguration();
    omMetadataManager = new OmMetadataManagerImpl(conf,
        tempDir.toFile(), "om.db", false);

    createTestData();
    
    omMetadataManager.getStore().close();
  }

  @AfterEach
  public void tearDown() throws Exception {
    if (omMetadataManager != null && !omMetadataManager.getStore().isClosed()) {
      omMetadataManager.getStore().close();
    }
  }

  @Test
  public void testContainerToKeyMapping() {
    int exitCode = execute("--containers", CONTAINER_ID_1 + "," + CONTAINER_ID_2);
    assertEquals(0, exitCode);

    String output = outWriter.toString();

    // Check FSO key
    assertThat(output).contains("\"" + CONTAINER_ID_1 + "\"");
    assertThat(output).contains("/vol1/fso-bucket/dir1/file1");

    // Check OBS key
    assertThat(output).contains("\"" + CONTAINER_ID_2 + "\"");
    assertThat(output).contains("/vol1/obs-bucket/key1");
  }

  @Test
  public void testContainerToKeyMappingWithOnlyFileNames() {
    int exitCode = execute("--containers", CONTAINER_ID_1 + "," + CONTAINER_ID_2, "--onlyFileNames");
    assertEquals(0, exitCode);

    String output = outWriter.toString();

    // Check FSO key - should show only filename
    assertThat(output).contains("\"" + CONTAINER_ID_1 + "\"");
    assertThat(output).contains("file1");
    assertThat(output).doesNotContain("/vol1/fso-bucket/dir1/file1");

    // Check OBS key - should also show only key name
    assertThat(output).contains("\"" + CONTAINER_ID_2 + "\"");
    assertThat(output).contains("key1");
    assertThat(output).doesNotContain("/vol1/obs-bucket/key1");
  }

  @Test
  public void testContainerToKeyMappingWithOpenKeys() {
    int exitCode = execute("--containers", CONTAINER_ID_1 + "," + CONTAINER_ID_2, "--in-progress");
    assertEquals(0, exitCode);

    String output = outWriter.toString();

    // Check open FSO file in openKeys
    assertThat(output).contains("\"" + CONTAINER_ID_1 + "\"");
    assertThat(output).contains("\"openKeys\"");
    assertThat(output).contains("/100/200/300/openFile/1");

    // Check open OBS key in openKeys
    assertThat(output).contains("\"" + CONTAINER_ID_2 + "\"");
    assertThat(output).contains("\"openKeys\"");
    assertThat(output).contains("/vol1/obs-bucket/openKey");
  }

  @Test
  public void testContainerToKeyMappingWithMPU() {
    int exitCode = execute("--containers", String.valueOf(CONTAINER_ID_4), "--in-progress");
    assertEquals(0, exitCode);

    String output = outWriter.toString();

    // Check MPU parts in openKeys
    assertThat(output).contains("\"" + CONTAINER_ID_4 + "\"");
    assertThat(output).contains("\"openKeys\"");
    assertThat(output).contains("/vol1/obs-bucket/mpuKey/test-upload-id");
  }

  @Test
  public void testContainerToKeyMappingWithMPUOnlyFileNames() {
    int exitCode = execute("--containers", String.valueOf(CONTAINER_ID_4), "--in-progress", "--onlyFileNames");
    assertEquals(0, exitCode);

    String output = outWriter.toString();

    // Open keys and MPU always show table key (not affected by --onlyFileNames)
    assertThat(output).contains("\"" + CONTAINER_ID_4 + "\"");
    assertThat(output).contains("\"openKeys\"");
    assertThat(output).contains("/vol1/obs-bucket/mpuKey/test-upload-id");
  }

  @Test

  public void testNonExistentContainer() {
    long nonExistentContainerId = 999L;
    
    int exitCode = execute("--containers", String.valueOf(nonExistentContainerId));
    assertEquals(0, exitCode);
    
    String output = outWriter.toString();
    assertThat(output).contains("\"" + nonExistentContainerId + "\"");
    assertThat(output).contains("\"totalKeys\" : 0");
  }

  @Test
  public void testUnreferencedKeys() {
    int exitCode = execute("--containers", String.valueOf(CONTAINER_ID_3));
    assertEquals(0, exitCode);
    
    String output = outWriter.toString();

    assertThat(output).contains("\"" + CONTAINER_ID_3 + "\"");
    assertThat(output).contains("\"totalKeys\" : 1");
    assertThat(output).contains("\"unreferencedKeys\" : 1");
  }

  private void createTestData() throws Exception {
    // Create volume
    OmVolumeArgs volumeArgs = OmVolumeArgs.newBuilder()
        .setVolume(VOLUME_NAME)
        .setOwnerName("testUser")
        .setAdminName("admin")
        .setObjectID(VOLUME_ID)
        .build();
    omMetadataManager.getVolumeTable().put(
        omMetadataManager.getVolumeKey(VOLUME_NAME), volumeArgs);

    // Create FSO bucket
    OmBucketInfo bucketInfo = OmBucketInfo.newBuilder()
        .setVolumeName(VOLUME_NAME)
        .setBucketName(FSO_BUCKET_NAME)
        .setBucketLayout(BucketLayout.FILE_SYSTEM_OPTIMIZED)
        .setObjectID(FSO_BUCKET_ID)
        .build();
    omMetadataManager.getBucketTable().put(
        omMetadataManager.getBucketKey(VOLUME_NAME, FSO_BUCKET_NAME), bucketInfo);

    // Create OBS bucket
    OmBucketInfo obsBucketInfo = OmBucketInfo.newBuilder()
        .setVolumeName(VOLUME_NAME)
        .setBucketName(OBS_BUCKET_NAME)
        .setBucketLayout(BucketLayout.OBJECT_STORE)
        .setObjectID(OBS_BUCKET_ID)
        .build();
    omMetadataManager.getBucketTable().put(
        omMetadataManager.getBucketKey(VOLUME_NAME, OBS_BUCKET_NAME), obsBucketInfo);

    // Create directory
    OmDirectoryInfo dirInfo = OmDirectoryInfo.newBuilder()
        .setName("dir1")
        .setObjectID(DIR_ID)
        .setParentObjectID(FSO_BUCKET_ID)
        .setUpdateID(1)
        .build();
    String dirKey = omMetadataManager.getOzonePathKey(VOLUME_ID, FSO_BUCKET_ID, FSO_BUCKET_ID, "dir1");
    omMetadataManager.getDirectoryTable().put(dirKey, dirInfo);

    // Create FSO file with a block in container 1
    OmKeyInfo keyInfo = createKeyInfo(
        "file1", FILE_ID, DIR_ID, CONTAINER_ID_1);
    String fileKey = omMetadataManager.getOzonePathKey(
        VOLUME_ID, FSO_BUCKET_ID, DIR_ID, "file1");
    omMetadataManager.getFileTable().put(fileKey, keyInfo);

    // Create OBS key with a block in container 2
    OmKeyInfo obsKeyInfo = createOBSKeyInfo(
        "key1", KEY_ID, CONTAINER_ID_2);
    String obsKey = omMetadataManager.getOzoneKey(
        VOLUME_NAME, OBS_BUCKET_NAME, "key1");
    omMetadataManager.getKeyTable(BucketLayout.OBJECT_STORE).put(obsKey, obsKeyInfo);

    // Create unreferenced file (parent directory doesn't exist)
    OmKeyInfo unreferencedKey = createKeyInfo(
        "unreferencedFile", UNREFERENCED_FILE_ID, MISSING_DIR_ID, CONTAINER_ID_3);
    String unreferencedFileKey = omMetadataManager.getOzonePathKey(
        VOLUME_ID, FSO_BUCKET_ID, MISSING_DIR_ID, "unreferencedFile");
    omMetadataManager.getFileTable().put(unreferencedFileKey, unreferencedKey);

    // Create open FSO file with a block in container 1
    OmKeyInfo openFileInfo = createKeyInfo(
        "openFile", OPEN_FILE_ID, DIR_ID, CONTAINER_ID_1);
    String openFileKey = omMetadataManager.getOpenFileName(
        VOLUME_ID, FSO_BUCKET_ID, DIR_ID, "openFile", 1L);
    omMetadataManager.getOpenKeyTable(BucketLayout.FILE_SYSTEM_OPTIMIZED).put(openFileKey, openFileInfo);

    // Create open OBS key with a block in container 2
    OmKeyInfo openKeyInfo = createOBSKeyInfo(
        "openKey", OPEN_KEY_ID, CONTAINER_ID_2);
    String openKey = omMetadataManager.getOzoneKey(
        VOLUME_NAME, OBS_BUCKET_NAME, "openKey");
    omMetadataManager.getOpenKeyTable(BucketLayout.OBJECT_STORE).put(openKey, openKeyInfo);

    // Create MPU (multipart upload) for OBS bucket with parts in container 5
    createMultipartUpload();
  }

  /**
   * Helper method to create a multipart upload with parts.
   */
  private void createMultipartUpload() throws Exception {
    String mpuKeyName = "mpuKey";
    String uploadId = "test-upload-id";

    // Create part 1 with a block in container 4
    OmKeyInfo part1Info = createOBSKeyInfo(
        mpuKeyName + "/" + uploadId + "/part-1", MPU_PART1_ID, CONTAINER_ID_4);
    KeyInfo part1Proto = part1Info.getProtobuf(true, 0);
    PartKeyInfo partKeyInfo1 = PartKeyInfo.newBuilder()
        .setPartName(mpuKeyName + "/" + uploadId + "/part-1")
        .setPartNumber(1)
        .setPartKeyInfo(part1Proto)
        .build();

    // Create part 2 with a block in container 4
    OmKeyInfo part2Info = createOBSKeyInfo(
        mpuKeyName + "/" + uploadId + "/part-2", MPU_PART2_ID, CONTAINER_ID_4);
    KeyInfo part2Proto = part2Info.getProtobuf(true, 0);
    PartKeyInfo partKeyInfo2 = PartKeyInfo.newBuilder()
        .setPartName(mpuKeyName + "/" + uploadId + "/part-2")
        .setPartNumber(2)
        .setPartKeyInfo(part2Proto)
        .build();

    // Create OmMultipartKeyInfo - for OBS, parentID should be 0 (not the bucket ID)
    OmMultipartKeyInfo mpuInfo = new OmMultipartKeyInfo.Builder()
        .setUploadID(uploadId)
        .setCreationTime(System.currentTimeMillis())
        .setReplicationConfig(StandaloneReplicationConfig.getInstance(HddsProtos.ReplicationFactor.ONE))
        .addPartKeyInfoList(1, partKeyInfo1)
        .addPartKeyInfoList(2, partKeyInfo2)
        .setObjectID(MPU_KEY_ID)
        .setParentID(0)  // OBS keys have parentID = 0
        .setUpdateID(1)
        .build();

    // Put into multipartInfoTable
    String mpuKey = omMetadataManager.getMultipartKey(
        VOLUME_NAME, OBS_BUCKET_NAME, mpuKeyName, uploadId);
    omMetadataManager.getMultipartInfoTable().put(mpuKey, mpuInfo);
  }

  /**
   * Helper method to create OmKeyInfo with a block in specified container (FSO).
   */
  private OmKeyInfo createKeyInfo(String keyName, long objectId, long parentId, long containerId) {
    OmKeyLocationInfo locationInfo = new OmKeyLocationInfo.Builder()
        .setBlockID(new BlockID(containerId, 1L))
        .setLength(1024)
        .setOffset(0)
        .build();

    OmKeyLocationInfoGroup locationGroup = new OmKeyLocationInfoGroup(0,
        Collections.singletonList(locationInfo));

    return new OmKeyInfo.Builder()
        .setVolumeName(VOLUME_NAME)
        .setBucketName(FSO_BUCKET_NAME)
        .setKeyName(keyName)
        .setReplicationConfig(StandaloneReplicationConfig.getInstance(HddsProtos.ReplicationFactor.ONE))
        .setDataSize(1024)
        .setObjectID(objectId)
        .setParentObjectID(parentId)
        .setUpdateID(1)
        .addOmKeyLocationInfoGroup(locationGroup)
        .build();
  }

  /**
   * Helper method to create OmKeyInfo for OBS keys with a block in specified container.
   */
  private OmKeyInfo createOBSKeyInfo(String keyName, long objectId, long containerId) {
    OmKeyLocationInfo locationInfo = new OmKeyLocationInfo.Builder()
        .setBlockID(new BlockID(containerId, 1L))
        .setLength(1024)
        .setOffset(0)
        .build();

    OmKeyLocationInfoGroup locationGroup = new OmKeyLocationInfoGroup(0,
        Collections.singletonList(locationInfo));

    return new OmKeyInfo.Builder()
        .setVolumeName(VOLUME_NAME)
        .setBucketName(OBS_BUCKET_NAME)
        .setKeyName(keyName)
        .setReplicationConfig(StandaloneReplicationConfig.getInstance(HddsProtos.ReplicationFactor.ONE))
        .setDataSize(1024)
        .setObjectID(objectId)
        .setUpdateID(1)
        .addOmKeyLocationInfoGroup(locationGroup)
        .build();
  }

  private int execute(String... args) {
    List<String> argList = new ArrayList<>(Arrays.asList("om", "container-key-mapping", "--db", dbPath));
    argList.addAll(Arrays.asList(args));

    return cmd.execute(argList.toArray(new String[0]));
  }
}
