import grammar.AntlrLspAdapter;
import grammar.ProgressLexer;
import grammar.ProgressParser;
import grammar.SemanticToken;
import grammar.FoldingRangeCollectorParserVisitor;
import org.antlr.v4.runtime.CharStream;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.eclipse.lsp4j.*;
import org.eclipse.lsp4j.jsonrpc.messages.Either;
import org.eclipse.lsp4j.services.TextDocumentService;

import java.util.*;
import java.util.concurrent.CompletableFuture;

/**
 * TextDocumentService implementation for the Progress language server.
 *
 * <p>This class implements {@link TextDocumentService} and provides support for:
 * <ul>
 *     <li>Document open, change, close, and save notifications</li>
 *     <li>Code completion</li>
 *     <li>Semantic tokens (full document and range)</li>
 *     <li>Folding ranges</li>
 *     <li>Hover</li>
 *     <li>Goto definition</li>
 * </ul>
 * </p>
 *
 * <p>It uses ANTLR-generated lexer and parser to analyze the document content
 * and produce semantic tokens and folding ranges, and it stores document
 * contents in-memory for quick access.</p>
 */
public class ProgressTextDocumentService implements TextDocumentService {

   private final ProgressLanguageServer languageServer;
   private final LSClientLogger clientLogger;
   private final AntlrLspAdapter antlrLspAdapter = new AntlrLspAdapter();

   /** We store the doc contents with their URI as keys, */
   private final Map<String, String> documentContents = new HashMap<>();

   /**
    * Constructs a new ProgressTextDocumentService.
    *
    * @param languageServer the parent {@link ProgressLanguageServer} instance
    */
   public ProgressTextDocumentService(ProgressLanguageServer languageServer) {
      this.languageServer = languageServer;
      this.clientLogger = LSClientLogger.getInstance();
   }

   /**
    * Called when a text document is opened in the client.
    *
    * @param params parameters including document URI and content
    */
   @Override
   public void didOpen(DidOpenTextDocumentParams params) {
      String uri = params.getTextDocument().getUri();
      String content = params.getTextDocument().getText();
      documentContents.put(uri, content);

      clientLogger.logMessage("Document opened: " + uri, MessageType.Info);
   }

   /**
    * Called when a text document is changed.
    * <p>This implementation assumes full synchronization and replaces
    * the entire document content with the new text.</p>
    *
    * @param params parameters including changes
    */
   @Override
   public void didChange(DidChangeTextDocumentParams params) {
      String uri = params.getTextDocument().getUri();
      List<TextDocumentContentChangeEvent> changes = params.getContentChanges();

      if (!changes.isEmpty()) {
         String newContent = changes.get(0).getText();
         documentContents.put(uri, newContent);
      }
   }

   /**
    * Called when a text document is closed.
    *
    * @param params parameters including document URI
    */
   @Override
   public void didClose(DidCloseTextDocumentParams params) {
      String uri = params.getTextDocument().getUri();
      documentContents.remove(uri);
      clientLogger.logMessage("Document closed: " + uri, MessageType.Error);
   }

   /**
    * Called when a text document is saved.
    *
    * @param params parameters including document URI
    */
   @Override
   public void didSave(DidSaveTextDocumentParams params) {
      clientLogger.logMessage("Document saved: " + params.getTextDocument().getUri(), MessageType.Error);
   }

   /**
    * Provides code completions at the cursor position.
    *
    * @param position completion parameters including document URI and cursor position
    * @return a future with either a list of {@link CompletionItem} or a {@link CompletionList}
    */
   @Override
   public CompletableFuture<Either<List<CompletionItem>, CompletionList>> completion(CompletionParams position) {
      var tokens = documentContents.get(position.getTextDocument().getUri()).split("\\s+");
      String lastToken = tokens[tokens.length - 1];

      return CompletableFuture.supplyAsync(() -> {
         clientLogger.logMessage("Completion token: " + lastToken, MessageType.Info);
         return Either.forLeft(antlrLspAdapter.getCompletions(lastToken));
      });
   }

