import org.eclipse.lsp4j.*;
import org.eclipse.lsp4j.jsonrpc.messages.Either;
import org.eclipse.lsp4j.services.LanguageClient;
import org.eclipse.lsp4j.services.LanguageClientAware;
import org.eclipse.lsp4j.services.LanguageServer;
import org.eclipse.lsp4j.services.TextDocumentService;
import org.eclipse.lsp4j.services.WorkspaceService;

import java.util.List;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;

import static grammar.AntlrLspAdapter.*;


/**
 * Implementation of the Language Server Protocol (LSP) for the Progress language.
 *
 * <p>This class implements {@link LanguageServer} and {@link LanguageClientAware},
 * providing core LSP features such as text document synchronization, completion,
 * semantic tokens, folding ranges, definitions, and hover support.</p>
 *
 * <p>The server uses custom {@link TextDocumentService} and {@link WorkspaceService}
 * implementations ({@link ProgressTextDocumentService} and {@link ProgressWorkSpaceService})
 * to handle document and workspace-related requests from the LSP client.</p>
 */
public class ProgressLanguageServer implements LanguageServer, LanguageClientAware {

   private TextDocumentService textDocumentService;
   private WorkspaceService workspaceService;
   private ClientCapabilities clientCapabilities;
   private LanguageClient languageClient;
   private int shutdown = 1;

   /**
    * Constructs a new ProgressLanguageServer instance.
    *
    * <p>Initializes the text document service and workspace service.</p>
    */
   public ProgressLanguageServer() {
      this.textDocumentService = new ProgressTextDocumentService(this);
      this.workspaceService = new ProgressWorkSpaceService(this);
   }

   /**
    * Configures server capabilities related to semantic tokens and folding ranges.
    *
    * @param capabilities the {@link ServerCapabilities} object to configure
    */
   private void setupSemanticTokensCapabilities(ServerCapabilities capabilities) {
      SemanticTokensWithRegistrationOptions semanticTokensCapabilities = new SemanticTokensWithRegistrationOptions();

      // Set up the legend for token types and modifiers
      SemanticTokensLegend legend = new SemanticTokensLegend(TOKEN_TYPES, TOKEN_MODIFIERS);
      semanticTokensCapabilities.setLegend(legend);

      // Enable full document and range semantic tokens
      semanticTokensCapabilities.setFull(Either.forLeft(true));
      semanticTokensCapabilities.setRange(Either.forLeft(true));

      // Enable other capabilities
      capabilities.setDefinitionProvider(true);
      capabilities.setHoverProvider(true);
      capabilities.setFoldingRangeProvider(true);
      capabilities.setSemanticTokensProvider(semanticTokensCapabilities);
   }

   /**
    * Initializes the language server.
    *
    * @param initializeParams initialization parameters from the client
    * @return a {@link CompletableFuture} that resolves to {@link InitializeResult} with server capabilities
    */
   @Override
   public CompletableFuture<InitializeResult> initialize(InitializeParams initializeParams) {
      final InitializeResult response = new InitializeResult(new ServerCapabilities());

      // Set document synchronization mode
      response.getCapabilities().setTextDocumentSync(TextDocumentSyncKind.Full);

      // Store client capabilities for later use
      this.clientCapabilities = initializeParams.getCapabilities();

      // If dynamic completion registration is not supported, register completion provider statically
      if (!isDynamicCompletionRegistration()) {
         response.getCapabilities().setCompletionProvider(new CompletionOptions());
      }

      // Setup semantic tokens and folding range capabilities
      setupSemanticTokensCapabilities(response.getCapabilities());

      return CompletableFuture.supplyAsync(() -> response);
   }

   /**
    * Called after initialization is complete.
    * Handles dynamic registration of completion if supported by the client.
    *
    * @param params parameters from the client
    */
   @Override
   public void initialized(InitializedParams params) {
      if (isDynamicCompletionRegistration()) {
         CompletionRegistrationOptions completionRegistrationOptions = new CompletionRegistrationOptions();
         Registration completionRegistration = new Registration(UUID.randomUUID().toString(),
                 "textDocument/completion", completionRegistrationOptions);

         languageClient.registerCapability(new RegistrationParams(List.of(completionRegistration)));
      }
   }

   /**
    * Shuts down the language server.
    *
    * @return a {@link CompletableFuture} that completes when shutdown is processed
    */
   @Override
   public CompletableFuture<Object> shutdown() {
      shutdown = 0;
      return CompletableFuture.supplyAsync(Object::new);
   }

   /**
    * Exits the language server process.
    */
   @Override
   public void exit() {
      System.exit(shutdown);
   }

   /**
    * Returns the {@link TextDocumentService} used by this server.
    *
    * @return text document service
    */
   @Override
   public TextDocumentService getTextDocumentService() {
      return this.textDocumentService;
   }

   /**
    * Returns the {@link WorkspaceService} used by this server.
    *
    * @return workspace service
    */
   @Override
   public WorkspaceService getWorkspaceService() {
      return this.workspaceService;
   }

   /**
    * Connects a {@link LanguageClient} to this server.
    *
    * @param languageClient the client to connect
    */
   @Override
   public void connect(LanguageClient languageClient) {
      this.languageClient = languageClient;
      LSClientLogger.getInstance().initialize(this.languageClient);
   }

   /**
    * Checks if the client supports dynamic registration for completion.
    *
    * @return true if dynamic completion registration is supported, false otherwise
    */
   private boolean isDynamicCompletionRegistration() {
      TextDocumentClientCapabilities textDocumentCapabilities =
              clientCapabilities.getTextDocument();
      return textDocumentCapabilities != null
              && textDocumentCapabilities.getCompletion() != null
              && Boolean.FALSE.equals(textDocumentCapabilities.getCompletion().getDynamicRegistration());
   }
}