package org.apache.seatunnel.transform.empty;

import com.google.auto.service.AutoService;
import lombok.NoArgsConstructor;
import lombok.NonNull;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.seatunnel.api.configuration.ReadonlyConfig;
import org.apache.seatunnel.api.configuration.util.ConfigValidator;
import org.apache.seatunnel.api.table.catalog.CatalogTable;
import org.apache.seatunnel.api.table.catalog.Column;
import org.apache.seatunnel.api.table.type.BasicType;
import org.apache.seatunnel.api.table.type.SeaTunnelDataType;
import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
import org.apache.seatunnel.api.transform.SeaTunnelTransform;
import org.apache.seatunnel.shade.com.typesafe.config.Config;
import org.apache.seatunnel.transform.common.MultipleFieldOutputTransform;
import org.apache.seatunnel.transform.common.SeaTunnelRowAccessor;
import org.apache.seatunnel.transform.replace.ReplaceTransformFactory;

import java.math.BigDecimal;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@AutoService(SeaTunnelTransform.class)
@NoArgsConstructor
public class ReplaceNullTransform extends MultipleFieldOutputTransform {

    private ReadonlyConfig config;

    private Map<String, Integer> inputFieldIndex;

    public ReplaceNullTransform(
            @NonNull ReadonlyConfig config, @NonNull CatalogTable inputCatalogTable) {
        super(inputCatalogTable);
        this.config = config;
        initOutputFields(
                inputCatalogTable.getTableSchema().toPhysicalRowDataType(),
                this.config.get(ReplaceNullConfig.KEY_REPLACE_FIELD_LIST));
    }

    @Override
    public String getPluginName() {
        return "ReplaceNull";
    }

    @Override
    protected void setConfig(Config pluginConfig) {
        ConfigValidator.of(ReadonlyConfig.fromConfig(pluginConfig))
                .validate(new ReplaceNullTransformFactory().optionRule());
        this.config = ReadonlyConfig.fromConfig(pluginConfig);
    }

    @Override
    protected void setInputRowType(SeaTunnelRowType inputRowType) {
        initOutputFields(inputRowType, config.get(ReplaceNullConfig.KEY_REPLACE_FIELD_LIST));
    }

    @Override
    protected String[] getOutputFieldNames() {
        List<ReplaceNullEntity> list = config.get(ReplaceNullConfig.KEY_REPLACE_FIELD_LIST);
        return list.stream().map(ReplaceNullEntity::getColumn).toArray(String[]::new);
    }

    @Override
    protected SeaTunnelDataType[] getOutputFieldDataTypes() {
        BasicType[] basicTypes = new BasicType[inputFieldIndex.size()];
        for (int i = 0; i < basicTypes.length; i++) {
            basicTypes[i] = BasicType.STRING_TYPE;
        }
        return basicTypes;
    }

    @Override
    protected Object[] getOutputFieldValues(SeaTunnelRowAccessor inputRow) {
        Object[] objects = new Object[inputFieldIndex.size()];
        List<ReplaceNullEntity> list = config.get(ReplaceNullConfig.KEY_REPLACE_FIELD_LIST);
        int i = 0;
        for (ReplaceNullEntity replaceNullEntity : list) {
            if (inputFieldIndex.get(replaceNullEntity.getColumn()) == null || inputFieldIndex.get(replaceNullEntity.getColumn()) < 0) {
                continue;
            }
            Object inputFieldValue = inputRow.getField(inputFieldIndex.get(replaceNullEntity.getColumn()));
            if (inputFieldValue == null || inputFieldValue.equals("") || inputFieldValue.equals("null")) {
                if (replaceNullEntity.getReplacement() != null && !replaceNullEntity.getReplacement().equals("")) {
                    objects[i] = replaceNullEntity.getReplacement();
                } else {
                    objects[i] = "0";
                }
            } else {
                objects[i] = inputFieldValue instanceof BigDecimal ? ((BigDecimal)inputFieldValue).toPlainString() : inputFieldValue.toString();
            }
            i++;
        }
        return objects;
    }

    @Override
    protected Column[] getOutputColumns() {
        Column[] columnList = new Column[inputFieldIndex.size()];
        List<ReplaceNullEntity> list = config.get(ReplaceNullConfig.KEY_REPLACE_FIELD_LIST);
        for (int i = 0; i < list.size(); i++) {
            List<Column> columns = inputCatalogTable.getTableSchema().getColumns();
            int finalI = i;
            List<Column> collect =
                    columns.stream().filter(column ->
                                    column.getName().equals(list.get(finalI).getColumn()))
                            .collect(Collectors.toList());
            if (CollectionUtils.isEmpty(collect)) {
                throw new IllegalArgumentException(
                        "Cannot find ["
                                + list.get(finalI).getColumn()
                                + "] field in input catalog table");
            }
            columnList[i] = collect.get(0).copy(BasicType.STRING_TYPE);
        }

        return columnList;
    }

    private void initOutputFields(SeaTunnelRowType inputRowType, List<ReplaceNullEntity> replaceFields) {
        inputFieldIndex = new HashMap<>();
        for (ReplaceNullEntity replaceNullEntity : replaceFields) {
            if (inputRowType.indexOf(replaceNullEntity.getColumn()) > 0) {
                inputFieldIndex.put(replaceNullEntity.getColumn(), inputRowType.indexOf(replaceNullEntity.getColumn()));
            } else{
                throw new IllegalArgumentException(
                        "Cannot find [" + replaceNullEntity.getColumn() + "] field in input row type");
            }
        }
    }
}
