package jalview.ws2.actions.annotation;

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

import jalview.analysis.AlignmentAnnotationUtils;
import jalview.api.AlignViewportI;
import jalview.api.FeatureColourI;
import jalview.datamodel.AlignmentAnnotation;
import jalview.datamodel.AlignmentI;
import jalview.datamodel.AnnotatedCollectionI;
import jalview.datamodel.Annotation;
import jalview.datamodel.ContiguousI;
import jalview.datamodel.Mapping;
import jalview.datamodel.SequenceI;
import jalview.datamodel.features.FeatureMatcherSetI;
import jalview.util.MapList;
import jalview.ws.params.ArgumentI;
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.AnnotationWebServiceClientI;

public class AnnotationTask extends BaseTask<AnnotationJob, AnnotationResult>
{
  private AnnotationWebServiceClientI client;

  private final AnnotationAction action;

  private final AlignViewportI viewport;

  public AnnotationTask(AnnotationWebServiceClientI client,
      AnnotationAction action, List<ArgumentI> args, Credentials credentials,
      AlignViewportI viewport)
  {
    super(client, args, credentials);
    this.client = client;
    this.action = action;
    this.viewport = viewport;
  }

  /**
   * Create and return a list of annotation jobs from the current state of the
   * viewport. Returned job are not started by this method and should be stored
   * in a field and started separately.
   * 
   * @return list of annotation jobs
   * @throws ServiceInputInvalidException
   *           input data is not valid
   */
  @Override
  public List<AnnotationJob> prepareJobs() throws ServiceInputInvalidException
  {
    AlignmentI alignment = viewport.getAlignment();
    if (alignment == null || alignment.getWidth() <= 0 ||
        alignment.getSequences() == null)
      throw new ServiceInputInvalidException("Alignment does not contain sequences");
    if (alignment.isNucleotide() && !action.doAllowNucleotide())
      throw new ServiceInputInvalidException(
          action.getFullName() + " does not allow nucleotide sequences");
    if (!alignment.isNucleotide() && !action.doAllowProtein())
      throw new ServiceInputInvalidException(
          action.getFullName() + " does not allow protein sequences");
    boolean bySequence = !action.isAlignmentAnalysis();
    AnnotatedCollectionI inputSeqs = bySequence ? viewport.getSelectionGroup() : null;
    if (inputSeqs == null || inputSeqs.getWidth() <= 0 ||
        inputSeqs.getSequences() == null || inputSeqs.getSequences().size() < 1)
      inputSeqs = alignment;
    boolean submitGaps = action.isAlignmentAnalysis();
    boolean requireAligned = action.getRequireAlignedSequences();
    boolean filterSymbols = action.getFilterSymbols();
    int minSize = action.getMinSequences();
    AnnotationJob job = AnnotationJob.create(inputSeqs, bySequence,
        submitGaps, requireAligned, filterSymbols, minSize);
    if (!job.isInputValid())
    {
      job.setStatus(JobStatus.INVALID);
      throw new ServiceInputInvalidException("Annotation job has invalid input");
    }
    job.setStatus(JobStatus.READY);
    return List.of(job);
  }

