/*
 * 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.paimon.spark;

import org.apache.paimon.catalog.CatalogContext;
import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.BinaryVector;
import org.apache.paimon.data.Blob;
import org.apache.paimon.data.Decimal;
import org.apache.paimon.data.InternalArray;
import org.apache.paimon.data.InternalMap;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.data.InternalVector;
import org.apache.paimon.data.Timestamp;
import org.apache.paimon.data.variant.Variant;
import org.apache.paimon.spark.util.shim.TypeUtils;
import org.apache.paimon.types.ArrayType;
import org.apache.paimon.types.DataType;
import org.apache.paimon.types.DateType;
import org.apache.paimon.types.MapType;
import org.apache.paimon.types.RowKind;
import org.apache.paimon.types.RowType;
import org.apache.paimon.types.VectorType;
import org.apache.paimon.utils.DateTimeUtils;
import org.apache.paimon.utils.UriReaderFactory;

import org.apache.spark.sql.Row;
import org.apache.spark.sql.paimon.shims.SparkShimLoader;

import java.io.Serializable;
import java.sql.Date;
import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

import scala.collection.JavaConverters;

/** An {@link InternalRow} wraps spark {@link Row} for v1 write. */
public class SparkRow implements InternalRow, Serializable {

    private final RowType type;
    private final Row row;
    private final RowKind rowKind;
    private final UriReaderFactory uriReaderFactory;

    public SparkRow(RowType type, Row row) {
        this(type, row, RowKind.INSERT, (CatalogContext) null);
    }

    public SparkRow(RowType type, Row row, RowKind rowkind, CatalogContext catalogContext) {
        this(type, row, rowkind, new UriReaderFactory(catalogContext));
    }

    public static SparkRow fromUriReaderFactory(
            RowType type, Row row, RowKind rowkind, UriReaderFactory uriReaderFactory) {
        return new SparkRow(type, row, rowkind, uriReaderFactory);
    }

    private SparkRow(RowType type, Row row, RowKind rowkind, UriReaderFactory uriReaderFactory) {
        this.type = type;
        this.row = row;
        this.rowKind = rowkind;
        this.uriReaderFactory = uriReaderFactory;
    }

    @Override
    public int getFieldCount() {
        return type.getFieldCount();
    }

    @Override
    public RowKind getRowKind() {
        return rowKind;
    }

    @Override
    public void setRowKind(RowKind rowKind) {
        throw new UnsupportedOperationException(
                "Spark row does not support modifying rowkind field");
    }

    @Override
    public boolean isNullAt(int i) {
        return row.isNullAt(i);
    }

    @Override
    public boolean getBoolean(int i) {
        return row.getBoolean(i);
    }

    @Override
    public byte getByte(int i) {
        return row.getByte(i);
    }

    @Override
    public short getShort(int i) {
        return row.getShort(i);
    }

    @Override
    public int getInt(int i) {
        if (type.getTypeAt(i) instanceof DateType) {
            return toPaimonDate(row.get(i));
        }
        return row.getInt(i);
    }

    @Override
    public long getLong(int i) {
        return row.getLong(i);
    }

    @Override
    public float getFloat(int i) {
        return row.getFloat(i);
    }

    @Override
    public double getDouble(int i) {
        return row.getDouble(i);
    }

    @Override
    public BinaryString getString(int i) {
        return BinaryString.fromString(row.getString(i));
    }

    @Override
    public Decimal getDecimal(int i, int precision, int scale) {
        return Decimal.fromBigDecimal(row.getDecimal(i), precision, scale);
    }

    @Override
    public Timestamp getTimestamp(int i, int precision) {
        return toPaimonTimestamp(row.get(i));
    }

    @Override
    public byte[] getBinary(int i) {
        return row.getAs(i);
    }

    @Override
    public Variant getVariant(int i) {
        return SparkShimLoader.shim().toPaimonVariant(row.getAs(i));
    }

    @Override
    public Blob getBlob(int i) {
        return Blob.fromBytes(row.getAs(i), uriReaderFactory, null);
    }

    @Override
    public InternalArray getArray(int i) {
        return new PaimonArray(
                ((ArrayType) type.getTypeAt(i)).getElementType(), row.getList(i), uriReaderFactory);
    }

    @Override
    public InternalVector getVector(int pos) {
        if (row.isNullAt(pos)) {
            return null;
        }
        return toPaimonVector((VectorType) type.getTypeAt(pos), row.get(pos));
    }

    @Override
    public InternalMap getMap(int i) {
        return toPaimonMap((MapType) type.getTypeAt(i), row.getJavaMap(i), uriReaderFactory);
    }

    @Override
    public InternalRow getRow(int i, int i1) {
        return new SparkRow(
                (RowType) type.getTypeAt(i), row.getStruct(i), RowKind.INSERT, uriReaderFactory);
    }

    private static int toPaimonDate(Object object) {
        if (object instanceof Date) {
            return DateTimeUtils.toInternal((Date) object);
        } else {
            return DateTimeUtils.toInternal((LocalDate) object);
        }
    }

