package jalview.ws2.actions.annotation;

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

import jalview.analysis.AlignmentAnnotationUtils;
import jalview.api.AlignCalcManagerI2;
import jalview.api.AlignCalcWorkerI;
import jalview.api.AlignViewportI;
import jalview.api.FeatureColourI;
import jalview.api.PollableAlignCalcWorkerI;
import jalview.bin.Cache;
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.schemes.FeatureSettingsAdapter;
import jalview.util.ArrayUtils;
import jalview.util.MapList;
import jalview.util.MathUtils;
import jalview.util.Pair;
import jalview.workers.AlignCalcWorker;
import jalview.ws.params.ArgumentI;
import jalview.ws2.actions.AbstractPollableTask;
import jalview.ws2.actions.BaseJob;
import jalview.ws2.actions.ServiceInputInvalidException;
import jalview.ws2.actions.api.JobI;
import jalview.ws2.actions.api.TaskEventListener;
import jalview.ws2.actions.api.TaskI;
import jalview.ws2.api.Credentials;
import jalview.ws2.api.JobStatus;
import jalview.ws2.api.WebServiceJobHandle;
import jalview.ws2.client.api.AnnotationWebServiceClientI;
import jalview.ws2.helpers.DelegateJobEventListener;
import jalview.ws2.helpers.TaskEventSupport;

import static java.util.Objects.requireNonNullElse;

public class AnnotationTask implements TaskI<AnnotationResult>
{
  private final long uid = MathUtils.getUID();

  private AnnotationWebServiceClientI client;

  private final AnnotationAction action;

  private final List<ArgumentI> args;

  private final Credentials credentials;

  private final AlignViewportI viewport;

  private final TaskEventSupport<AnnotationResult> eventHandler;

  private JobStatus taskStatus = null;

  private AlignCalcWorkerAdapter worker = null;

  private List<AnnotationJob> jobs = Collections.emptyList();

  private AnnotationResult result = null;

  private DelegateJobEventListener<AnnotationResult> jobEventHandler;

  private class AlignCalcWorkerAdapter extends AlignCalcWorker
      implements PollableAlignCalcWorkerI
  {
    private boolean restarting = false;

    AlignCalcWorkerAdapter(AlignCalcManagerI2 calcMan)
    {
      super(viewport, null);
      this.calcMan = calcMan;
    }

    String getServiceName()
    {
      return action.getWebService().getName();
    }

    @Override
    public void startUp() throws Throwable
    {
      if (alignViewport.isClosed())
      {
        stop();
        throw new IllegalStateException("Starting annotation for closed viewport");
      }
      if (restarting)
        eventHandler.fireTaskRestarted();
      else
        restarting = true;
      jobs = Collections.emptyList();
      try
      {
        jobs = prepare();
      } catch (ServiceInputInvalidException e)
      {
        setStatus(JobStatus.INVALID);
        eventHandler.fireTaskException(e);
        throw e;
      }
      setStatus(JobStatus.READY);
      eventHandler.fireTaskStarted(jobs);
      for (var job : jobs)
      {
        job.addPropertyChagneListener(jobEventHandler);
      }
      try
      {
        startJobs();
      } catch (IOException e)
      {
        eventHandler.fireTaskException(e);
        cancelJobs();
        setStatus(JobStatus.SERVER_ERROR);
        throw e;
      }
      setStatus(JobStatus.SUBMITTED);
    }

    @Override
    public boolean poll() throws Throwable
    {
      boolean done = AnnotationTask.this.poll();
      updateGlobalStatus();
      if (done)
      {
        retrieveAndProcessResult();
        eventHandler.fireTaskCompleted(result);
      }
      return done;
    }

    private void retrieveAndProcessResult() throws IOException
    {
      result = retrieveResult();
      updateOurAnnots(result.annotations);
      if (result.transferFeatures)
      {
        final var featureColours = result.featureColours;
        final var featureFilters = result.featureFilters;
        viewport.applyFeaturesStyle(new FeatureSettingsAdapter()
        {
          @Override
          public FeatureColourI getFeatureColour(String type)
          {
            return featureColours.get(type);
          }

          @Override
          public FeatureMatcherSetI getFeatureFilters(String type)
          {
            return featureFilters.get(type);
          }

          @Override
          public boolean isFeatureDisplayed(String type)
          {
            return featureColours.containsKey(type);
          }
        });
      }
    }

    @Override
    public void updateAnnotation()
    {
      var job = jobs.size() > 0 ? jobs.get(0) : null;
      if (!calcMan.isWorking(this) && job != null)
      {
        var ret = updateResultAnnotation(job, job.returnedAnnotations);
        updateOurAnnots(ret.get0());
      }
    }

    private void updateOurAnnots(List<AlignmentAnnotation> newAnnots)
    {
      List<AlignmentAnnotation> oldAnnots = requireNonNullElse(ourAnnots, Collections.emptyList());
      ourAnnots = newAnnots;
      AlignmentI alignment = viewport.getAlignment();
      for (AlignmentAnnotation an : oldAnnots)
      {
        if (!newAnnots.contains(an))
        {
          alignment.deleteAnnotation(an);
        }
      }
      oldAnnots.clear();
      for (AlignmentAnnotation an : ourAnnots)
      {
        viewport.getAlignment().validateAnnotation(an);
      }
    }

    @Override
    public void cancel()
    {
      cancelJobs();
    }

    void stop()
    {
      calcMan.disableWorker(this);
      super.abortAndDestroy();
    }

    @Override
    public void done()
    {
      for (var job : jobs)
      {
        if (job.isInputValid() && !job.isCompleted())
        {
          /* if done was called but job is not completed then it
           * must have been stopped by an exception */
          job.setStatus(JobStatus.SERVER_ERROR);
        }
      }
      updateGlobalStatus();
      // dispose of unfinished jobs just in case
      cancelJobs();
    }

    @Override
    public String toString()
    {
      return AnnotationTask.this.toString() + "$AlignCalcWorker@"
          + Integer.toHexString(hashCode());
    }
  }

