package jalview.ws2.operations;

import static java.lang.String.format;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;

import jalview.analysis.SeqsetUtils;
import jalview.analysis.SeqsetUtils.SequenceInfo;
import jalview.api.AlignViewportI;
import jalview.bin.Cache;
import jalview.commands.RemoveGapsCommand;
import jalview.datamodel.Alignment;
import jalview.datamodel.AlignmentAnnotation;
import jalview.datamodel.AlignmentI;
import jalview.datamodel.AlignmentView;
import jalview.datamodel.HiddenColumns;
import jalview.datamodel.SeqCigar;
import jalview.datamodel.SequenceI;
import jalview.gui.AlignFrame;
import jalview.io.JnetAnnotationMaker;
import jalview.util.MessageManager;
import jalview.ws.params.ArgumentI;
import jalview.ws2.WSJob;
import jalview.ws2.WSJobStatus;
import jalview.ws2.operations.AlignmentWorker.AlignmentJob;

public class JPredWorker extends AbstractPollableWorker
{

  private class InputFormatParameter implements ArgumentI
  {
    String value = "";

    @Override
    public String getName()
    {
      return "format";
    }

    @Override
    public String getValue()
    {
      return value;
    }

    @Override
    public void setValue(String selectedItem)
    {
      value = selectedItem;
    }
  }

  private static class JobInput
  {
    List<SequenceI> msf;

    int[] delMap;

    Map<String, SequenceInfo> sequenceInfo;
  }

  public class JPredJob extends WSJob
  {
    List<SequenceI> msf;

    int[] delMap;

    Map<String, SequenceInfo> sequenceInfo;

    private JPredJob()
    {
      super(operation.service.getProviderName(), operation.getName(),
          operation.getHostName());
    }

    private void setInput(JobInput input)
    {
      msf = input.msf;
      delMap = input.delMap;
      sequenceInfo = input.sequenceInfo;
    }
  }

  public class PredictionResult
  {
    AlignmentI alignment;

    HiddenColumns hiddenCols;

    int firstSeq;

    public AlignmentI getAlignment()
    {
      return alignment;
    }

    public HiddenColumns getHiddenCols()
    {
      return hiddenCols;
    }
  }

  private JPredOperation operation;

  private Consumer<PredictionResult> resultConsumer;

  private AlignmentView view;

  private WSJobList<JPredJob> jobs = new WSJobList<>();

  private JPredJob job;

  private char gapChar;

  AlignmentI currentView;

  public JPredWorker(JPredOperation operation, AlignmentView alignView,
      AlignViewportI viewport)
  {
    this.operation = operation;
    this.view = alignView;
    this.gapChar = viewport.getGapCharacter();
    this.currentView = viewport.getAlignment();
  }

  @Override
  public Operation getOperation()
  {
    return operation;
  }

  @Override
  public WSJobList<? extends WSJob> getJobs()
  {
    return jobs;
  }

  public void setResultConsumer(Consumer<PredictionResult> consumer)
  {
    this.resultConsumer = consumer;
  }

  @Override
  public void start() throws IOException
  {
    var input = prepareInputData(view, true);
    job = new JPredJob();
    job.setInput(input);
    jobs.add(job);
    listeners.fireJobCreated(job);

    var formatArg = new InputFormatParameter();
    formatArg.setValue(input.msf.size() > 1 ? "fasta" : "seq");
    List<ArgumentI> args = List.of(formatArg);
    int exceptionCount = MAX_RETRY;
    String jobId = null;
    do
    {
      try
      {
        jobId = operation.getWebService().submit(job.msf, args);
      } catch (IOException e)
      {
        Cache.log.warn(format("%s failed to submit sequences to the server %s.",
            operation.getName(), operation.getHostName()), e);
        exceptionCount--;
      }
    } while (jobId == null && exceptionCount > 0);
    if (jobId != null)
    {
      job.setJobId(jobId);
      job.setStatus(WSJobStatus.SUBMITTED);
      listeners.fireWorkerStarted();
    }
    else
    {
      job.setStatus(WSJobStatus.SERVER_ERROR);
      listeners.fireWorkerNotStarted();
    }
  }

  private static JobInput prepareInputData(AlignmentView view, boolean viewOnly)
  {
    SeqCigar[] msf = view.getSequences();
    SequenceI seq = msf[0].getSeq('-');
    int[] delMap = null;
    if (viewOnly)
      delMap = view.getVisibleContigMapFor(seq.gapMap());
    SequenceI[] aln = new SequenceI[msf.length];
    for (int i = 0; i < msf.length; i++)
      aln[i] = msf[i].getSeq('-');
    var sequenceInfo = msf.length > 1 ? SeqsetUtils.uniquify(aln, true)
        : Map.of("Sequence", SeqsetUtils.SeqCharacterHash(seq));
    if (viewOnly)
    {
      // Remove hidden regions from sequence objects.
      String seqs[] = view.getSequenceStrings('-');
      for (int i = 0; i < msf.length; i++)
        aln[i].setSequence(seqs[i]);
      seq.setSequence(seqs[0]);
    }
    var input = new JobInput();
    input.msf = List.of(aln);
    input.delMap = delMap;
    input.sequenceInfo = sequenceInfo;
    return input;
  }

