package jalview.ws2.operations;

import static java.lang.String.format;

import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Hashtable;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletionStage;
import java.util.function.Consumer;

import javax.swing.JMenu;
import javax.swing.JMenuItem;
import javax.swing.ToolTipManager;

import org.eclipse.jetty.http.HttpGenerator.Result;

import jalview.analysis.AlignSeq;
import jalview.analysis.AlignmentSorter;
import jalview.analysis.SeqsetUtils;
import jalview.bin.Cache;
import jalview.datamodel.AlignedCodonFrame;
import jalview.datamodel.Alignment;
import jalview.datamodel.AlignmentI;
import jalview.datamodel.AlignmentOrder;
import jalview.datamodel.AlignmentView;
import jalview.datamodel.HiddenColumns;
import jalview.datamodel.SequenceI;
import jalview.datamodel.Sequence;
import jalview.gui.AlignFrame;
import jalview.gui.AlignViewport;
import jalview.gui.Desktop;
import jalview.gui.JvSwingUtils;
import jalview.gui.WebserviceInfo;
import jalview.gui.WsJobParameters;
import jalview.util.MathUtils;
import jalview.util.MessageManager;
import jalview.ws.params.ArgumentI;
import jalview.ws.params.ParamDatastoreI;
import jalview.ws.params.WsParamSetI;
import jalview.ws2.MenuEntryProviderI;
import jalview.ws2.ResultSupplier;
import jalview.ws2.WSJob;
import jalview.ws2.WSJobStatus;
import jalview.ws2.PollingTaskExecutor;
import jalview.ws2.WebServiceI;
import jalview.ws2.WebServiceInfoUpdater;
import jalview.ws2.WebServiceWorkerI;
import jalview.ws2.WebServiceWorkerListener;
import jalview.ws2.WebServiceWorkerListenersList;
import jalview.ws2.gui.AlignmentMenuBuilder;
import jalview.ws2.utils.WSJobList;

/**
 *
 * @author mmwarowny
 *
 */
public class AlignmentOperation implements Operation
{
  private final WebServiceI service;

  private final ResultSupplier<AlignmentI> supplier;
  

  public AlignmentOperation(
      WebServiceI service,
      ResultSupplier<AlignmentI> supplier)
  {
    this.service = service;
    this.supplier = supplier;
  }

  @Override
  public String getName()
  {
    return service.getName();
  }
  
  @Override
  public String getDescription()
  {
    return service.getDescription();
  }

  @Override
  public String getTypeName()
  {
    return "Multiple Sequence Alignment";
  }

  @Override
  public String getHostName()
  {
    return service.getHostName();
  }
  
  @Override
  public boolean hasParameters()
  {
    return service.hasParameters();
  }
  
  @Override
  public ParamDatastoreI getParamStore()
  {
    return service.getParamStore();
  }

  @Override
  public int getMinSequences()
  {
    return 2;
  }

  @Override
  public int getMaxSequences()
  {
    return Integer.MAX_VALUE;
  }

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

  @Override
  public boolean isNucleotideOperation()
  {
    return true;
  }
  
  @Override
  public boolean isAlignmentAnalysis()
  {
    return false;
  }

  @Override
  public boolean canSubmitGaps()
  {
    // hack copied from original jabaws code, don't blame me
    return getName().contains("lustal");
  }

  @Override
  public boolean isInteractive()
  {
    return false;
  }
  
  @Override
  public boolean getFilterNonStandardSymbols()
  {
    return true;
  }
  
  @Override
  public boolean getNeedsAlignedSequences()
  {
    return false;
  }

  @Override
  public MenuEntryProviderI getMenuBuilder()
  {
    return new AlignmentMenuBuilder(this);
  }


  /**
   * Implementation of the web service worker performing multiple sequence
   * alignment.
   *
   * @author mmwarowny
   *
   */
  public class AlignmentWorker implements WebServiceWorkerI
  {

    private long uid = MathUtils.getUID();

    private final AlignmentView msa;

    private final AlignmentI dataset;

    private final AlignViewport viewport;

