package org.genomicsdb.model;

import gnu.getopt.Getopt;
import gnu.getopt.LongOpt;
import htsjdk.tribble.AbstractFeatureReader;
import htsjdk.tribble.FeatureReader;
import htsjdk.tribble.readers.LineIterator;
import htsjdk.variant.variantcontext.VariantContext;
import htsjdk.variant.vcf.VCFCodec;
import htsjdk.variant.vcf.VCFHeader;
import htsjdk.variant.vcf.VCFUtils;

import java.util.*;
import java.util.function.Consumer;
import java.util.stream.IntStream;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;

import static java.util.stream.Collectors.toList;

public class CommandLineImportConfig extends ImportConfig {
    //TODO: remove this constant once C++ layer makes use of protobuf structures
    protected static final long DEFAULT_SIZE_PER_COLUMN_PARTITION = 16384L;
    private GenomicsDBImportConfiguration.ImportConfiguration.Builder configurationBuilder =
            GenomicsDBImportConfiguration.ImportConfiguration.newBuilder();
    private List<GenomicsDBImportConfiguration.Partition.Builder> partitionList = new ArrayList<>();
    private String workspace = "";
    private String array = "";
    private String vcfOutputFilename = "";
    private boolean readVcfUsingHtslib = false;
    private int coalesceContigs = 0;
    private int consolidateOnly = -1;

    private Consumer<String[]> chromInterToPartitionList = par -> {
        GenomicsDBImportConfiguration.Partition.Builder partitionBuilder = GenomicsDBImportConfiguration.Partition.newBuilder();
        Coordinates.ContigPosition.Builder contigPositionBuilder = Coordinates.ContigPosition.newBuilder();
        Coordinates.GenomicsDBColumn.Builder columnBuilder = Coordinates.GenomicsDBColumn.newBuilder();
        //begin
        contigPositionBuilder.setContig(par[0]).setPosition(Long.parseLong(par[1].split("-")[0]));
        columnBuilder.setContigPosition(contigPositionBuilder.build());
        partitionBuilder.setBegin(columnBuilder.build());
        //end
        contigPositionBuilder.setPosition(Long.parseLong(par[1].split("-")[1]));
        columnBuilder.setContigPosition(contigPositionBuilder.build());
        partitionBuilder.setEnd(columnBuilder.build());
        partitionList.add(partitionBuilder);
    };

    public CommandLineImportConfig(final String command, final String[] commandArgs) {
        Getopt getOpt = new Getopt(command, commandArgs, "w:A:L:", resolveLongOpt());
        //TODO: remove next line once C++ layer makes use of protobuf structures. Making size per column partition explicit.
        configurationBuilder.setSizePerColumnPartition(DEFAULT_SIZE_PER_COLUMN_PARTITION);
        resolveCommandArgs(getOpt);
        partitionList.forEach(partition -> {
                    partition.setWorkspace(this.workspace);
                    if(array.isEmpty())
                         partition.setGenerateArrayNameFromPartitionBounds(true);
                    else
                         partition.setArrayName(array);
                    if (!vcfOutputFilename.isEmpty()) partition.setVcfOutputFilename(vcfOutputFilename);
                });
        partitionList.forEach(partition -> configurationBuilder.addColumnPartitions(partition));
        this.setImportConfiguration(configurationBuilder.build());
        this.validateChromosomeIntervals();
        int numPositionalArgs = commandArgs.length - getOpt.getOptind();
        if (numPositionalArgs <= 0 
                || this.getImportConfiguration().getColumnPartitions(0).getWorkspace().isEmpty()
                || this.getImportConfiguration().getColumnPartitionsList().isEmpty()) {
            throwIllegalArgumentException();
        }
        List<String> files = IntStream.range(getOpt.getOptind(), commandArgs.length).mapToObj(
                i -> commandArgs[i]).collect(toList());
        try {
            this.resolveHeaders(files);
        }
        catch(IOException e) {
            System.err.println("IOException thrown "+e);
        }
        if(!this.readVcfUsingHtslib)
            this.setSampleToReaderMapCreator(this::createSampleToReaderMap);
    }