  public AnnotationTask(AnnotationWebServiceClientI client,
      AnnotationAction action, List<ArgumentI> args, Credentials credentials,
      AlignViewportI viewport,
      TaskEventListener<AnnotationResult> eventListener)
  {
    this.client = client;
    this.action = action;
    this.args = args;
    this.credentials = credentials;
    this.viewport = viewport;
    this.eventHandler = new TaskEventSupport<>(this, eventListener);
    this.jobEventHandler = new DelegateJobEventListener<>(this.eventHandler);
  }

  @Override
  public long getUid()
  {
    return uid;
  }

  public void start(AlignCalcManagerI2 calcManager)
  {
    if (this.worker != null)
      throw new IllegalStateException("task already started");
    this.worker = new AlignCalcWorkerAdapter(calcManager);
    if (taskStatus != JobStatus.CANCELLED)
    {
      List<AlignCalcWorkerI> oldWorkers = calcManager.getWorkersOfClass(
          AlignCalcWorkerAdapter.class);
      for (var worker : oldWorkers)
      {
        if (action.getWebService().getName().equalsIgnoreCase(
            ((AlignCalcWorkerAdapter) worker).getServiceName()))
        {
          // remove interactive workers for the same service.
          calcManager.removeWorker(worker);
          calcManager.cancelWorker(worker);
        }
      }
      if (action.getWebService().isInteractive())
        calcManager.registerWorker(worker);
      else
        calcManager.startWorker(worker);
    }
  }

  /*
   * The following methods are mostly copied from the {@link AbstractPollableTask}
   * TODO: move common functionality to a base class
   */
  @Override
  public JobStatus getStatus()
  {
    return taskStatus;
  }

  private void setStatus(JobStatus status)
  {
    if (this.taskStatus != status)
    {
      Cache.log.debug(String.format("%s status change to %s", this, status.name()));
      this.taskStatus = status;
      eventHandler.fireTaskStatusChanged(status);
    }
  }

  private void updateGlobalStatus()
  {
    int precedence = -1;
    for (BaseJob job : jobs)
    {
      JobStatus status = job.getStatus();
      int jobPrecedence = ArrayUtils.indexOf(JobStatus.statusPrecedence, status);
      if (precedence < jobPrecedence)
        precedence = jobPrecedence;
    }
    if (precedence >= 0)
    {
      setStatus(JobStatus.statusPrecedence[precedence]);
    }
  }

  @Override
  public List<? extends JobI> getSubJobs()
  {
    return jobs;
  }