    private final List<AlignedCodonFrame> codonFrame = new ArrayList<>();

    private List<ArgumentI> args = Collections.emptyList();

    private String alnTitle = "";

    private boolean submitGaps = false;

    private boolean preserveOrder = false;

    private char gapCharacter;

    private WSJobList jobs = new WSJobList();

    private Map<Long, JobInput> inputs = new LinkedHashMap<>();

    private Map<Long, Integer> exceptionCount = new HashMap<>();
    
    private final int MAX_RETRY = 5;

    public AlignmentWorker(AlignmentView msa, List<ArgumentI> args,
            String alnTitle, boolean submitGaps, boolean preserveOrder,
            AlignViewport viewport)
    {
      this.msa = msa;
      this.dataset = viewport.getAlignment().getDataset();
      List<AlignedCodonFrame> cf = Objects.requireNonNullElse(
          viewport.getAlignment().getCodonFrames(), Collections.emptyList());
      this.codonFrame.addAll(cf);
      this.args = args;
      this.alnTitle = alnTitle;
      this.submitGaps = submitGaps;
      this.preserveOrder = preserveOrder;
      this.viewport = viewport;
      this.gapCharacter = viewport.getGapCharacter();
    }

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

    @Override
    public WebServiceI getWebService()
    {
      return service;
    }

    @Override
    public WSJobList getJobs()
    {
      return jobs;
    }

    @Override
    public void start() throws IOException
    {
      Cache.log.info(format("Starting new %s job.", service.getName()));
      SequenceI[][] conmsa = msa.getVisibleContigs('-');
      if (conmsa == null)
      {
        return;
      }
      int numValid = 0;
      for (int i = 0; i < conmsa.length; i++)
      {
        JobInput input = JobInput.create(conmsa[i], 2, submitGaps);
        WSJob job = new WSJob(service.getProviderName(), service.getName(),
                service.getHostName());
        job.setJobNum(i);
        inputs.put(job.getUid(), input);
        jobs.add(job);
        listeners.fireJobCreated(job);
        if (input.isInputValid())
        {
          int count;
          String jobId = null;
          do
          {
            count = exceptionCount.getOrDefault(job.getUid(), MAX_RETRY);
            try
            {
              jobId = service.submit(input.inputSequences, args);
              Cache.log.debug((format("Job %s submitted", job)));
              exceptionCount.remove(job.getUid());
            } catch (IOException e)
            {
              exceptionCount.put(job.getUid(), --count);
            }
          } while (jobId == null && count > 0);
          if (jobId != null)
          {
            job.setJobId(jobId);
            job.setStatus(WSJobStatus.SUBMITTED);
            numValid++;
          }
          else
          {
            job.setStatus(WSJobStatus.SERVER_ERROR);
          }
        }
        else
        {
          job.setStatus(WSJobStatus.INVALID);
          job.setErrorLog(
                  MessageManager.getString("label.empty_alignment_job"));
        }
      }
      if (numValid > 0)
      {
        listeners.fireWorkerStarted();
      }
      else
      {
        listeners.fireWorkerNotStarted();
      }
    }

    @Override
    public boolean poll()
    {
      boolean done = true;
      for (WSJob job : getJobs())
      {
        if (!job.getStatus().isDone() && !job.getStatus().isFailed())
        {
          Cache.log.debug(format("Polling job %s.", job));
          try
          {
            service.updateProgress(job);
            exceptionCount.remove(job.getUid());
          } catch (IOException e)
          {
            Cache.log.error(format("Polling job %s failed.", job), e);
            listeners.firePollException(job, e);
            int count = exceptionCount.getOrDefault(job.getUid(),
                    MAX_RETRY);
            if (--count <= 0)
            {
              job.setStatus(WSJobStatus.SERVER_ERROR);
              Cache.log.warn(format(
                      "Attempts limit exceeded. Droping job %s.", job));
            }
            exceptionCount.put(job.getUid(), count);
          } catch (OutOfMemoryError e)
          {
            job.setStatus(WSJobStatus.BROKEN);
            Cache.log.error(
                    format("Out of memory when retrieving job %s", job), e);
          }
          Cache.log.debug(
                  format("Job %s status is %s", job, job.getStatus()));
        }
        done &= job.getStatus().isDone() || job.getStatus().isFailed();
      }
      return done;
    }


