#ifndef RUBY_TYPER_LSPTASK_H
#define RUBY_TYPER_LSPTASK_H

#include "absl/synchronization/notification.h"
#include "main/lsp/AbstractRewriter.h"
#include "main/lsp/LSPMessage.h"
#include "main/lsp/LSPTypechecker.h"
#include "main/lsp/json_types.h"

namespace sorbet::realmain::lsp {
class LSPIndexer;
class LSPPreprocessor;
class DocumentHighlight;

/**
 * A work unit that needs to execute on the typechecker thread. Subclasses implement `run`.
 * Contains miscellaneous helper methods that are useful in multiple tasks.
 *
 * NOTE: If `enableMultithreading` is set to `true`, then this task cannot preempt slow path typechecking.
 */
class LSPTask {
protected:
    const LSPConfiguration &config;

    // Task helper methods.

    std::vector<std::unique_ptr<core::lsp::QueryResponse>>
    getReferencesToSymbols(LSPTypecheckerDelegate &typechecker, core::lsp::Query::Symbol::STORAGE &&symbols) const;

    std::vector<std::unique_ptr<core::lsp::QueryResponse>>
    getReferencesToSymbolsInPackage(LSPTypecheckerDelegate &typechecker, core::packages::MangledName packageName,
                                    core::lsp::Query::Symbol::STORAGE &&symbols) const;

    std::vector<std::unique_ptr<core::lsp::QueryResponse>>
    getReferencesToSymbolsInFile(LSPTypecheckerDelegate &typechecker, core::FileRef file,
                                 core::lsp::Query::Symbol::STORAGE &&symbols) const;

    std::vector<std::unique_ptr<DocumentHighlight>>
    getHighlights(LSPTypecheckerDelegate &typechecker,
                  const std::vector<std::unique_ptr<core::lsp::QueryResponse>> &responses) const;
    void addLocIfExists(const core::GlobalState &gs, std::vector<std::unique_ptr<Location>> &locs, core::Loc loc) const;
    std::vector<std::unique_ptr<Location>>
    extractLocations(const core::GlobalState &gs,
                     const std::vector<std::unique_ptr<core::lsp::QueryResponse>> &queryResponses,
                     std::vector<std::unique_ptr<Location>> locations = {}) const;

    LSPTask(const LSPConfiguration &config, LSPMethod method);

public:
    virtual ~LSPTask();

    const LSPMethod method;

    // Measures the end-to-end latency of this task from the first message received from the client. Can be nullptr
    // if not relevant.
    std::unique_ptr<Timer> latencyTimer;

    enum class Phase {
        PREPROCESS = 1,
        INDEX = 2,
        RUN = 3,
    };

    // Get this task's method as a string that can be used in timers and counters.
    ConstExprStr methodString() const;

    /**
     * If `true`, this task can be delayed in favor of processing other tasks sooner. Defaults to `false`.
     */
    virtual bool isDelayable() const;

    // Attempts to cancel the task if it corresponds to the given request ID. Returns `true` if cancelation succeeds.
    // The default implementation returns `false`.
    virtual bool cancel(const MessageId &id);

    virtual bool canPreempt(const LSPIndexer &) const;

    // Returns the phase at which the task is complete. Some tasks only need to interface with the preprocessor or the
    // indexer. The default implementation returns RUN.
    virtual Phase finalPhase() const;

    // Some tasks, like request cancellations, need to interface with the preprocessor. The default implementation is
    // a no-op. Is only ever invoked from the preprocessor thread.
    virtual void preprocess(LSPPreprocessor &preprocessor);

    // Some tasks, like edits, need to interface with the indexer. Is only ever invoked from the processing
    // thread, and is guaranteed to be invoked exactly once. The default implementation is a no-op.
    // May be run from the processing thread (normally) or the typechecking thread (if preempting).
    virtual void index(LSPIndexer &indexer);

    // Runs the task. Is only ever invoked from the typechecker thread. Since it is exceedingly rare for a request to
    // not need to interface with the typechecker, this method must be implemented by all subclasses.
    virtual void run(LSPTypecheckerDelegate &typechecker) = 0;

    // The stratum that this task is runnable at, if running from the preemption scheduler. By default this returns the
    // last stratum, which is a conservatively correct choice for any task that supports preemption.
    virtual core::packages::Stratum preemptionStratum(FileStratumMapping info) const;
};

/**
 * A specialized version of LSPTask for LSP requests (which must be responded to).
 */
class LSPRequestTask : public LSPTask {
protected:
    const MessageId id;

    LSPRequestTask(const LSPConfiguration &config, MessageId id, LSPMethod method);

    virtual std::unique_ptr<ResponseMessage> runRequest(LSPTypecheckerDelegate &typechecker) = 0;

public:
    void run(LSPTypecheckerDelegate &typechecker) override;

    // Requests cannot override this method, as runRequest must be run (and it only runs during the RUN phase).
    Phase finalPhase() const override;

    bool cancel(const MessageId &id) override;

    core::packages::Stratum preemptionStratum(FileStratumMapping info) const override = 0;
};

// Doubles as the `methodString` for a `TextDocumentCompletion` LSPTask and also as
// the prefix for any metrics collected during completion itself.
#define LSP_COMPLETION_METRICS_PREFIX "textDocument.completion"

/**
 * A special form of LSPTask that has direct access to the typechecker and controls its own scheduling.
 * Is only used for slow path-related tasks. Do not use for anything else.
 */
class LSPDangerousTypecheckerTask : public LSPTask {
protected:
    LSPDangerousTypecheckerTask(const LSPConfiguration &config, LSPMethod method);

public:
    // Should never be called; throws an exception. May be overridden by tasks that can be dangerous or not dangerous.
    virtual void run(LSPTypecheckerDelegate &typechecker) override;
    // Performs the actual work on the task.
    virtual void runSpecial(LSPTypechecker &typechecker, WorkerPool &worker) = 0;
    // Tells the scheduler how long to wait before it can schedule more tasks.
    virtual void schedulerWaitUntilReady() = 0;
};

} // namespace sorbet::realmain::lsp

#endif