    private static Timestamp toPaimonTimestamp(Object object) {
        if (object instanceof java.sql.Timestamp) {
            java.sql.Timestamp ts = (java.sql.Timestamp) object;
            if (TypeUtils.treatPaimonTimestampTypeAsSparkTimestampType()) {
                return Timestamp.fromSQLTimestamp(ts);
            } else {
                return Timestamp.fromInstant(ts.toInstant());
            }
        } else if (object instanceof java.time.Instant) {
            Instant instant = (Instant) object;
            if (TypeUtils.treatPaimonTimestampTypeAsSparkTimestampType()) {
                LocalDateTime localDateTime =
                        LocalDateTime.ofInstant((Instant) object, ZoneId.systemDefault());
                return Timestamp.fromLocalDateTime(localDateTime);
            } else {
                return Timestamp.fromInstant(instant);
            }
        } else {
            return Timestamp.fromLocalDateTime((LocalDateTime) object);
        }
    }

    private static InternalMap toPaimonMap(
            MapType mapType, Map<Object, Object> map, UriReaderFactory uriReaderFactory) {
        List<Object> keys = new ArrayList<>();
        List<Object> values = new ArrayList<>();
        map.forEach(
                (k, v) -> {
                    keys.add(k);
                    values.add(v);
                });

        PaimonArray key = new PaimonArray(mapType.getKeyType(), keys, uriReaderFactory);
        PaimonArray value = new PaimonArray(mapType.getValueType(), values, uriReaderFactory);
        return new InternalMap() {
            @Override
            public int size() {
                return map.size();
            }

            @Override
            public InternalArray keyArray() {
                return key;
            }

            @Override
            public InternalArray valueArray() {
                return value;
            }
        };
    }

    private static class PaimonArray implements InternalArray {

        private final DataType elementType;
        private final List<Object> list;
        private final UriReaderFactory uriReaderFactory;

        private PaimonArray(DataType elementType, List<Object> list) {
            this(elementType, list, null);
        }

        private PaimonArray(
                DataType elementType, List<Object> list, UriReaderFactory uriReaderFactory) {
            this.list = list;
            this.elementType = elementType;
            this.uriReaderFactory = uriReaderFactory;
        }

        @Override
        public int size() {
            return list.size();
        }

        @Override
        public boolean isNullAt(int i) {
            return list.get(i) == null;
        }

        @SuppressWarnings("unchecked")
        private <T> T getAs(int i) {
            return (T) list.get(i);
        }

        @Override
        public boolean getBoolean(int i) {
            return getAs(i);
        }

        @Override
        public byte getByte(int i) {
            return getAs(i);
        }

        @Override
        public short getShort(int i) {
            return getAs(i);
        }

        @Override
        public int getInt(int i) {
            if (elementType instanceof DateType) {
                return toPaimonDate(getAs(i));
            }
            return getAs(i);
        }

        @Override
        public long getLong(int i) {
            return getAs(i);
        }

        @Override
        public float getFloat(int i) {
            return getAs(i);
        }

        @Override
        public double getDouble(int i) {
            return getAs(i);
        }

        @Override
        public BinaryString getString(int i) {
            return BinaryString.fromString(getAs(i));
        }

        @Override
        public Decimal getDecimal(int i, int precision, int scale) {
            return Decimal.fromBigDecimal(getAs(i), precision, scale);
        }

        @Override
        public Timestamp getTimestamp(int i, int precision) {
            return toPaimonTimestamp(getAs(i));
        }

        @Override
        public byte[] getBinary(int i) {
            return getAs(i);
        }

        @Override
        public Variant getVariant(int i) {
            return SparkShimLoader.shim().toPaimonVariant(getAs(i));
        }

        @Override
        public Blob getBlob(int i) {
            return Blob.fromBytes(getAs(i), uriReaderFactory, null);
        }

        @Override
        public InternalArray getArray(int i) {
            Object o = getAs(i);
            List<Object> array =
                    o instanceof scala.collection.Seq
                            ? JavaConverters.seqAsJavaList((scala.collection.Seq<Object>) o)
                            : (List<Object>) o;
            return new PaimonArray(
                    ((ArrayType) elementType).getElementType(), array, uriReaderFactory);
        }

        @Override
        public InternalVector getVector(int pos) {
            throw new UnsupportedOperationException("Not support VectorType yet.");
        }

        @Override
        public InternalMap getMap(int i) {
            Object o = getAs(i);
            Map<Object, Object> map =
                    o instanceof scala.collection.Map
                            ? JavaConverters.mapAsJavaMap((scala.collection.Map<Object, Object>) o)
                            : (Map<Object, Object>) o;
            return toPaimonMap((MapType) elementType, map, uriReaderFactory);
        }

        @Override
        public InternalRow getRow(int i, int i1) {
            return new SparkRow((RowType) elementType, getAs(i), RowKind.INSERT, uriReaderFactory);
        }

        @Override
        public boolean[] toBooleanArray() {
            boolean[] res = new boolean[size()];
            for (int i = 0; i < size(); i++) {
                res[i] = getBoolean(i);
            }
            return res;
        }