    @Override
    public void done()
    {
      listeners.fireWorkerCompleting();
      Map<Long, AlignmentI> results = new LinkedHashMap<>();
      for (WSJob job : getJobs())
      {
        if (job.getStatus().isFailed())
          continue;
        try
        {
          AlignmentI alignment = supplier.getResult(job, 
              dataset.getSequences(), viewport);
          if (alignment != null)
          {
            results.put(job.getUid(), alignment);
          }
        } catch (Exception e)
        {
          if (!service.handleCollectionError(job, e))
          {
            Cache.log.error("Couldn't get alignment for job.", e);
            // TODO: Increment exception count and retry.
            job.setStatus(WSJobStatus.SERVER_ERROR);
          }
        }
      }
      if (results.size() > 0)
      {
        AlignmentResult out = prepareResult(results);
        resultConsumer.accept(out);
      }
      else
      {
        resultConsumer.accept(null);
      }
      listeners.fireWorkerCompleted();
    }

    private AlignmentResult prepareResult(Map<Long, AlignmentI> alignments)
    {
      List<AlignmentOrder> alorders = new ArrayList<>();
      SequenceI[][] results = new SequenceI[jobs.size()][];
      AlignmentOrder[] orders = new AlignmentOrder[jobs.size()];
      for (int i = 0; i < jobs.size(); i++)
      {
        WSJob job = jobs.get(i);
        AlignmentI aln = alignments.get(job.getUid());
        if (aln != null) // equivalent of job.hasResults()
        {
          /* Get the alignment including any empty sequences in the original
           * order with original ids. */
          JobInput input = inputs.get(job.getUid());
          char gapChar = aln.getGapCharacter();
          List<SequenceI> emptySeqs = input.emptySequences;
          List<SequenceI> alnSeqs = aln.getSequences();
          // find the width of the longest sequence
          int width = 0;
          for (var seq : alnSeqs)
            width = Integer.max(width, seq.getLength());
          for (var emptySeq : emptySeqs)
            width = Integer.max(width, emptySeq.getLength());
          // pad shorter sequences with gaps
          String gapSeq = String.join("",
                  Collections.nCopies(width, Character.toString(gapChar)));
          List<SequenceI> seqs = new ArrayList<>(
                  alnSeqs.size() + emptySeqs.size());
          seqs.addAll(alnSeqs);
          seqs.addAll(emptySeqs);
          for (var seq : seqs)
          {
            if (seq.getLength() < width)
              seq.setSequence(seq.getSequenceAsString()
                      + gapSeq.substring(seq.getLength()));
          }
          SequenceI[] result = seqs.toArray(new SequenceI[0]);
          AlignmentOrder msaOrder = new AlignmentOrder(result);
          AlignmentSorter.recoverOrder(result);
          // temporary workaround for deuniquify
          @SuppressWarnings({ "rawtypes", "unchecked" })
          Hashtable names = new Hashtable(input.sequenceNames);
          // FIXME first call to deuniquify alters original alignment
          SeqsetUtils.deuniquify(names, result);
          alorders.add(msaOrder);
          results[i] = result;
          orders[i] = msaOrder;
        }
        else
        {
          results[i] = null;
        }
      }

      Object[] newView = msa.getUpdatedView(results, orders, gapCharacter);
      // free references to original data
      for (int i = 0; i < jobs.size(); i++)
      {
        results[i] = null;
        orders[i] = null;
      }
      SequenceI[] alignment = (SequenceI[]) newView[0];
      HiddenColumns hidden = (HiddenColumns) newView[1];
      Alignment aln = new Alignment(alignment);
      aln.setProperty("Alignment Program", service.getName());
      if (dataset != null)
        aln.setDataset(dataset);

      propagateDatasetMappings(aln);
      return new AlignmentResult(aln, alorders, hidden);
      // displayNewFrame(aln, alorders, hidden);
    }

