/*
 * 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.iceberg.encryption;

import static org.apache.iceberg.types.Types.NestedField.optional;
import static org.apache.iceberg.types.Types.NestedField.required;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;
import java.util.Map;
import org.apache.avro.generic.IndexedRecord;
import org.apache.iceberg.Schema;
import org.apache.iceberg.avro.AvroSchemaUtil;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
import org.apache.iceberg.types.Types;

class StandardKeyMetadata implements NativeEncryptionKeyMetadata, IndexedRecord {
  private static final byte V1 = 1;
  private static final Schema SCHEMA_V1 =
      new Schema(
          required(0, "encryption_key", Types.BinaryType.get()),
          optional(1, "aad_prefix", Types.BinaryType.get()),
          optional(2, "file_length", Types.LongType.get()));
  private static final org.apache.avro.Schema AVRO_SCHEMA_V1 =
      AvroSchemaUtil.convert(SCHEMA_V1, StandardKeyMetadata.class.getCanonicalName());

  private static final Map<Byte, Schema> SCHEMA_VERSIONS = ImmutableMap.of(V1, SCHEMA_V1);
  private static final Map<Byte, org.apache.avro.Schema> AVRO_SCHEMA_VERSIONS =
      ImmutableMap.of(V1, AVRO_SCHEMA_V1);

  private static final KeyMetadataEncoder KEY_METADATA_ENCODER = new KeyMetadataEncoder(V1);
  private static final KeyMetadataDecoder KEY_METADATA_DECODER = new KeyMetadataDecoder(V1);

  private ByteBuffer encryptionKey;
  private ByteBuffer aadPrefix;
  private Long fileLength;

  /** Used by Avro reflection to instantiate this class * */
  StandardKeyMetadata() {}

  StandardKeyMetadata(byte[] key, byte[] aad) {
    this(key, aad, null);
  }

  StandardKeyMetadata(byte[] key, byte[] aad, Long fileLength) {
    this.encryptionKey = ByteBuffer.wrap(key);
    this.aadPrefix = ByteBuffer.wrap(aad);
    this.fileLength = fileLength;
  }

  /**
   * Copy constructor.
   *
   * @param toCopy a StandardKeymetadata to copy
   * @param fileLength file length that overrides toCopy if not null
   */
  private StandardKeyMetadata(StandardKeyMetadata toCopy, Long fileLength) {
    this.encryptionKey = toCopy.encryptionKey;
    this.aadPrefix = toCopy.aadPrefix;
    this.fileLength = fileLength != null ? fileLength : toCopy.fileLength;
  }

  static Map<Byte, Schema> supportedSchemaVersions() {
    return SCHEMA_VERSIONS;
  }

  static Map<Byte, org.apache.avro.Schema> supportedAvroSchemaVersions() {
    return AVRO_SCHEMA_VERSIONS;
  }

  @Override
  public ByteBuffer encryptionKey() {
    return encryptionKey;
  }

  @Override
  public ByteBuffer aadPrefix() {
    return aadPrefix;
  }

  @Override
  public Long fileLength() {
    return fileLength;
  }

  static StandardKeyMetadata castOrParse(EncryptionKeyMetadata keyMetadata) {
    if (keyMetadata instanceof StandardKeyMetadata) {
      return (StandardKeyMetadata) keyMetadata;
    }

    ByteBuffer kmBuffer = keyMetadata.buffer();

    if (kmBuffer == null) {
      throw new IllegalStateException("Null key metadata buffer");
    }

    return parse(kmBuffer);
  }

  static StandardKeyMetadata parse(ByteBuffer buffer) {
    try {
      return KEY_METADATA_DECODER.decode(buffer);
    } catch (IOException e) {
      throw new UncheckedIOException("Failed to parse envelope encryption metadata", e);
    }
  }

  @Override
  public ByteBuffer buffer() {
    try {
      return KEY_METADATA_ENCODER.encode(this);
    } catch (IOException e) {
      throw new UncheckedIOException("Failed to serialize envelope key metadata", e);
    }
  }

  @Override
  public EncryptionKeyMetadata copy() {
    return new StandardKeyMetadata(this, null);
  }

  @Override
  public NativeEncryptionKeyMetadata copyWithLength(long length) {
    return new StandardKeyMetadata(this, length);
  }

  @Override
  public void put(int i, Object v) {
    switch (i) {
      case 0:
        this.encryptionKey = (ByteBuffer) v;
        return;
      case 1:
        this.aadPrefix = (ByteBuffer) v;
        return;
      case 2:
        this.fileLength = (Long) v;
        return;
      default:
        // ignore the object, it must be from a newer version of the format
    }
  }

  @Override
  public Object get(int i) {
    switch (i) {
      case 0:
        return encryptionKey;
      case 1:
        return aadPrefix;
      case 2:
        return fileLength;
      default:
        throw new UnsupportedOperationException("Unknown field ordinal: " + i);
    }
  }

  @Override
  public org.apache.avro.Schema getSchema() {
    return AVRO_SCHEMA_V1;
  }
}
