// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.data.cosmos.serialization.hybridrow.layouts;

import com.azure.data.cosmos.core.Out;
import com.azure.data.cosmos.serialization.hybridrow.SchemaId;
import com.azure.data.cosmos.serialization.hybridrow.schemas.ArrayPropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.MapPropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.Namespace;
import com.azure.data.cosmos.serialization.hybridrow.schemas.ObjectPropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.PrimitivePropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.Property;
import com.azure.data.cosmos.serialization.hybridrow.schemas.PropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.Schema;
import com.azure.data.cosmos.serialization.hybridrow.schemas.ScopePropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.SetPropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.TaggedPropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.TuplePropertyType;
import com.azure.data.cosmos.serialization.hybridrow.schemas.TypeKind;
import com.azure.data.cosmos.serialization.hybridrow.schemas.UdtPropertyType;
import com.google.common.base.Strings;
import tangible.ListHelper;

import javax.annotation.Nonnull;
import java.util.List;
import java.util.Optional;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Strings.lenientFormat;

/**
 * Converts a logical schema into a physical layout.
 */
public final class LayoutCompiler {
    /**
     * Compiles a logical schema into a physical layout that can be used to read and write rows
     *
     * @param ns     The namespace within which {@code schema} is defined
     * @param schema The logical schema to produce a layout for
     * @return The layout for the schema
     */
    @Nonnull
    public static Layout compile(@Nonnull final Namespace ns, @Nonnull final Schema schema) {

        checkNotNull(ns, "expected non-null ns");
        checkNotNull(schema, "expected non-null schema");
        checkArgument(schema.type() == TypeKind.Schema);
        checkArgument(!Strings.isNullOrEmpty(schema.name()));
        checkArgument(ns.schemas().contains(schema));

        LayoutBuilder builder = new LayoutBuilder(schema.name(), schema.schemaId());
        LayoutCompiler.addProperties(builder, ns, LayoutCode.SCHEMA, schema.properties());

        return builder.build();
    }

    private static void addProperties(LayoutBuilder builder, Namespace ns, LayoutCode scope, List<Property> properties) {

        final Out<TypeArgumentList> typeArgs = new Out<>();

        for (Property p : properties) {

            LayoutType type = LayoutCompiler.logicalToPhysicalType(ns, p.propertyType(), typeArgs);

            switch (LayoutCodeTraits.clearImmutableBit(type.layoutCode())) {

                case OBJECT_SCOPE: {
                    if (!p.propertyType().nullable()) {
                        throw new LayoutCompilationException("Non-nullable sparse column are not supported.");
                    }
                    ObjectPropertyType op = (ObjectPropertyType)p.propertyType();
                    builder.addObjectScope(p.path(), type);
                    LayoutCompiler.addProperties(builder, ns, type.layoutCode(), op.properties());
                    builder.EndObjectScope();
                    break;
                }

                case ARRAY_SCOPE:
                case TYPED_ARRAY_SCOPE:
                case SET_SCOPE:
                case TYPED_SET_SCOPE:
                case MAP_SCOPE:
                case TYPED_MAP_SCOPE:
                case TUPLE_SCOPE:
                case TYPED_TUPLE_SCOPE:
                case TAGGED_SCOPE:
                case TAGGED2_SCOPE:
                case SCHEMA: {
                    if (!p.propertyType().nullable()) {
                        throw new LayoutCompilationException("Non-nullable sparse column are not supported.");
                    }
                    builder.addTypedScope(p.path(), type, typeArgs.get());
                    break;
                }

                case NULLABLE_SCOPE: {
                    throw new LayoutCompilationException("Nullables cannot be explicitly declared as columns.");
                }

                default: {

                    if (p.propertyType() instanceof PrimitivePropertyType) {

                        PrimitivePropertyType pp = (PrimitivePropertyType) p.propertyType();

                        switch (pp.storage()) {

                            case FIXED:
                                if (LayoutCodeTraits.clearImmutableBit(scope) != LayoutCode.SCHEMA) {
                                    throw new LayoutCompilationException(
                                        "Cannot have fixed storage within a sparse scope.");
                                }
                                if (type.isNull() && !pp.nullable()) {
                                    throw new LayoutCompilationException(
                                        "Non-nullable null columns are not supported.");
                                }
                                builder.addFixedColumn(p.path(), type, pp.nullable(), pp.length());
                                break;

                            case VARIABLE:
                                if (LayoutCodeTraits.clearImmutableBit(scope) != LayoutCode.SCHEMA) {
                                    throw new LayoutCompilationException(
                                        "Cannot have variable storage within a sparse scope.");
                                }
                                if (!pp.nullable()) {
                                    throw new LayoutCompilationException(
                                        "Non-nullable variable columns are not supported.");
                                }
                                builder.addVariableColumn(p.path(), type, pp.length());
                                break;

                            case SPARSE:
                                if (!pp.nullable()) {
                                    throw new LayoutCompilationException(
                                        "Non-nullable sparse columns are not supported.");
                                }
                                builder.addSparseColumn(p.path(), type);
                                break;

                            default:
                                throw new LayoutCompilationException(
                                    lenientFormat("Unknown storage specification: %s", pp.storage()));
                        }
                    } else {
                        throw new LayoutCompilationException(
                            lenientFormat("Unknown property type: %s", type.name()));
                    }

                    break;
                }
            }
        }
    }