    /*
     * conserves dataset references to sequence objects returned from web
     * services. propagate codon frame data to alignment.
     */
    private void propagateDatasetMappings(Alignment aln)
    {
      if (codonFrame != null)
      {
        SequenceI[] alignment = aln.getSequencesArray();
        for (SequenceI seq : alignment)
        {
          for (AlignedCodonFrame acf : codonFrame)
          {
            if (acf != null && acf.involvesSequence(seq))
            {
              aln.addCodonFrame(acf);
              break;
            }
          }
        }
      }
    }
    
    private Consumer<AlignmentResult> resultConsumer;
    
    public void setResultConsumer(Consumer<AlignmentResult> consumer)
    {
      this.resultConsumer = consumer;
    }

    private WebServiceWorkerListenersList listeners = 
        new WebServiceWorkerListenersList(this);
    
    @Override
    public void addListener(WebServiceWorkerListener listener)
    {
      listeners.addListener(listener);
    }
  }
  
  public class AlignmentResult
  {
    AlignmentI aln;

    List<AlignmentOrder> alorders;

    HiddenColumns hidden;

    AlignmentResult(AlignmentI aln, List<AlignmentOrder> alorders,
            HiddenColumns hidden)
    {
      this.aln = aln;
      this.alorders = alorders;
      this.hidden = hidden;
    }

    public AlignmentI getAln()
    {
      return aln;
    }

    public List<AlignmentOrder> getAlorders()
    {
      return alorders;
    }

    public HiddenColumns getHidden()
    {
      return hidden;
    }
  }

  private static class JobInput
  {
    final List<SequenceI> inputSequences;

    final List<SequenceI> emptySequences;

    @SuppressWarnings("rawtypes")
    final Map<String, ? extends Map> sequenceNames;

    private JobInput(int numSequences, List<SequenceI> inputSequences,
            List<SequenceI> emptySequences,
            @SuppressWarnings("rawtypes") Map<String, ? extends Map> names)
    {
      this.inputSequences = Collections.unmodifiableList(inputSequences);
      this.emptySequences = Collections.unmodifiableList(emptySequences);
      this.sequenceNames = names;
    }

    boolean isInputValid()
    {
      return inputSequences.size() >= 2;
    }

    static JobInput create(SequenceI[] sequences, int minLength,
            boolean submitGaps)
    {
      assert minLength >= 0 : MessageManager.getString(
              "error.implementation_error_minlen_must_be_greater_zero");
      int numSeq = 0;
      for (SequenceI seq : sequences)
      {
        if (seq.getEnd() - seq.getStart() >= minLength)
        {
          numSeq++;
        }
      }

      List<SequenceI> inputSequences = new ArrayList<>();
      List<SequenceI> emptySequences = new ArrayList<>();
      @SuppressWarnings("rawtypes")
      Map<String, Hashtable> names = new LinkedHashMap<>();
      for (int i = 0; i < sequences.length; i++)
      {
        SequenceI seq = sequences[i];
        String newName = SeqsetUtils.unique_name(i);
        @SuppressWarnings("rawtypes")
        Hashtable hash = SeqsetUtils.SeqCharacterHash(seq);
        names.put(newName, hash);
        if (numSeq > 1 && seq.getEnd() - seq.getStart() >= minLength)
        {
          String seqString = seq.getSequenceAsString();
          if (!submitGaps)
          {
            seqString = AlignSeq.extractGaps(
                    jalview.util.Comparison.GapChars, seqString);
          }
          inputSequences.add(new Sequence(newName, seqString));
        }
        else
        {
          String seqString = "";
          if (seq.getEnd() >= seq.getStart())  // true if gaps only
          {
            seqString = seq.getSequenceAsString();
            if (!submitGaps)
            {
              seqString = AlignSeq.extractGaps(
                      jalview.util.Comparison.GapChars, seqString);
            }
          }
          emptySequences.add(new Sequence(newName, seqString));
        }
      }

      return new JobInput(numSeq, inputSequences, emptySequences, names);
    }
  }

}