        @Override
        public byte[] toByteArray() {
            byte[] res = new byte[size()];
            for (int i = 0; i < size(); i++) {
                res[i] = getByte(i);
            }
            return res;
        }

        @Override
        public short[] toShortArray() {
            short[] res = new short[size()];
            for (int i = 0; i < size(); i++) {
                res[i] = getShort(i);
            }
            return res;
        }

        @Override
        public int[] toIntArray() {
            int[] res = new int[size()];
            for (int i = 0; i < size(); i++) {
                res[i] = getInt(i);
            }
            return res;
        }

        @Override
        public long[] toLongArray() {
            long[] res = new long[size()];
            for (int i = 0; i < size(); i++) {
                res[i] = getLong(i);
            }
            return res;
        }

        @Override
        public float[] toFloatArray() {
            float[] res = new float[size()];
            for (int i = 0; i < size(); i++) {
                res[i] = getFloat(i);
            }
            return res;
        }

        @Override
        public double[] toDoubleArray() {
            double[] res = new double[size()];
            for (int i = 0; i < size(); i++) {
                res[i] = getDouble(i);
            }
            return res;
        }
    }

    private static InternalVector toPaimonVector(VectorType vectorType, Object vector) {
        if (vector == null) {
            return null;
        }
        if (vector instanceof boolean[]) {
            return BinaryVector.fromPrimitiveArray((boolean[]) vector);
        } else if (vector instanceof byte[]) {
            return BinaryVector.fromPrimitiveArray((byte[]) vector);
        } else if (vector instanceof short[]) {
            return BinaryVector.fromPrimitiveArray((short[]) vector);
        } else if (vector instanceof int[]) {
            return BinaryVector.fromPrimitiveArray((int[]) vector);
        } else if (vector instanceof long[]) {
            return BinaryVector.fromPrimitiveArray((long[]) vector);
        } else if (vector instanceof float[]) {
            return BinaryVector.fromPrimitiveArray((float[]) vector);
        } else if (vector instanceof double[]) {
            return BinaryVector.fromPrimitiveArray((double[]) vector);
        }
        if (vector instanceof scala.collection.Seq) {
            vector = JavaConverters.seqAsJavaList((scala.collection.Seq<Object>) vector);
        } else if (vector.getClass().isArray()) {
            vector = Arrays.asList((Object[]) vector);
        }
        if (!(vector instanceof List)) {
            throw new UnsupportedOperationException(
                    "Unsupported vector object: " + vector.getClass().getName());
        }
        return toPaimonVector(vectorType, (List<?>) vector);
    }

    private static InternalVector toPaimonVector(VectorType vectorType, List<?> list) {
        int expectedLength = vectorType.getLength();
        if (list.size() != expectedLength) {
            throw new IllegalArgumentException(
                    String.format(
                            "Vector length mismatch. Expected %d but was %d.",
                            expectedLength, list.size()));
        }
        switch (vectorType.getElementType().getTypeRoot()) {
            case BOOLEAN:
                boolean[] booleanValues = new boolean[expectedLength];
                for (int i = 0; i < expectedLength; i++) {
                    booleanValues[i] = (Boolean) list.get(i);
                }
                return BinaryVector.fromPrimitiveArray(booleanValues);
            case TINYINT:
                byte[] byteValues = new byte[expectedLength];
                for (int i = 0; i < expectedLength; i++) {
                    byteValues[i] = ((Number) list.get(i)).byteValue();
                }
                return BinaryVector.fromPrimitiveArray(byteValues);
            case SMALLINT:
                short[] shortValues = new short[expectedLength];
                for (int i = 0; i < expectedLength; i++) {
                    shortValues[i] = ((Number) list.get(i)).shortValue();
                }
                return BinaryVector.fromPrimitiveArray(shortValues);
            case INTEGER:
                int[] intValues = new int[expectedLength];
                for (int i = 0; i < expectedLength; i++) {
                    intValues[i] = ((Number) list.get(i)).intValue();
                }
                return BinaryVector.fromPrimitiveArray(intValues);
            case BIGINT:
                long[] longValues = new long[expectedLength];
                for (int i = 0; i < expectedLength; i++) {
                    longValues[i] = ((Number) list.get(i)).longValue();
                }
                return BinaryVector.fromPrimitiveArray(longValues);
            case FLOAT:
                float[] floatValues = new float[expectedLength];
                for (int i = 0; i < expectedLength; i++) {
                    floatValues[i] = ((Number) list.get(i)).floatValue();
                }
                return BinaryVector.fromPrimitiveArray(floatValues);
            case DOUBLE:
                double[] doubleValues = new double[expectedLength];
                for (int i = 0; i < expectedLength; i++) {
                    doubleValues[i] = ((Number) list.get(i)).doubleValue();
                }
                return BinaryVector.fromPrimitiveArray(doubleValues);
            default:
                throw new UnsupportedOperationException(
                        "Unsupported element type for vector " + vectorType.getElementType());
        }
    }
}