  @Override
  public void done()
  {
    listeners.fireWorkerCompleting();
    PredictionResult result = null;
    try
    {
      result = (job.msf.size() > 1)
          ? prepareMultipleSequenceResult(job)
          : prepareSingleSequenceResult(job);
    } catch (Exception e)
    {
      Cache.log.error("Couldn't retrieve results for job.", e);
      job.setStatus(WSJobStatus.SERVER_ERROR);
    }
    if (result != null)
    {
      for (var annot : result.alignment.getAlignmentAnnotation())
      {
        if (annot.sequenceRef != null)
        {
          replaceAnnotationOnAlignmentWith(annot, annot.label,
              getClass().getName(), annot.sequenceRef);
        }
      }
    }
    resultConsumer.accept(result);
    listeners.fireWorkerCompleted();
  }

  private PredictionResult prepareMultipleSequenceResult(JPredJob job)
      throws Exception
  {
    AlignmentI alignment;
    HiddenColumns hiddenCols = null;
    var prediction = operation.predictionSupplier.getPrediction(job);
    if (job.delMap != null)
    {
      Object[] alandcolsel = view.getAlignmentAndHiddenColumns(gapChar);
      alignment = new Alignment((SequenceI[]) alandcolsel[0]);
      hiddenCols = (HiddenColumns) alandcolsel[1];
    }
    else
    {
      alignment = operation.predictionSupplier.getAlignment(job);
      var seqs = new SequenceI[alignment.getHeight()];
      for (int i = 0; i < alignment.getHeight(); i++)
      {
        seqs[i] = alignment.getSequenceAt(i);
      }
      SeqsetUtils.deuniquify(job.sequenceInfo, seqs);
    }
    int firstSeq = 0;
    alignment.setDataset(currentView.getDataset());
    JnetAnnotationMaker.add_annotation(prediction, alignment, firstSeq, false,
        job.delMap);
    var result = new PredictionResult();
    result.alignment = alignment;
    result.hiddenCols = hiddenCols;
    result.firstSeq = firstSeq;
    return result;
  }

  static final int msaIndex = 0;

  private PredictionResult prepareSingleSequenceResult(JPredJob job)
      throws Exception
  {
    var prediction = operation.predictionSupplier.getPrediction(job);
    AlignmentI alignment = new Alignment(prediction.getSeqsAsArray());
    HiddenColumns hiddenCols = null;
    int firstSeq = prediction.getQuerySeqPosition();
    if (job.delMap != null)
    {
      Object[] alanndcolsel = view.getAlignmentAndHiddenColumns(gapChar);
      SequenceI[] seqs = (SequenceI[]) alanndcolsel[0];
      new RemoveGapsCommand(MessageManager.getString("label.remove_gaps"),
          new SequenceI[]
          { seqs[msaIndex] }, currentView);
      SequenceI profileSeq = alignment.getSequenceAt(firstSeq);
      profileSeq.setSequence(seqs[msaIndex].getSequenceAsString());
    }
    SeqsetUtils.SeqCharacterUnhash(alignment.getSequenceAt(firstSeq),
        job.sequenceInfo.get("Sequence"));
    alignment.setDataset(currentView.getDataset());
    JnetAnnotationMaker.add_annotation(prediction, alignment, firstSeq, true,
        job.delMap);
    SequenceI profileSeq = alignment.getSequenceAt(0);
    if (job.delMap != null)
    {
      hiddenCols = alignment.propagateInsertions(profileSeq, view);
    }
    var result = new PredictionResult();
    result.alignment = alignment;
    result.hiddenCols = hiddenCols;
    result.firstSeq = firstSeq;
    return result;
  }

  private static void replaceAnnotationOnAlignmentWith(
      AlignmentAnnotation newAnnot, String typeName, String calcId,
      SequenceI aSeq)
  {
    SequenceI dsseq = aSeq.getDatasetSequence();
    while (dsseq.getDatasetSequence() != null)
    {
      dsseq = dsseq.getDatasetSequence();
    }
    // look for same annotation on dataset and lift this one over
    List<AlignmentAnnotation> dsan = dsseq.getAlignmentAnnotations(calcId,
        typeName);
    if (dsan != null && dsan.size() > 0)
    {
      for (AlignmentAnnotation dssan : dsan)
      {
        dsseq.removeAlignmentAnnotation(dssan);
      }
    }
    AlignmentAnnotation dssan = new AlignmentAnnotation(newAnnot);
    dsseq.addAlignmentAnnotation(dssan);
    dssan.adjustForAlignment();
  }

}