/*
 * 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.beam.sdk.extensions.sbe;

import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalTime;
import java.time.OffsetDateTime;
import java.time.OffsetTime;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.Schema.LogicalType;
import org.apache.beam.sdk.schemas.logicaltypes.PassThroughLogicalType;
import org.checkerframework.checker.initialization.qual.Initialized;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.checkerframework.checker.nullness.qual.UnknownKeyFor;

/**
 * Classes that represent various SBE semantic types.
 *
 * <p>Not all of SBE's semantic types are represented here, as some can be reasonably represented in
 * Beam schemas with just a primitive.
 */
@Experimental(Kind.SCHEMAS)
public final class SbeLogicalTypes {
  // Default argument type values
  private static final String DEFAULT_STRING_ARG = "";

  private SbeLogicalTypes() {}

  // Unsigned types are all stored at the next highest value. This prevents unexpected behavior
  // when reading and likely has negligible space impact.

  /** Represents SBE's uint8 type. */
  public static final class Uint8 extends PassThroughLogicalType<Short> {
    public static final String IDENTIFIER = "uint8";

    public Uint8() {
      super(IDENTIFIER, FieldType.STRING, DEFAULT_STRING_ARG, FieldType.INT16);
    }
  }

  /** Represents SBE's uint16 type. */
  public static final class Uint16 extends PassThroughLogicalType<Integer> {
    public static final String IDENTIFIER = "uint16";

    public Uint16() {
      super(IDENTIFIER, FieldType.STRING, DEFAULT_STRING_ARG, FieldType.INT32);
    }
  }

  /** Represents SBE's uint32 type. */
  public static final class Uint32 extends PassThroughLogicalType<Long> {
    public static final String IDENTIFIER = "uint32";

    public Uint32() {
      super(IDENTIFIER, FieldType.STRING, DEFAULT_STRING_ARG, FieldType.INT64);
    }
  }

  /** Represents SBE's uint64 type. */
  public static final class Uint64 extends PassThroughLogicalType<String> {
    // Unknown if anyone will ever use this as a BigInteger, so we're keeping it as a String for
    // now.

    public static final String IDENTIFIER = "uint64";

    public Uint64() {
      super(IDENTIFIER, FieldType.STRING, DEFAULT_STRING_ARG, FieldType.STRING);
    }
  }

  // SBE time-based composite and logical types.

  /** Represents SBE's UTCTimestamp composite type. */
  public static final class UTCTimestamp implements LogicalType<Instant, String> {
    public static final String IDENTIFIER = "UTCTimestamp";

    @Override
    public @UnknownKeyFor @NonNull @Initialized String getIdentifier() {
      return IDENTIFIER;
    }

    @Override
    public @Nullable @UnknownKeyFor @Initialized FieldType getArgumentType() {
      return null;
    }

    @Override
    public @UnknownKeyFor @NonNull @Initialized FieldType getBaseType() {
      return FieldType.INT64;
    }

    @Override
    public @NonNull String toBaseType(@NonNull Instant input) {
      return input.toString();
    }

    @Override
    public @NonNull Instant toInputType(@NonNull String base) {
      return Instant.parse(base);
    }
  }

  /** Represents SBE's UTCTimeOnly composite type. */
  public static final class UTCTimeOnly implements LogicalType<LocalTime, String> {
    public static final String IDENTIFIER = "UTCTimeOnly";

    @Override
    public @UnknownKeyFor @NonNull @Initialized String getIdentifier() {
      return IDENTIFIER;
    }

    @Override
    public @Nullable @UnknownKeyFor @Initialized FieldType getArgumentType() {
      return null;
    }

    @Override
    public @UnknownKeyFor @NonNull @Initialized FieldType getBaseType() {
      return FieldType.INT64;
    }

    @Override
    public @NonNull String toBaseType(@NonNull LocalTime input) {
      return input.toString();
    }

    @Override
    public @NonNull LocalTime toInputType(@NonNull String base) {
      return LocalTime.parse(base);
    }
  }

  /** Represents SBE's TZTimestamp composite type. */
  public static final class TZTimestamp implements LogicalType<OffsetDateTime, String> {
    public static final String IDENTIFIER = "TZTimestamp";

    @Override
    public @UnknownKeyFor @NonNull @Initialized String getIdentifier() {
      return IDENTIFIER;
    }

    @Override
    public @Nullable @UnknownKeyFor @Initialized FieldType getArgumentType() {
      return null;
    }

    @Override
    public @UnknownKeyFor @NonNull @Initialized FieldType getBaseType() {
      return FieldType.STRING;
    }

    @Override
    public @NonNull String toBaseType(@NonNull OffsetDateTime input) {
      return input.toString();
    }

    @Override
    public @NonNull OffsetDateTime toInputType(@NonNull String base) {
      return OffsetDateTime.parse(base);
    }
  }

  /** Represents SBE's TimeOnly composite type. */
  public static final class TZTimeOnly implements LogicalType<OffsetTime, String> {
    public static final String IDENTIFIER = "TZTimeOnly";

    @Override
    public @UnknownKeyFor @NonNull @Initialized String getIdentifier() {
      return IDENTIFIER;
    }

    @Override
    public @Nullable @UnknownKeyFor @Initialized FieldType getArgumentType() {
      return null;
    }

    @Override
    public @UnknownKeyFor @NonNull @Initialized FieldType getBaseType() {
      return FieldType.STRING;
    }

    @Override
    public @NonNull String toBaseType(@NonNull OffsetTime input) {
      return input.toString();
    }

    @Override
    public @NonNull OffsetTime toInputType(@NonNull String base) {
      return OffsetTime.parse(base);
    }
  }

  /** Helper type for SBE's date types. */
  private static class SbeDateType implements LogicalType<LocalDate, String> {
    private final String identifier;

    SbeDateType(String identifier) {
      this.identifier = identifier;
    }

    @Override
    public @UnknownKeyFor @NonNull @Initialized String getIdentifier() {
      return identifier;
    }

    @Override
    public @Nullable @UnknownKeyFor @Initialized FieldType getArgumentType() {
      return null;
    }

    @Override
    public @UnknownKeyFor @NonNull @Initialized FieldType getBaseType() {
      return FieldType.INT32;
    }

    @Override
    public @NonNull String toBaseType(@NonNull LocalDate input) {
      return input.toString();
    }

    @Override
    public @NonNull LocalDate toInputType(@NonNull String base) {
      return LocalDate.parse(base);
    }
  }

  /** Representation of SBE's UTCDateOnly. */
  public static final class UTCDateOnly extends SbeDateType {
    public static final String IDENTIFIER = "UTCDateOnly";

    public UTCDateOnly() {
      super(IDENTIFIER);
    }
  }

  /** Representation of SBE's LocalMktDate. */
  public static final class LocalMktDate extends SbeDateType {
    public static final String IDENTIFIER = "LocalMktDate";

    public LocalMktDate() {
      super(IDENTIFIER);
    }
  }
}
