#include "main/lsp/LSPTask.h"
#include "absl/strings/match.h"
#include "absl/synchronization/notification.h"
#include "common/sort/sort.h"
#include "core/FileHash.h"
#include "core/lsp/QueryResponse.h"
#include "main/lsp/LSPLoop.h"
#include "main/lsp/LSPOutput.h"
#include "main/lsp/LSPQuery.h"
#include "main/lsp/json_types.h"

namespace sorbet::realmain::lsp {
using namespace std;

namespace {

// Enum used for tracking success/failure for requests. Basically only used to increment the
// respective counter
enum class ResponseMessageStatus {
    // Unknown status (will not increment any counter)
    Unknown = 0,
    // Request errored, and returned a RPC error (ResponseMessage::error)
    Errored,
    // Request succeeded, but return a known "empty" result for that response type.
    // (This could be a null result, an empty list, etc. depending on the response type)
    EmptyResult,
    // Request succeeded, and returned a meaningful result.
    Succeeded,
};

ResponseMessageStatus statusForResponse(const ResponseMessage &response) {
    if (response.error.has_value()) {
        ENFORCE(response.error.value() != nullptr);
        return ResponseMessageStatus::Errored;
    } else if (response.result.has_value()) {
        return visit(
            [](auto &&res) -> ResponseMessageStatus {
                using T = decay_t<decltype(res)>;
                // Note: Many of these are Unknown just out of laziness / not needing or wanting to
                // track them. Feel free to implement any case here more sensibly.
                if constexpr (is_same_v<T, unique_ptr<SorbetCounters>>) {
                    // __GETCOUNTERS__
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, unique_ptr<InitializeResult>>) {
                    // initialize
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, JSONNullObject>) {
                    // shutdown
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, vector<unique_ptr<DocumentHighlight>>>>) {
                    // textDocument/documentHighlight
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, vector<unique_ptr<DocumentSymbol>>>>) {
                    // textDocument/documentSymbol
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, vector<unique_ptr<Location>>>>) {
                    // textDocument/definition, textDocument/typeDefinition, textDocument/references, and
                    // textDocument/implementation
                    if (const auto *locationsPtr = get_if<vector<unique_ptr<Location>>>(&res)) {
                        return locationsPtr->empty() ? ResponseMessageStatus::EmptyResult
                                                     : ResponseMessageStatus::Succeeded;
                    } else {
                        return ResponseMessageStatus::EmptyResult;
                    }
                } else if constexpr (is_same_v<T, variant<JSONNullObject, unique_ptr<Hover>>>) {
                    // textDocument/hover
                    if (const auto *hoverPtr = get_if<unique_ptr<Hover>>(&res)) {
                        auto &hover = *hoverPtr;
                        return hover->contents->value.empty() ? ResponseMessageStatus::EmptyResult
                                                              : ResponseMessageStatus::Succeeded;
                    } else {
                        return ResponseMessageStatus::EmptyResult;
                    }
                } else if constexpr (is_same_v<T, unique_ptr<CompletionList>>) {
                    // textDocument/completion
                    return res->items.empty() ? ResponseMessageStatus::EmptyResult : ResponseMessageStatus::Succeeded;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, unique_ptr<PrepareRenameResult>>>) {
                    // textDocument/prepareRename
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, unique_ptr<WorkspaceEdit>>>) {
                    // textDocument/rename
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, unique_ptr<SignatureHelp>>>) {
                    // textDocument/signatureHelp
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, vector<unique_ptr<CodeAction>>>>) {
                    // textDocument/codeAction
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, optional<unique_ptr<CodeAction>>>) {
                    //  codeAction/resolve
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, vector<unique_ptr<SymbolInformation>>>>) {
                    // workspace/symbol
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, unique_ptr<SorbetErrorParams>>) {
                    // sorbet/error
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, unique_ptr<TextDocumentItem>>) {
                    // sorbet/readFile
                    return ResponseMessageStatus::Unknown;
                } else if constexpr (is_same_v<T, variant<JSONNullObject, unique_ptr<SymbolInformation>>>) {
                    // sorbet/showSymbol
                    return holds_alternative<JSONNullObject>(res) ? ResponseMessageStatus::EmptyResult
                                                                  : ResponseMessageStatus::Succeeded;
                } else {
                    static_assert(always_false_v<T>, "non-exhaustive visitor!");
                }
            },
            response.result.value());
    } else {
        ENFORCE(false, "ResponseMessage should have had eithe error or result set");
        return ResponseMessageStatus::Unknown;
    }
}

} // namespace

LSPTask::LSPTask(const LSPConfiguration &config, LSPMethod method) : config(config), method(method) {}

LSPTask::~LSPTask() {
    if (!latencyTimer) {
        return;
    }
    latencyTimer->setTag("method", methodString());
}

// Used for timers/counters, so we use `.` instead of `/` as a separator.
ConstExprStr LSPTask::methodString() const {
    switch (method) {
        case LSPMethod::$CancelRequest:
            return "cancelRequest";
        case LSPMethod::Exit:
            return "exit";
        case LSPMethod::PAUSE:
            return "PAUSE";
        case LSPMethod::RESUME:
            return "RESUME";
        case LSPMethod::GETCOUNTERS:
            return "GETCOUNTERS";
        case LSPMethod::Shutdown:
            return "shutdown";
        case LSPMethod::SorbetError:
            return "sorbet.error";
        case LSPMethod::SorbetFence:
            return "sorbet.fence";
        case LSPMethod::SorbetShowOperation:
        case LSPMethod::SorbetTypecheckRunInfo:
        case LSPMethod::TextDocumentPublishDiagnostics:
        case LSPMethod::WindowShowMessage:
            // Sorbet sends these to the client; the server shouldn't receive them.
            return "other";
        case LSPMethod::Initialize:
            return "initialize";
        case LSPMethod::Initialized:
            return "initialized";
        case LSPMethod::SorbetReadFile:
            return "sorbet.readFile";
        case LSPMethod::SorbetIndexerInitialization:
            return "sorbet.indexerInitialization";
        case LSPMethod::SorbetShowSymbol:
            return "sorbet.showSymbol";
        case LSPMethod::SorbetHierarchyReferences:
            return "sorbet.hierarchyReferences";
        case LSPMethod::SorbetWatchmanFileChange:
            return "sorbet.watchmanFileChange";
        case LSPMethod::SorbetWatchmanStateEnter:
            return "sorbet.watchmanStateEnter";
        case LSPMethod::SorbetWatchmanStateLeave:
            return "sorbet.watchmanStateLeave";
        case LSPMethod::SorbetWorkspaceEdit:
        case LSPMethod::TextDocumentDidChange:
        case LSPMethod::TextDocumentDidClose:
        case LSPMethod::TextDocumentDidOpen:
            return "sorbet.workspaceEdit";
        case LSPMethod::TextDocumentCodeAction:
            return "textDocument.codeAction";
        case LSPMethod::TextDocumentCompletion:
            return LSP_COMPLETION_METRICS_PREFIX;
        case LSPMethod::TextDocumentDefinition:
            return "textDocument.definition";
        case LSPMethod::TextDocumentDocumentHighlight:
            return "textDocument.documentHighlight";
        case LSPMethod::TextDocumentDocumentSymbol:
            return "textDocument.documentSymbol";
        case LSPMethod::TextDocumentHover:
            return "textDocument.hover";
        case LSPMethod::TextDocumentPrepareRename:
            return "textDocument.prepareRename";
        case LSPMethod::TextDocumentReferences:
            return "textDocument.references";
        case LSPMethod::TextDocumentRename:
            return "textDocument.rename";
        case LSPMethod::TextDocumentSignatureHelp:
            return "textDocument.signatureHelp";
        case LSPMethod::TextDocumentTypeDefinition:
            return "textDocument.typeDefinition";
        case LSPMethod::WorkspaceSymbol:
            return "workspace.symbol";
        case LSPMethod::WorkspaceDidChangeConfiguration:
            return "worksapce.didChangeConfiguration";
        case LSPMethod::TextDocumentImplementation:
            return "textDocument.implementation";
        case LSPMethod::CodeActionResolve:
            return "codeAction.resolve";
    }
}

LSPTask::Phase LSPTask::finalPhase() const {
    return Phase::RUN;
}

void LSPTask::preprocess(LSPPreprocessor &preprocessor) {}

void LSPTask::index(LSPIndexer &indexer) {}

core::packages::Stratum LSPTask::preemptionStratum(FileStratumMapping info) const {
    return info.getLastStratum();
}

LSPRequestTask::LSPRequestTask(const LSPConfiguration &config, MessageId id, LSPMethod method)
    : LSPTask(config, method), id(move(id)) {}

void LSPRequestTask::run(LSPTypecheckerDelegate &typechecker) {
    auto response = runRequest(typechecker);
    ENFORCE(response != nullptr);

    switch (statusForResponse(*response)) {
        case ResponseMessageStatus::Unknown:
            break;
        case ResponseMessageStatus::Errored:
            prodCategoryCounterInc("lsp.messages.run.errored", methodString());
            break;
        case ResponseMessageStatus::EmptyResult:
            prodCategoryCounterInc("lsp.messages.run.emptyresult", methodString());
            break;
        case ResponseMessageStatus::Succeeded:
            prodCategoryCounterInc("lsp.messages.run.succeeded", methodString());
            break;
    }

    config.output->write(move(response));
}

LSPTask::Phase LSPRequestTask::finalPhase() const {
    return Phase::RUN;
}

bool LSPRequestTask::cancel(const MessageId &id) {
    if (this->id.equals(id)) {
        if (latencyTimer) {
            latencyTimer->cancel();
            latencyTimer = nullptr;
        }
        auto response = make_unique<ResponseMessage>("2.0", id, method);
        prodCategoryCounterInc("lsp.messages.canceled", methodString());
        response->error = make_unique<ResponseError>((int)LSPErrorCodes::RequestCancelled, "Request was canceled");
        config.output->write(move(response));
        return true;
    }
    return false;
}

bool LSPTask::isDelayable() const {
    return false;
}

bool LSPTask::cancel(const MessageId &id) {
    return false;
}

bool LSPTask::canPreempt(const LSPIndexer &indexer) const {
    return true;
}

vector<unique_ptr<Location>>
LSPTask::extractLocations(const core::GlobalState &gs,
                          const vector<unique_ptr<core::lsp::QueryResponse>> &queryResponses,
                          vector<unique_ptr<Location>> locations) const {
    auto queryResponsesFiltered = LSPQuery::filterAndDedup(gs, queryResponses);
    for (auto &q : queryResponsesFiltered) {
        if (auto *send = q->isSend()) {
            addLocIfExists(gs, locations, send->funLoc());
        } else {
            addLocIfExists(gs, locations, q->getLoc());
        }
    }
    return locations;
}

vector<unique_ptr<core::lsp::QueryResponse>>
LSPTask::getReferencesToSymbols(LSPTypecheckerDelegate &typechecker,
                                core::lsp::Query::Symbol::STORAGE &&symbols) const {
    auto it = remove_if(symbols.begin(), symbols.end(), [](auto symbol) { return !symbol.exists(); });
    symbols.erase(it, symbols.end());

    if (!symbols.empty()) {
        auto run2 = LSPQuery::bySymbol(config, typechecker, move(symbols));
        return move(run2.responses);
    }
    return {};
}

vector<unique_ptr<core::lsp::QueryResponse>>
LSPTask::getReferencesToSymbolsInPackage(LSPTypecheckerDelegate &typechecker, core::packages::MangledName packageName,
                                         core::lsp::Query::Symbol::STORAGE &&symbols) const {
    auto it = remove_if(symbols.begin(), symbols.end(), [](auto symbol) { return !symbol.exists(); });
    symbols.erase(it, symbols.end());

    if (!symbols.empty()) {
        auto run2 = LSPQuery::bySymbol(config, typechecker, move(symbols), packageName);
        return move(run2.responses);
    }
    return {};
}

vector<unique_ptr<core::lsp::QueryResponse>>
LSPTask::getReferencesToSymbolsInFile(LSPTypecheckerDelegate &typechecker, core::FileRef fref,
                                      core::lsp::Query::Symbol::STORAGE &&symbols) const {
    auto it = remove_if(symbols.begin(), symbols.end(), [](auto symbol) { return !symbol.exists(); });
    symbols.erase(it, symbols.end());

    if (!symbols.empty() && fref.exists()) {
        auto run2 = LSPQuery::bySymbolsInFiles(config, typechecker, move(symbols), {fref});
        // Ignore results in other files (which may have been picked up for typechecking purposes)
        erase_if(run2.responses, [fref](auto &resp) { return resp->getLoc().file() != fref; });
        return move(run2.responses);
    }
    return {};
}

vector<unique_ptr<DocumentHighlight>>
LSPTask::getHighlights(LSPTypecheckerDelegate &typechecker,
                       const vector<unique_ptr<core::lsp::QueryResponse>> &queryResponses) const {
    vector<unique_ptr<DocumentHighlight>> highlights;
    auto locations = extractLocations(typechecker.state(), queryResponses);
    for (auto const &location : locations) {
        auto highlight = make_unique<DocumentHighlight>(move(location->range));
        highlights.emplace_back(move(highlight));
    }
    return highlights;
}

void LSPTask::addLocIfExists(const core::GlobalState &gs, vector<unique_ptr<Location>> &locs, core::Loc loc) const {
    auto location = config.loc2Location(gs, loc);
    if (location != nullptr) {
        locs.push_back(std::move(location));
    }
}

LSPDangerousTypecheckerTask::LSPDangerousTypecheckerTask(const LSPConfiguration &config, LSPMethod method)
    : LSPTask(config, method) {}

void LSPDangerousTypecheckerTask::run(LSPTypecheckerDelegate &tc) {
    Exception::raise("Bug: Dangerous typechecker tasks are expected to run specially");
}

} // namespace sorbet::realmain::lsp
