/*
 * Copyright DataStax, Inc.
 *
 * Licensed 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 ai.langstream.ai.agents.commons;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.JsonNodeFactory;
import com.fasterxml.jackson.databind.node.ObjectNode;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalTime;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import org.apache.avro.Conversion;
import org.apache.avro.Conversions;
import org.apache.avro.Schema;
import org.apache.avro.data.TimeConversions;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericFixed;
import org.apache.avro.generic.GenericRecord;

/** Convert an AVRO GenericRecord to a JsonNode. */
public class JsonConverter {

    private static final Map<String, LogicalTypeConverter<?>> logicalTypeConverters =
            new HashMap<>();
    private static final JsonNodeFactory jsonNodeFactory =
            JsonNodeFactory.withExactBigDecimals(true);

    public static JsonNode toJson(GenericRecord genericRecord) {
        if (genericRecord == null) {
            return null;
        }
        ObjectNode objectNode = jsonNodeFactory.objectNode();
        for (Schema.Field field : genericRecord.getSchema().getFields()) {
            objectNode.set(field.name(), toJson(field.schema(), genericRecord.get(field.name())));
        }
        return objectNode;
    }

    public static JsonNode toJson(Schema schema, Object value) {
        if (schema.getLogicalType() != null
                && logicalTypeConverters.containsKey(schema.getLogicalType().getName())) {
            return logicalTypeConverters
                    .get(schema.getLogicalType().getName())
                    .toJson(schema, value);
        }
        if (value == null) {
            return jsonNodeFactory.nullNode();
        }
        switch (schema.getType()) {
            case NULL: // this should not happen
                return jsonNodeFactory.nullNode();
            case INT:
                return jsonNodeFactory.numberNode((Integer) value);
            case LONG:
                return jsonNodeFactory.numberNode((Long) value);
            case DOUBLE:
                return jsonNodeFactory.numberNode((Double) value);
            case FLOAT:
                return jsonNodeFactory.numberNode((Float) value);
            case BOOLEAN:
                return jsonNodeFactory.booleanNode((Boolean) value);
            case BYTES:
                byte[] bytes;
                if (value instanceof byte[]) {
                    bytes = (byte[]) value;
                } else if (value instanceof ByteBuffer) {
                    bytes = getBytes((ByteBuffer) value);
                } else {
                    throw new IllegalArgumentException(
                            "Invalid type for field of type BYTES, expected byte[] or ByteBuffer but was "
                                    + value.getClass());
                }
                return jsonNodeFactory.binaryNode(bytes);
            case FIXED:
                return jsonNodeFactory.binaryNode(((GenericFixed) value).bytes());
            case ENUM: // GenericEnumSymbol
            case STRING:
                return jsonNodeFactory.textNode(
                        value.toString()); // can be a String or org.apache.avro.util.Utf8
            case ARRAY:
                {
                    Schema elementSchema = schema.getElementType();
                    ArrayNode arrayNode = jsonNodeFactory.arrayNode();
                    Object[] iterable;
                    if (value instanceof GenericData.Array) {
                        iterable = ((GenericData.Array) value).toArray();
                    } else {
                        iterable = (Object[]) value;
                    }
                    for (Object elem : iterable) {
                        JsonNode fieldValue = toJson(elementSchema, elem);
                        arrayNode.add(fieldValue);
                    }
                    return arrayNode;
                }
            case MAP:
                {
                    Map<Object, Object> map = (Map<Object, Object>) value;
                    ObjectNode objectNode = jsonNodeFactory.objectNode();
                    for (Map.Entry<Object, Object> entry : map.entrySet()) {
                        JsonNode jsonNode = toJson(schema.getValueType(), entry.getValue());
                        // can be a String or org.apache.avro.util.Utf8
                        final String entryKey =
                                entry.getKey() == null ? null : entry.getKey().toString();
                        objectNode.set(entryKey, jsonNode);
                    }
                    return objectNode;
                }
            case RECORD:
                return toJson((GenericRecord) value);
            case UNION:
                for (Schema s : schema.getTypes()) {
                    if (s.getType() == Schema.Type.NULL) {
                        continue;
                    }
                    return toJson(s, value);
                }
                // this case should not happen
                return jsonNodeFactory.textNode(value.toString());
            default:
                throw new UnsupportedOperationException(
                        "Unknown AVRO schema type=" + schema.getType());
        }
    }

