package jalview.ws2.actions.secstructpred;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

import jalview.analysis.AlignmentAnnotationUtils;
import jalview.analysis.SeqsetUtils;
import jalview.analysis.SeqsetUtils.SequenceInfo;
import jalview.api.AlignViewportI;
import jalview.bin.Console;
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.io.AlignFile;
import jalview.io.JPredFile;
import jalview.io.JnetAnnotationMaker;
import jalview.ws.params.ArgumentI;
import jalview.ws2.actions.BaseJob;
import jalview.ws2.actions.BaseTask;
import jalview.ws2.actions.ServiceInputInvalidException;
import jalview.ws2.api.Credentials;
import jalview.ws2.api.JobStatus;
import jalview.ws2.client.api.SecStructPredWebServiceClientI;

public class SecStructPredMsaTask
        extends BaseTask<SecStructPredMsaTask.SecStructPredJob, AlignmentI>
{
  private final SecStructPredWebServiceClientI client;

  private final AlignmentView alignmentView;

  private final AlignmentI currentView;

  private final char gapChar;

  SecStructPredMsaTask(SecStructPredWebServiceClientI client,
          List<ArgumentI> args, Credentials credentials,
          AlignViewportI viewport)
  {
    super(client, args, credentials);
    this.client = client;
    this.alignmentView = viewport.getAlignmentView(true);
    this.currentView = viewport.getAlignment();
    this.gapChar = viewport.getGapCharacter();
  }

  @Override
  protected List<SecStructPredJob> prepareJobs()
          throws ServiceInputInvalidException
  {
    SeqCigar[] msf = alignmentView.getSequences();
    SequenceI referenceSeq = msf[0].getSeq('-');
    int[] delMap = alignmentView
            .getVisibleContigMapFor(referenceSeq.gapMap());

    // TODO: assume MSA for now
    SequenceI[] sequences = new SequenceI[msf.length];
    for (int i = 0; i < msf.length; i++)
      sequences[i] = msf[i].getSeq('-');
    var sequenceInfo = SeqsetUtils.uniquify(sequences, true);
    referenceSeq.setSequence(alignmentView.getASequenceString('-', 0));
    for (int i = 0; i < sequences.length; i++)
      sequences[i].setSequence(alignmentView.getASequenceString('-', i));
    var nonEmptySeqs = SeqsetUtils.getNonEmptySequenceSet(sequences);
    var job = new SecStructPredJob(Arrays.asList(nonEmptySeqs),
            referenceSeq, delMap, sequenceInfo);
    job.setStatus(JobStatus.READY);
    return List.of(job);
  }

  @Override
  protected AlignmentI collectResult(List<SecStructPredJob> jobs)
          throws IOException
  {
    var job = jobs.get(0); // There shouldn't be more than one job
    var status = job.getStatus();
    Console.info(
            String.format("sec str pred job \"%s\" finished with status %s",
                    job.getServerJob().getJobId(), status));
    if (status != JobStatus.COMPLETED)
      return null;
    JPredFile predictionFile = client.getPredictionFile(job.getServerJob());
    AlignFile alignmentFile = client.getAlignmentFile(job.getServerJob());

    Object[] alnAndHiddenCols = alignmentView
            .getAlignmentAndHiddenColumns(gapChar);
    Alignment aln = new Alignment((SequenceI[]) alnAndHiddenCols[0]);
    aln.setDataset(currentView.getDataset());
    aln.setHiddenColumns((HiddenColumns) alnAndHiddenCols[1]);
    try
    {
      JnetAnnotationMaker.add_annotation(predictionFile, aln, 0, false,
              job.getDelMap());
    } catch (Exception e)
    {
      throw new IOException(e);
    }

    for (AlignmentAnnotation alnAnnot : aln.getAlignmentAnnotation())
    {
      if (alnAnnot.sequenceRef != null)
      {
        AlignmentAnnotationUtils.replaceAnnotationOnAlignmentWith(alnAnnot,
                alnAnnot.label, getClass().getSimpleName());
      }
    }
    aln.setSeqrep(aln.getSequenceAt(0));
    return aln;
  }

  public static class SecStructPredJob extends BaseJob
  {
    private final int[] delMap;

    private final SequenceI refSeq;

    final Map<String, SequenceInfo> seqNames;

    SecStructPredJob(List<SequenceI> sequences, SequenceI refSeq,
            int[] delMap, Map<String, SequenceInfo> seqNames)
    {
      super(sequences);
      this.refSeq = refSeq;
      this.delMap = delMap;
      this.seqNames = seqNames;
    }

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

    public int[] getDelMap()
    {
      return this.delMap;
    }
  }

}