    private LongOpt[] resolveLongOpt() {
        LongOpt[] longopts = new LongOpt[16];
        longopts[0] = new LongOpt("use_samples_in_order", LongOpt.NO_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_USE_SAMPLES_IN_ORDER.idx());
        longopts[1] = new LongOpt("fail_if_updating", LongOpt.NO_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_FAIL_IF_UPDATING.idx());
        longopts[2] = new LongOpt("interval", LongOpt.REQUIRED_ARGUMENT, null, 'L');
        longopts[3] = new LongOpt("workspace", LongOpt.REQUIRED_ARGUMENT, null, 'w');
        longopts[4] = new LongOpt("batchsize", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_BATCHSIZE.idx());
        longopts[5] = new LongOpt("vidmap-output", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_VIDMAP_OUTPUT.idx());
        longopts[6] = new LongOpt("callset-output", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_CALLSET_OUTPUT.idx());
        longopts[7] = new LongOpt("pass-as-bcf", LongOpt.NO_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_PASS_AS_BCF.idx());
        longopts[8] = new LongOpt("vcf-header-output", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_VCF_HEADER_OUTPUT.idx());
        longopts[9] = new LongOpt("size_per_column_partition", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_SIZE_PER_COLUMN_PARTITION.idx());
        longopts[10] = new LongOpt("segment_size", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_SEGMENT_SIZE.idx());
        longopts[11] = new LongOpt("array", LongOpt.REQUIRED_ARGUMENT, null, 'A');
        longopts[12] = new LongOpt("htslib", LongOpt.NO_ARGUMENT, null,
            ArgsIdxEnum.ARGS_IDX_READ_INPUT_VCF_USING_HTSLIB.idx());
        longopts[13] = new LongOpt("incremental_import", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_INCREMENTAL_IMPORT.idx());
        longopts[14] = new LongOpt("coalesce-multiple-contigs", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_COALESCE_CONTIGS.idx());
        longopts[15] = new LongOpt("consolidate-only", LongOpt.REQUIRED_ARGUMENT, null,
                ArgsIdxEnum.ARGS_IDX_CONSOLIDATE_ONLY.idx());
        return longopts;
    }

    private void resolveCommandArgs(final Getopt commandArgs) {
        int c;
        final int firstEnumIdx = ArgsIdxEnum.ARGS_IDX_USE_SAMPLES_IN_ORDER.idx();
        final ArgsIdxEnum[] enumArray = ArgsIdxEnum.values();
        while ((c = commandArgs.getopt()) != -1) {
            switch (c) {
                case 'w':
                    workspace = commandArgs.getOptarg();
                    break;
                case 'A':
                    array = commandArgs.getOptarg();
                    break;
                case 'L':
                    chromInterToPartitionList.accept(commandArgs.getOptarg().split(":"));
                    break;
                default: {
                    if (c >= firstEnumIdx && c < ArgsIdxEnum.ARGS_IDX_AFTER_LAST_ARG_IDX.idx()) {
                        int offset = c - firstEnumIdx;
                        assert offset < enumArray.length;
                        switch (enumArray[offset]) {
                            case ARGS_IDX_USE_SAMPLES_IN_ORDER:
                                this.setUseSamplesInOrder(true);
                                break;
                            case ARGS_IDX_FAIL_IF_UPDATING:
                                configurationBuilder.setFailIfUpdating(true);
                                break;
                            case ARGS_IDX_BATCHSIZE:
                                setBatchSize(Integer.parseInt(commandArgs.getOptarg()));
                                break;
                            case ARGS_IDX_VIDMAP_OUTPUT:
                                this.setOutputVidmapJsonFile(commandArgs.getOptarg());
                                break;
                            case ARGS_IDX_CALLSET_OUTPUT:
                                this.setOutputCallsetmapJsonFile(commandArgs.getOptarg());
                                break;
                            case ARGS_IDX_VCF_HEADER_OUTPUT:
                                this.setOutputVcfHeaderFile(commandArgs.getOptarg());
                                break;
                            case ARGS_IDX_PASS_AS_BCF:
                                setPassAsVcf(false);
                                break;
                            case ARGS_IDX_SIZE_PER_COLUMN_PARTITION:
                                configurationBuilder.setSizePerColumnPartition(Long.parseLong(commandArgs.getOptarg()));
                                break;
                            case ARGS_IDX_SEGMENT_SIZE:
                                configurationBuilder.setSegmentSize(Long.parseLong(commandArgs.getOptarg()));
                                break;
                            case ARGS_IDX_READ_INPUT_VCF_USING_HTSLIB:
                                readVcfUsingHtslib = true;
                                break;
                            case ARGS_IDX_INCREMENTAL_IMPORT:
                                this.setIncrementalImport(true);
                                configurationBuilder.setLbCallsetRowIdx(Long.parseLong(commandArgs.getOptarg()));
                                break;
                            case ARGS_IDX_COALESCE_CONTIGS:
                                coalesceContigs = Integer.parseInt(commandArgs.getOptarg());
                                break;
                            case ARGS_IDX_CONSOLIDATE_ONLY:
                                consolidateOnly = Integer.parseInt(commandArgs.getOptarg());
                                break;
                            default:
                                throwIllegalArgumentException(commandArgs);
                                return;
                        }
                    } else {
                        throwIllegalArgumentException(commandArgs);
                    }
                }
            }
        }
    }