    abstract static class LogicalTypeConverter<T> {
        final Conversion<T> conversion;

        public LogicalTypeConverter(Conversion<T> conversion) {
            this.conversion = conversion;
        }

        abstract JsonNode toJson(Schema schema, Object value);
    }

    static {
        logicalTypeConverters.put(
                "decimal",
                new LogicalTypeConverter<BigDecimal>(new Conversions.DecimalConversion()) {
                    @Override
                    JsonNode toJson(Schema schema, Object value) {
                        if (!(value instanceof BigDecimal)) {
                            throw new IllegalArgumentException(
                                    "Invalid type for Decimal, expected BigDecimal but was "
                                            + value.getClass());
                        }
                        BigDecimal decimal = (BigDecimal) value;
                        return jsonNodeFactory.numberNode(decimal);
                    }
                });
        logicalTypeConverters.put(
                "date",
                new LogicalTypeConverter<LocalDate>(new TimeConversions.DateConversion()) {
                    @Override
                    JsonNode toJson(Schema schema, Object value) {
                        if (!(value instanceof Integer)) {
                            throw new IllegalArgumentException(
                                    "Invalid type for date, expected Integer but was "
                                            + value.getClass());
                        }
                        Integer daysFromEpoch = (Integer) value;
                        return jsonNodeFactory.numberNode(daysFromEpoch);
                    }
                });
        logicalTypeConverters.put(
                "time-millis",
                new LogicalTypeConverter<LocalTime>(new TimeConversions.TimeMillisConversion()) {
                    @Override
                    JsonNode toJson(Schema schema, Object value) {
                        if (!(value instanceof Integer)) {
                            throw new IllegalArgumentException(
                                    "Invalid type for time-millis, expected Integer but was "
                                            + value.getClass());
                        }
                        Integer timeMillis = (Integer) value;
                        return jsonNodeFactory.numberNode(timeMillis);
                    }
                });
        logicalTypeConverters.put(
                "time-micros",
                new LogicalTypeConverter<LocalTime>(new TimeConversions.TimeMicrosConversion()) {
                    @Override
                    JsonNode toJson(Schema schema, Object value) {
                        if (!(value instanceof Long)) {
                            throw new IllegalArgumentException(
                                    "Invalid type for time-micros, expected Long but was "
                                            + value.getClass());
                        }
                        Long timeMicro = (Long) value;
                        return jsonNodeFactory.numberNode(timeMicro);
                    }
                });
        logicalTypeConverters.put(
                "timestamp-millis",
                new LogicalTypeConverter<Instant>(new TimeConversions.TimestampMillisConversion()) {
                    @Override
                    JsonNode toJson(Schema schema, Object value) {
                        if (!(value instanceof Long)) {
                            throw new IllegalArgumentException(
                                    "Invalid type for timestamp-millis, expected Long but was "
                                            + value.getClass());
                        }
                        Long epochMillis = (Long) value;
                        return jsonNodeFactory.numberNode(epochMillis);
                    }
                });
        logicalTypeConverters.put(
                "timestamp-micros",
                new LogicalTypeConverter<Instant>(new TimeConversions.TimestampMicrosConversion()) {
                    @Override
                    JsonNode toJson(Schema schema, Object value) {
                        if (!(value instanceof Long)) {
                            throw new IllegalArgumentException(
                                    "Invalid type for timestamp-micros, expected Long but was "
                                            + value.getClass());
                        }
                        Long epochMillis = (Long) value;
                        return jsonNodeFactory.numberNode(epochMillis);
                    }
                });
        logicalTypeConverters.put(
                "uuid",
                new LogicalTypeConverter<UUID>(new Conversions.UUIDConversion()) {
                    @Override
                    JsonNode toJson(Schema schema, Object value) {
                        return jsonNodeFactory.textNode(value == null ? null : value.toString());
                    }
                });
    }

    public static byte[] getBytes(ByteBuffer byteBuffer) {
        if (byteBuffer == null) {
            return null;
        }
        if (byteBuffer.hasArray()
                && byteBuffer.arrayOffset() == 0
                && byteBuffer.array().length == byteBuffer.remaining()) {
            return byteBuffer.array();
        }
        // Direct buffer is not backed by array and it needs to be read from direct memory
        byte[] array = new byte[byteBuffer.remaining()];
        byteBuffer.get(array);
        return array;
    }
}