    private static LayoutType logicalToPhysicalType(Namespace ns, PropertyType logicalType,
                                                    Out<TypeArgumentList> typeArgs) {

        typeArgs.set(TypeArgumentList.EMPTY);
        boolean immutable = logicalType instanceof ScopePropertyType && ((ScopePropertyType) logicalType).immutable();

        switch (logicalType.type()) {

            case Null:
                return LayoutTypes.NULL;

            case Boolean:
                return LayoutTypes.BOOLEAN;

            case Int8:
                return LayoutTypes.INT_8;

            case Int16:
                return LayoutTypes.INT_16;

            case Int32:
                return LayoutTypes.INT_32;

            case Int64:
                return LayoutTypes.INT_64;

            case UInt8:
                return LayoutTypes.UINT_8;

            case UInt16:
                return LayoutTypes.UINT_16;

            case UInt32:
                return LayoutTypes.UINT_32;

            case UInt64:
                return LayoutTypes.UINT_64;

            case Float32:
                return LayoutTypes.FLOAT_32;

            case Float64:
                return LayoutTypes.FLOAT_64;

            case Float128:
                return LayoutTypes.FLOAT_128;

            case Decimal:
                return LayoutTypes.DECIMAL;

            case DateTime:
                return LayoutTypes.DATE_TIME;

            case UnixDateTime:
                return LayoutTypes.UNIX_DATE_TIME;

            case Guid:
                return LayoutTypes.GUID;

            case MongoDbObjectId:
                throw new UnsupportedOperationException();
                // return LayoutTypes.MONGO_DB_OBJECT_ID;

            case Utf8:
                return LayoutTypes.UTF_8;

            case Binary:
                return LayoutTypes.BINARY;

            case VarInt:
                return LayoutTypes.VAR_INT;

            case VarUInt:
                return LayoutTypes.VAR_UINT;

            case Object:
                return immutable ? LayoutTypes.IMMUTABLE_OBJECT : LayoutTypes.OBJECT;

            case Array: {

                assert logicalType instanceof ArrayPropertyType;
                ArrayPropertyType ap = (ArrayPropertyType) logicalType;

                if (ap.items() != null && (ap.items().type() != TypeKind.Any)) {

                    final Out<TypeArgumentList> out = new Out<>();

                    LayoutType itemType = LayoutCompiler.logicalToPhysicalType(ns, ap.items(), out);
                    TypeArgumentList itemTypeArgs = out.get();

                    if (ap.items().nullable()) {
                        itemTypeArgs = new TypeArgumentList(new TypeArgument(itemType, itemTypeArgs));
                        itemType = itemType.isImmutable() ? LayoutTypes.IMMUTABLE_NULLABLE : LayoutTypes.NULLABLE;
                    }

                    typeArgs.set(new TypeArgumentList(new TypeArgument(itemType, itemTypeArgs)));
                    return immutable ? LayoutTypes.IMMUTABLE_TYPED_ARRAY : LayoutTypes.TYPED_ARRAY;
                }

                return immutable ? LayoutTypes.IMMUTABLE_ARRAY : LayoutTypes.ARRAY;
            }
            case SET: {

                assert logicalType instanceof SetPropertyType;
                SetPropertyType sp = (SetPropertyType) logicalType;

                if ((sp.items() != null) && (sp.items().type() != TypeKind.Any)) {

                    final Out<TypeArgumentList> out = new Out<>();

                    LayoutType itemType = LayoutCompiler.logicalToPhysicalType(ns, sp.items(), out);
                    TypeArgumentList itemTypeArgs = out.get();

                    if (sp.items().nullable()) {
                        itemTypeArgs = new TypeArgumentList(new TypeArgument(itemType, itemTypeArgs));
                        itemType = itemType.isImmutable() ? LayoutTypes.IMMUTABLE_NULLABLE : LayoutTypes.NULLABLE;
                    }

                    typeArgs.set(new TypeArgumentList(new TypeArgument(itemType, itemTypeArgs)));
                    return immutable ? LayoutTypes.IMMUTABLE_TYPED_SET : LayoutTypes.TYPED_SET;
                }

                // TODO(283638): implement sparse set

                throw new LayoutCompilationException(lenientFormat(
                    "Unknown property type: %s",
                    logicalType.type()
                ));
            }
            case MAP: {

                assert logicalType instanceof MapPropertyType;
                MapPropertyType mp = (MapPropertyType) logicalType;

                if (mp.keys() != null && (mp.keys().type() != TypeKind.Any) && (mp.values() != null) && (mp.values().type() != TypeKind.Any)) {

                    final Out<TypeArgumentList> out = new Out<>();

                    LayoutType keyType = LayoutCompiler.logicalToPhysicalType(ns, mp.keys(), out);
                    TypeArgumentList keyTypeArgs = out.get();

                    if (mp.keys().nullable()) {
                        keyTypeArgs = new TypeArgumentList(new TypeArgument(keyType, keyTypeArgs));
                        keyType = keyType.isImmutable() ? LayoutTypes.IMMUTABLE_NULLABLE : LayoutTypes.NULLABLE;
                    }

                    LayoutType valueType = LayoutCompiler.logicalToPhysicalType(ns, mp.values(), out);
                    TypeArgumentList valueTypeArgs = out.get();

                    if (mp.values().nullable()) {
                        valueTypeArgs = new TypeArgumentList(new TypeArgument(valueType, valueTypeArgs));
                        valueType = valueType.isImmutable() ? LayoutTypes.IMMUTABLE_NULLABLE : LayoutTypes.NULLABLE;
                    }

                    typeArgs.set(new TypeArgumentList(
                        new TypeArgument(keyType, keyTypeArgs),
                        new TypeArgument(valueType, valueTypeArgs)
                    ));

                    return immutable ? LayoutTypes.IMMUTABLE_TYPED_MAP : LayoutTypes.TYPED_MAP;
                }

                // TODO(283638): implement sparse map

                throw new LayoutCompilationException(lenientFormat(
                    "Unknown property type: %s", logicalType.type())
                );
            }
            case Tuple: {

                assert logicalType instanceof TuplePropertyType;
                final TuplePropertyType tp = (TuplePropertyType) logicalType;

                final TypeArgument[] args = new TypeArgument[tp.items().size()];
                final Out<TypeArgumentList> out = new Out<>();

                for (int i = 0; i < tp.items().size(); i++) {

                    LayoutType itemType = LayoutCompiler.logicalToPhysicalType(ns, tp.items().get(i), out);
                    TypeArgumentList itemTypeArgs = out.get();

                    if (tp.items().get(i).nullable()) {
                        itemTypeArgs = new TypeArgumentList(new TypeArgument(itemType, itemTypeArgs));
                        itemType = itemType.isImmutable() ? LayoutTypes.IMMUTABLE_NULLABLE : LayoutTypes.NULLABLE;
                    }

                    args[i] = new TypeArgument(itemType, itemTypeArgs);
                }

                typeArgs.setAndGet(new TypeArgumentList(args));
                return immutable ? LayoutTypes.IMMUTABLE_TYPED_TUPLE : LayoutTypes.TYPED_TUPLE;
            }
            case TAGGED: {

                assert logicalType instanceof TaggedPropertyType;
                TaggedPropertyType tg = (TaggedPropertyType) logicalType;

                if (tg.items().size() < TaggedPropertyType.MinTaggedArguments || (tg.items().size() > TaggedPropertyType.MaxTaggedArguments)) {
                    throw new LayoutCompilationException(lenientFormat(
                        "Invalid number of arguments in Tagged: %s <= %s <= %s",
                        TaggedPropertyType.MinTaggedArguments,
                        tg.items().size(),
                        TaggedPropertyType.MaxTaggedArguments
                    ));
                }

                final Out<TypeArgumentList> out = new Out<>();
                final TypeArgument[] tgArgs = new TypeArgument[tg.items().size() + 1];

                tgArgs[0] = new TypeArgument(LayoutTypes.UINT_8, TypeArgumentList.EMPTY);

                for (int i = 0; i < tg.items().size(); i++) {

                    LayoutType itemType = LayoutCompiler.logicalToPhysicalType(ns, tg.items().get(i), out);
                    TypeArgumentList itemTypeArgs = out.get();

                    if (tg.items().get(i).nullable()) {
                        itemTypeArgs = new TypeArgumentList(new TypeArgument(itemType, itemTypeArgs));
                        itemType = itemType.isImmutable() ? LayoutTypes.IMMUTABLE_NULLABLE : LayoutTypes.NULLABLE;
                    }

                    tgArgs[i + 1] = new TypeArgument(itemType, itemTypeArgs);
                }

                typeArgs.set(new TypeArgumentList(tgArgs));

                switch (tg.items().size()) {
                    case 1:
                        return immutable ? LayoutTypes.IMMUTABLE_TAGGED : LayoutTypes.TAGGED;
                    case 2:
                        return immutable ? LayoutTypes.IMMUTABLE_TAGGED_2 : LayoutTypes.TAGGED_2;
                    default:
                        throw new LayoutCompilationException("Unexpected tagged arity");
                }
            }
            case Schema: {

                assert logicalType instanceof UdtPropertyType;
                UdtPropertyType up = (UdtPropertyType) logicalType;

                final Optional<Schema> udtSchema;

                if (up.schemaId() == SchemaId.INVALID) {
                    udtSchema = ns.schemas().stream()
                        .filter(schema -> up.name().equals(schema.name()))
                        .findFirst();
                } else {
                    udtSchema = ns.schemas().stream()
                        .filter(schema -> up.schemaId().equals(schema.schemaId()))
                        .findFirst();
                    if (up.name().equals(udtSchema.map(Schema::name).orElse(null))) {
                        throw new LayoutCompilationException(lenientFormat(
                            "Ambiguous schema reference: '%s:%s'", up.name(), up.schemaId()
                        ));
                    }
                }

                if (!udtSchema.isPresent()) {
                    throw new LayoutCompilationException(lenientFormat(
                        "Cannot resolve schema reference '%s:%s'", up.name(), up.schemaId()
                    ));
                }

                typeArgs.set(new TypeArgumentList(udtSchema.get().schemaId()));
                return immutable ? LayoutTypes.IMMUTABLE_UDT : LayoutTypes.UDT;
            }
            default:
                throw new LayoutCompilationException(Strings.lenientFormat(
                    "Unknown property type: %s", logicalType.type()
                ));
        }
    }
}