   /**
    * Provides semantic tokens for the full document.
    *
    * @param params semantic token parameters including document URI
    * @return a future with {@link SemanticTokens} for the document
    */
   @Override
   public CompletableFuture<SemanticTokens> semanticTokensFull(SemanticTokensParams params) {
      String uri = params.getTextDocument().getUri();
      String content = documentContents.get(uri);

      if (content == null) {
         return CompletableFuture.completedFuture(new SemanticTokens(List.of()));
      }

      return CompletableFuture.supplyAsync(() -> {
         try {
            List<SemanticToken> tokens = antlrLspAdapter.analyzeText(content);
            List<Integer> data = antlrLspAdapter.tokensToLSPFormat(tokens);
            return new SemanticTokens(data);
         } catch (Exception e) {
            clientLogger.logMessage("Error analyzing semantic tokens: " + e.getMessage(), MessageType.Error);
            return new SemanticTokens(List.of());
         }
      });
   }

   /**
    * Provides semantic tokens for a specified range within a document.
    *
    * @param params range parameters including document URI and range
    * @return a future with {@link SemanticTokens} for the range
    */
   @Override
   public CompletableFuture<SemanticTokens> semanticTokensRange(SemanticTokensRangeParams params) {
      String uri = params.getTextDocument().getUri();
      Range range = params.getRange();

      return CompletableFuture.supplyAsync(() -> {
         List<SemanticToken> tokens = antlrLspAdapter.analyzeTextRange(documentContents.get(uri), range);
         List<Integer> data = antlrLspAdapter.tokensToLSPFormat(tokens);
         return new SemanticTokens(data);
      });
   }

   /**
    * Provides folding ranges for a document by parsing its content using ANTLR.
    *
    * @param params folding range request parameters including document URI
    * @return a future with a list of {@link FoldingRange} objects
    */
   @Override
   public CompletableFuture<List<FoldingRange>> foldingRange(FoldingRangeRequestParams params) {
      return CompletableFuture.supplyAsync(() -> {
         String content = documentContents.get(params.getTextDocument().getUri());
         CharStream input = CharStreams.fromString(content);

         ProgressLexer lexer = new ProgressLexer(input);
         CommonTokenStream tokens = new CommonTokenStream(lexer);
         ProgressParser parser = new ProgressParser(tokens);

         ProgressParser.ProgramContext tree = parser.program();
         return new FoldingRangeCollectorParserVisitor().visit(tree);
      });
   }

   /**
    * Provides delta semantic tokens (not implemented).
    *
    * @param params delta parameters
    * @return a future with either {@link SemanticTokens} or {@link SemanticTokensDelta}
    */
   @Override
   public CompletableFuture<Either<SemanticTokens, SemanticTokensDelta>> semanticTokensFullDelta(SemanticTokensDeltaParams params) {
      return TextDocumentService.super.semanticTokensFullDelta(params);
   }

   /**
    * Provides hover information at a given position.
    *
    * @param params hover parameters
    * @return a future with {@link Hover} content
    */
   @Override
   public CompletableFuture<Hover> hover(HoverParams params) {
      Hover hover = new Hover();
      hover.setContents(Collections.singletonList(Either.forLeft("This is plain text hover")));
      return CompletableFuture.completedFuture(hover);
   }

   /**
    * Provides a "goto definition" location for a given position.
    *
    * <p>This mock implementation always points to the first character of the file.</p>
    *
    * @param params definition parameters
    * @return a future with either a list of {@link Location} or {@link LocationLink}
    */
   @Override
   public CompletableFuture<Either<List<? extends Location>, List<? extends LocationLink>>> definition(DefinitionParams params) {
      String uri = params.getTextDocument().getUri();
      Range range = new Range(new Position(0, 0), new Position(0, 1));
      Location location = new Location(uri, range);
      return CompletableFuture.completedFuture(Either.forLeft(List.of(location)));
   }
}