  /**
   * 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
   */
  private List<AnnotationJob> prepare() 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);
  }

  private void startJobs() throws IOException
  {
    for (BaseJob job : jobs)
    {
      if (job.isInputValid() && job.getStatus() == JobStatus.READY)
      {
        var serverJob = client.submit(job.getInputSequences(),
            args, credentials);
        job.setServerJob(serverJob);
        job.setStatus(JobStatus.SUBMITTED);
      }
    }
  }

  private boolean poll() throws IOException
  {
    boolean allDone = true;
    for (BaseJob job : jobs)
    {
      if (job.isInputValid() && !job.getStatus().isDone())
      {
        WebServiceJobHandle serverJob = job.getServerJob();
        job.setStatus(client.getStatus(serverJob));
        job.setLog(client.getLog(serverJob));
        job.setErrorLog(client.getErrorLog(serverJob));
      }
      allDone &= job.isCompleted();
    }
    return allDone;
  }

  private AnnotationResult retrieveResult() 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 */

    // update calcId if it is not already set on returned annotation
    for (AlignmentAnnotation annot : returnedAnnot)
    {
      if (annot.getCalcId() == null || annot.getCalcId().isEmpty())
      {
        annot.setCalcId(action.getFullName());
      }
      annot.autoCalculated = action.isAlignmentAnalysis() &&
          action.getWebService().isInteractive();
    }
    job.returnedAnnotations = returnedAnnot;
    job.featureColours = featureColours;
    job.featureFilters = featureFilters;
    var ret = updateResultAnnotation(job, returnedAnnot);
    var annotations = ret.get0();
    var transferFeatures = ret.get1();
    return new AnnotationResult(annotations, transferFeatures, featureColours,
        featureFilters);
  }

  private Pair<List<AlignmentAnnotation>, Boolean> updateResultAnnotation(
      AnnotationJob job, List<AlignmentAnnotation> annotations)
  {
    List<AlignmentAnnotation> newAnnots = new ArrayList<>();
    // update graphGroup for all annotation
    /* find a graphGroup greater than any existing one, could be moved
     * to Alignment#getNewGraphGroup() - returns next unused graph group */
    int graphGroup = 1;
    if (viewport.getAlignment().getAlignmentAnnotation() != null)
    {
      for (var ala : viewport.getAlignment().getAlignmentAnnotation())
      {
        graphGroup = Math.max(graphGroup, ala.graphGroup);
      }
    }
    // update graphGroup in the annotation rows returned form service'
    /* TODO: look at sequence annotation rows and update graph groups in the
     * case of reference annotation */
    for (AlignmentAnnotation ala : annotations)
    {
      if (ala.graphGroup > 0)
        ala.graphGroup += graphGroup;
      SequenceI aseq = null;
      // transfer sequence refs and adjust gapMap
      if (ala.sequenceRef != null)
      {
        aseq = job.seqNames.get(ala.sequenceRef.getName());
      }
      ala.sequenceRef = aseq;

      Annotation[] resAnnot = ala.annotations;
      boolean[] gapMap = job.gapMap;
      Annotation[] gappedAnnot = new Annotation[Math.max(
          viewport.getAlignment().getWidth(), gapMap.length)];
      for (int p = 0, ap = job.start; ap < gappedAnnot.length; ap++)
      {
        if (gapMap.length > ap && !gapMap[ap])
          gappedAnnot[ap] = new Annotation("", "", ' ', Float.NaN);
        else if (p < resAnnot.length)
          gappedAnnot[ap] = resAnnot[p++];
        // is this loop exhaustive of resAnnot?
      }
      ala.annotations = gappedAnnot;

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

    boolean transferFeatures = false;
    for (SequenceI sq : job.getInputSequences())
    {
      if (!sq.getFeatures().hasFeatures() &&
          (sq.getDBRefs() == null || sq.getDBRefs().size() == 0))
        continue;
      transferFeatures = true;
      SequenceI seq = job.seqNames.get(sq.getName());
      SequenceI dseq;
      int start = job.start, end = job.end;
      boolean[] gapMap = job.gapMap;
      ContiguousI seqRange = seq.findPositions(start, end);
      while ((dseq = seq).getDatasetSequence() != null)
      {
        seq = seq.getDatasetSequence();
      }
      List<ContiguousI> sourceRange = new ArrayList<>();
      if (gapMap.length >= end)
      {
        int lastcol = start, col = start;
        do
        {
          if (col == end || !gapMap[col])
          {
            if (lastcol <= (col - 1))
            {
              seqRange = seq.findPositions(lastcol, col);
              sourceRange.add(seqRange);
            }
            lastcol = col + 1;
          }
        } while (col++ < end);
      }
      else
      {
        sourceRange.add(seq.findPositions(start, end));
      }

      int i = 0;
      int sourceStartEnd[] = new int[sourceRange.size() * 2];
      for (ContiguousI range : sourceRange)
      {
        sourceStartEnd[i++] = range.getBegin();
        sourceStartEnd[i++] = range.getEnd();
      }
      Mapping mp = new Mapping(new MapList(
          sourceStartEnd, new int[]
          { seq.getStart(), seq.getEnd() }, 1, 1));
      dseq.transferAnnotation(sq, mp);
    }

    return new Pair<>(newAnnots, transferFeatures);
  }

  @Override
  public AnnotationResult getResult()
  {
    return result;
  }

  @Override
  public void cancel()
  {
    setStatus(JobStatus.CANCELLED);
    if (worker != null)
    {
      worker.stop();
    }
    cancelJobs();
  }

  public void cancelJobs()
  {
    for (BaseJob job : jobs)
    {
      if (!job.isCompleted())
      {
        try
        {
          if (job.getServerJob() != null)
          {
            client.cancel(job.getServerJob());
          }
          job.setStatus(JobStatus.CANCELLED);
        } catch (IOException e)
        {
          Cache.log.error(String.format(
              "failed to cancel job %s", job.getServerJob()), e);
        }
      }
    }
  }

  @Override
  public String toString()
  {
    var status = taskStatus != null ? taskStatus.name() : "UNSET";
    return String.format("%s(%x, %s)", getClass().getSimpleName(), uid, status);
  }
}