  @Override
  protected AnnotationResult collectResult(List<AnnotationJob> jobs) throws IOException
  {
    final Map<String, FeatureColourI> featureColours = new HashMap<>();
    final Map<String, FeatureMatcherSetI> featureFilters = new HashMap<>();
    var job = jobs.get(0);
    List<AlignmentAnnotation> returnedAnnot = client.attachAnnotations(
        job.getServerJob(), job.getInputSequences(), featureColours,
        featureFilters);
    /* TODO
     * copy over each annotation row returned and also defined on each
     * sequence, excluding regions not annotated due to gapMap/column
     * visibility */

    udpateCalcId(returnedAnnot);
    int graphGroup = viewport.getAlignment().getLastGraphGroup();
    shiftGraphGroup(returnedAnnot, graphGroup);
    List<AlignmentAnnotation> annotations = new ArrayList<>();
    for (AlignmentAnnotation ala : returnedAnnot)
    {
      SequenceI aseq = null;
      if (ala.sequenceRef != null) {
        SequenceI seq = job.seqNames.get(ala.sequenceRef.getName());
        aseq = seq.getRootDatasetSequence();
      }
      ala.sequenceRef = aseq;
      Annotation[] gappedAnnots = createGappedAnnotations(ala.annotations, job.start, job.gapMap);
      ala.annotations = gappedAnnots;

      AlignmentAnnotation newAnnot = viewport.getAlignment()
          .updateFromOrCopyAnnotation(ala);
      if (aseq != null)
      {
        aseq.addAlignmentAnnotation(newAnnot);
        newAnnot.adjustForAlignment();
        AlignmentAnnotationUtils.replaceAnnotationOnAlignmentWith(
            newAnnot, newAnnot.label, newAnnot.getCalcId());
      }
      annotations.add(newAnnot);
    }

    boolean hasFeatures = false;
    for (SequenceI sq : job.getInputSequences())
    {
      if (!sq.getFeatures().hasFeatures() && (sq.getDBRefs() == null || sq.getDBRefs().isEmpty()))
        continue;
      hasFeatures = true;
      SequenceI seq = job.seqNames.get(sq.getName());
      SequenceI datasetSeq = seq.getRootDatasetSequence();
      List<ContiguousI> sourceRange = findContiguousRanges(datasetSeq, job.gapMap, job.start, job.end);
      int[] sourceStartEnd = ContiguousI.toStartEndArray(sourceRange);
      Mapping mp = new Mapping(new MapList(
          sourceStartEnd, new int[]
          { datasetSeq.getStart(), datasetSeq.getEnd() }, 1, 1));
      datasetSeq.transferAnnotation(sq, mp);
    }

    return new AnnotationResult(annotations, hasFeatures, featureColours, featureFilters);
  }

  /**
   * Updates calcId on provided annotations if not already set.
   */
  public void udpateCalcId(Iterable<AlignmentAnnotation> annotations)
  {
    for (var annotation : annotations)
    {
      if (annotation.getCalcId() == null || annotation.getCalcId().isEmpty())
      {
        annotation.setCalcId(action.getFullName());
      }
      annotation.autoCalculated = action.isAlignmentAnalysis() &&
          action.getWebService().isInteractive();
    }
  }

  private static void shiftGraphGroup(Iterable<AlignmentAnnotation> annotations, int shift)
  {
    for (AlignmentAnnotation ala : annotations)
    {
      if (ala.graphGroup > 0)
      {
        ala.graphGroup += shift;
      }
    }
  }

  private Annotation[] createGappedAnnotations(Annotation[] annotations, int start, boolean[] gapMap)
  {
    var size = Math.max(viewport.getAlignment().getWidth(), gapMap.length);
    Annotation[] gappedAnnotations = new Annotation[size];
    for (int p = 0, ap = start; ap < size; ap++)
    {
      if (gapMap != null && gapMap.length > ap && !gapMap[ap])
      {
        gappedAnnotations[ap] = new Annotation("", "", ' ', Float.NaN);
      }
      else if (p < annotations.length)
      {
        gappedAnnotations[ap] = annotations[p++];
      }
    }
    return gappedAnnotations;
  }

  private List<ContiguousI> findContiguousRanges(SequenceI seq, boolean[] gapMap, int start, int end)
  {
    if (gapMap == null || gapMap.length < end)
      return List.of(seq.findPositions(start, end));
    List<ContiguousI> ranges = new ArrayList<>();
    int lastcol = start, col = start;
    do
    {
      if (col == end || !gapMap[col])
      {
        if (lastcol < col)
          ranges.add(seq.findPositions(lastcol, col));
        lastcol = col + 1;
      }
    } while (++col <= end);
    return ranges;
  }
}