#pragma once
#include "LSPProxyServer.h"
#include "LanguageServerRouter.h"
#include "CrossLanguageDefinitionResolver.h"
#include "CrossLanguageSymbolTable.h"
#include "LSPWorkspaceIndex.h"
#include "LSPDiagnosticsAggregator.h"
#include "LSPCompletionMerger.h"
#include <nlohmann/json.hpp>
#include <string>
#include <vector>

namespace whetstone {

struct OrchestrationSummary {
    int openDocuments  = 0;
    int diagnosticUris = 0;
    int completionItems= 0;
    int errorCount     = 0;   // total error-severity diagnostics across all open docs
};

class LSPOrchestrator {
public:
    explicit LSPOrchestrator(const CrossLanguageSymbolTable& table)
        : resolver_(table) {
        // Wire cross-language definition resolution into the proxy
        proxy_.registerHandler("textDocument/definition",
            [this](const nlohmann::json& params) -> nlohmann::json {
                return handleDefinition(params);
            });
        // Wire didOpen/didChange/didClose into workspace index
        proxy_.registerHandler("textDocument/didOpen",
            [this](const nlohmann::json& params) -> nlohmann::json {
                handleDidOpen(params); return {};
            });
        proxy_.registerHandler("textDocument/didChange",
            [this](const nlohmann::json& params) -> nlohmann::json {
                handleDidChange(params); return {};
            });
        proxy_.registerHandler("textDocument/didClose",
            [this](const nlohmann::json& params) -> nlohmann::json {
                handleDidClose(params); return {};
            });
    }

    // Route any LSP message through the proxy.
    nlohmann::json handle(const nlohmann::json& message) {
        return proxy_.handle(message);
    }

    // Push diagnostics from a language server source.
    void publishDiagnostics(const std::string& source, const std::string& uri,
                            const std::vector<Diagnostic>& diags) {
        diags_.set(source, uri, diags);
    }

    // Register completion items from a language server source for the current position.
    void registerCompletions(const std::string& source,
                             const std::vector<CompletionItem>& items) {
        completions_.addItems(source, items);
    }

    // Get merged completions for a URI (source-agnostic).
    std::vector<CompletionItem> getCompletions() const {
        return completions_.getItems();
    }

    // Get merged diagnostics for a URI.
    std::vector<Diagnostic> getDiagnostics(const std::string& uri) const {
        return diags_.get(uri);
    }

    // Summary of current orchestrator state.
    OrchestrationSummary summary() const {
        OrchestrationSummary s;
        s.openDocuments   = index_.size();
        s.diagnosticUris  = static_cast<int>(diags_.urisWithDiagnostics().size());
        s.completionItems = completions_.uniqueCount();
        // Count total errors across all open docs
        for (const auto& uri : allOpenUris())
            s.errorCount += diags_.errorCount(uri);
        return s;
    }

    const LSPWorkspaceIndex& workspaceIndex() const { return index_; }

private:
    LSPProxyServer               proxy_;
    LanguageServerRouter         router_;
    CrossLanguageDefinitionResolver resolver_;
    LSPWorkspaceIndex            index_;
    LSPDiagnosticsAggregator     diags_;
    LSPCompletionMerger          completions_;

    nlohmann::json handleDefinition(const nlohmann::json& params) {
        std::string uri    = params.value("textDocument", nlohmann::json::object())
                                   .value("uri", "");
        std::string symbol = params.value("symbol", "");
        if (symbol.empty()) return nullptr;
        auto loc = resolver_.resolveFromUri(uri, symbol);
        if (!loc.found) return nullptr;
        return {{"uri",        "scip://" + loc.providerComponent},
                {"provider",   loc.providerComponent},
                {"language",   loc.providerLanguage},
                {"scipSymbol", loc.scipSymbol}};
    }

    void handleDidOpen(const nlohmann::json& params) {
        auto td = params.value("textDocument", nlohmann::json::object());
        std::string uri  = td.value("uri",        "");
        std::string lang = td.value("languageId", "");
        std::string text = td.value("text",       "");
        int ver          = td.value("version",    0);
        if (!uri.empty()) index_.open(uri, lang, text, ver);
    }

    void handleDidChange(const nlohmann::json& params) {
        auto td  = params.value("textDocument", nlohmann::json::object());
        std::string uri = td.value("uri", "");
        int ver         = td.value("version", 0);
        std::string text;
        if (params.contains("contentChanges") && params["contentChanges"].is_array()
            && !params["contentChanges"].empty())
            text = params["contentChanges"][0].value("text", "");
        if (!uri.empty()) index_.update(uri, text, ver);
    }

    void handleDidClose(const nlohmann::json& params) {
        auto td = params.value("textDocument", nlohmann::json::object());
        std::string uri = td.value("uri", "");
        if (!uri.empty()) index_.close(uri);
    }

    std::vector<std::string> allOpenUris() const {
        // Return all URIs from the workspace index
        std::vector<std::string> uris;
        for (auto& lang : index_.allLanguages())
            for (auto& u : index_.listUris(lang))
                uris.push_back(u);
        return uris;
    }
};

} // namespace whetstone