    private void throwIllegalArgumentException() {
        throw new IllegalArgumentException("Invalid usage. Correct way of using arguments: -L chromosome:interval " +
                "-w genomicsdbworkspace [-A array] --size_per_column_partition 10000 --segment_size 1048576 variantfile(s) " +
                "[--use_samples_in_order --fail_if_updating --batchsize=<N> --vidmap-output <path> --incremental_update]");
    }

    private void throwIllegalArgumentException(Getopt commandArgs) {
        throw new IllegalArgumentException("Unknown command line option " +
                commandArgs.getOptarg() + " - ignored");
    }

    private void resolveHeaders(final List<String> files) throws IOException {
        List<VCFHeader> headers = new ArrayList<>();
        Map<String, URI> sampleNameToVcfPath = new LinkedHashMap<>();

        //Get merged header first
        for (String file : files) {
            AbstractFeatureReader<VariantContext, LineIterator> reader =
                    AbstractFeatureReader.getFeatureReader(file, new VCFCodec(), false);
            if(reader == null)
              throw new IOException("Could not open "+file);
            headers.add((VCFHeader) reader.getHeader());
            final String sampleName = ((VCFHeader) reader.getHeader()).getGenotypeSamples().get(0);
            try {
                sampleNameToVcfPath.put(sampleName, new URI(file));
            }
            catch (URISyntaxException e) {
                throw new IOException("URISyntaxException : "+e);
            }
            reader.close();
            //Hopefully, GC kicks in and frees resources assigned to reader
        }

        this.setMergedHeader(VCFUtils.smartMergeHeaders(headers, true));
        this.setSampleNameToVcfPath(sampleNameToVcfPath);
    }

    private Map<String, FeatureReader<VariantContext>> createSampleToReaderMap(
            final Map<String, URI> sampleNameToVcfPath, final int batchSize, final int index) {
        final Map<String, FeatureReader<VariantContext>> sampleToReaderMap = new LinkedHashMap<>();
        for (int j = index; j < sampleNameToVcfPath.size() && j < index + batchSize; ++j) {
            final String sampleName = sampleNameToVcfPath.keySet().toArray()[j].toString(); //TODO: fix this.
            assert sampleNameToVcfPath.containsKey(sampleName);
            AbstractFeatureReader<VariantContext, LineIterator> reader = AbstractFeatureReader.getFeatureReader(
                    sampleNameToVcfPath.get(sampleName).toString(), new VCFCodec(), false);
            assert sampleName.equals(((VCFHeader) reader.getHeader()).getGenotypeSamples().get(0));
            sampleToReaderMap.put(sampleName, reader);
        }
        return sampleToReaderMap;
    }

    public int getNumberCoalesceContigs() {
        return coalesceContigs;
    }

    public int getConsolidateOnlyThreads() {
       return consolidateOnly;
    }

    public enum ArgsIdxEnum {
        ARGS_IDX_USE_SAMPLES_IN_ORDER(1000),
        ARGS_IDX_FAIL_IF_UPDATING(1001),
        ARGS_IDX_BATCHSIZE(1002),
        ARGS_IDX_VIDMAP_OUTPUT(1003),
        ARGS_IDX_CALLSET_OUTPUT(1004),
        ARGS_IDX_PASS_AS_BCF(1005),
        ARGS_IDX_VCF_HEADER_OUTPUT(1006),
        ARGS_IDX_SIZE_PER_COLUMN_PARTITION(1007),
        ARGS_IDX_SEGMENT_SIZE(1008),
        ARGS_IDX_READ_INPUT_VCF_USING_HTSLIB(1009),
        ARGS_IDX_INCREMENTAL_IMPORT(1010),
        ARGS_IDX_COALESCE_CONTIGS(1011),
        ARGS_IDX_CONSOLIDATE_ONLY(1012),
        ARGS_IDX_AFTER_LAST_ARG_IDX(1013);
        private final int mArgsIdx;

        ArgsIdxEnum(final int idx) {
            mArgsIdx = idx;
        }

        public int idx() {
            return mArgsIdx;
        }
    }
}
