/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  *  See the NOTICE file distributed with this work for additional
 *  *  information regarding copyright ownership.
 *  * 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.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */

package org.datavec.api.records.reader.impl.jackson;

import lombok.Getter;
import lombok.Setter;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.datavec.api.conf.Configuration;
import org.datavec.api.io.labels.PathLabelGenerator;
import org.datavec.api.records.Record;
import org.datavec.api.records.metadata.RecordMetaData;
import org.datavec.api.records.metadata.RecordMetaDataURI;
import org.datavec.api.records.reader.BaseRecordReader;
import org.datavec.api.split.FileSplit;
import org.datavec.api.split.InputSplit;
import org.datavec.api.writable.Writable;
import org.nd4j.shade.jackson.core.type.TypeReference;
import org.nd4j.shade.jackson.databind.ObjectMapper;

import java.io.*;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.*;

public class JacksonRecordReader extends BaseRecordReader {

    private static final TypeReference<Map<String, Object>> typeRef = new TypeReference<Map<String, Object>>() {};

    private FieldSelection selection;
    private ObjectMapper mapper;
    private boolean shuffle;
    private long rngSeed;
    private PathLabelGenerator labelGenerator;
    private int labelPosition;
    private InputSplit is;
    private Random r;
    @Getter @Setter
    private String charset = StandardCharsets.UTF_8.name(); //Using String as StandardCharsets.UTF_8 is not serializable

    private URI[] uris;
    private int cursor = 0;

    public JacksonRecordReader(FieldSelection selection, ObjectMapper mapper) {
        this(selection, mapper, false);
    }

    public JacksonRecordReader(FieldSelection selection, ObjectMapper mapper, boolean shuffle) {
        this(selection, mapper, shuffle, System.currentTimeMillis(), null);
    }

    public JacksonRecordReader(FieldSelection selection, ObjectMapper mapper, boolean shuffle, long rngSeed,
                    PathLabelGenerator labelGenerator) {
        this(selection, mapper, shuffle, rngSeed, labelGenerator, -1);
    }

    public JacksonRecordReader(FieldSelection selection, ObjectMapper mapper, boolean shuffle, long rngSeed,
                    PathLabelGenerator labelGenerator, int labelPosition) {
        this.selection = selection;
        this.mapper = mapper;
        this.shuffle = shuffle;
        this.rngSeed = rngSeed;
        if (shuffle)
            r = new Random(rngSeed);
        this.labelGenerator = labelGenerator;
        this.labelPosition = labelPosition;
    }

    @Override
    public void initialize(InputSplit split) throws IOException, InterruptedException {
        if (split instanceof FileSplit)
            throw new UnsupportedOperationException("Cannot use JacksonRecordReader with FileSplit");
        super.initialize(inputSplit);
        this.uris = split.locations();
        if (shuffle) {
            List<URI> list = Arrays.asList(uris);
            Collections.shuffle(list, r);
            uris = list.toArray(new URI[uris.length]);
        }
    }

    @Override
    public void initialize(Configuration conf, InputSplit split) throws IOException, InterruptedException {
        initialize(split);
    }

    @Override
    public List<Writable> next() {
        if (uris == null)
            throw new IllegalStateException("URIs are null. Not initialized?");
        if (!hasNext())
            throw new NoSuchElementException("No next element");

        URI uri = uris[cursor++];
        invokeListeners(uri);
        String fileAsString;
        try (InputStream s = streamCreatorFn.apply(uri)){
            fileAsString = IOUtils.toString(s, charset);
        } catch (IOException e) {
            throw new RuntimeException("Error reading URI file", e);
        }

        return readValues(uri, fileAsString);

    }

    @Override
    public boolean hasNext() {
        return cursor < uris.length;
    }

    @Override
    public List<String> getLabels() {
        throw new UnsupportedOperationException();
    }

    @Override
    public void reset() {
        cursor = 0;
        if (shuffle) {
            List<URI> list = Arrays.asList(uris);
            Collections.shuffle(list, r);
            uris = list.toArray(new URI[uris.length]);
        }
    }

    @Override
    public boolean resetSupported() {
        return true;
    }

    @Override
    public List<Writable> record(URI uri, DataInputStream dataInputStream) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(dataInputStream));
        StringBuilder sb = new StringBuilder();
        String line;
        while ((line = br.readLine()) != null) {
            sb.append(line).append("\n");
        }

        return readValues(uri, sb.toString());
    }

    @Override
    public void close() throws IOException {

    }

    @Override
    public void setConf(Configuration conf) {

    }

    @Override
    public Configuration getConf() {
        return null;
    }


    private List<Writable> readValues(URI uri, String fileContents) {
        List<Writable> out = JacksonReaderUtils.parseRecord(fileContents, selection, mapper);

        //Add label - if required
        if(labelGenerator != null){
            Writable label = labelGenerator.getLabelForPath(uri);
            List<String[]> paths = selection.getFieldPaths();
            if ((labelPosition >= paths.size() || labelPosition == -1)) {
                //Edge case: might want label as the last value
                out.add(label);
            } else {
                out.add(labelPosition, label);  //Add and shift existing to right
            }
        }

        return out;
    }

    @Override
    public Record nextRecord() {
        URI currentURI = uris[cursor];
        List<Writable> writables = next();
        RecordMetaData meta = new RecordMetaDataURI(currentURI, JacksonRecordReader.class);
        return new org.datavec.api.records.impl.Record(writables, meta);
    }

    @Override
    public Record loadFromMetaData(RecordMetaData recordMetaData) throws IOException {
        return loadFromMetaData(Collections.singletonList(recordMetaData)).get(0);
    }

    @Override
    public List<Record> loadFromMetaData(List<RecordMetaData> recordMetaDatas) throws IOException {

        List<Record> out = new ArrayList<>();
        for (RecordMetaData metaData : recordMetaDatas) {
            URI uri = metaData.getURI();

            String fileAsString;
            try {
                fileAsString = FileUtils.readFileToString(new File(uri.toURL().getFile()));
            } catch (IOException e) {
                throw new RuntimeException("Error reading URI file", e);
            }

            List<Writable> writables = readValues(uri, fileAsString);
            out.add(new org.datavec.api.records.impl.Record(writables, metaData));
        }

        return out;
    }
}
