// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

import { ILanguageService, log } from "qsharp-lang";
import * as vscode from "vscode";
import { isQsharpNotebookCell } from "../common.js";
import { findQSharpCellMagic, jupyterNotebookType } from "../notebook.js";

const qsharpConfigMimeType = "application/x.qsharp-config";

const openQSharpNotebooks = new Set<string>();

/**
 * Document update handlers for syncing notebook cell contents with the language service.
 */
export function registerQSharpNotebookCellUpdateHandlers(
  languageService: ILanguageService,
) {
  vscode.workspace.notebookDocuments.forEach((notebook) => {
    updateIfQsharpNotebook(notebook);
  });

  const subscriptions = [];
  subscriptions.push(
    vscode.workspace.onDidOpenNotebookDocument((notebook) => {
      updateIfQsharpNotebook(notebook);
    }),
  );

  subscriptions.push(
    vscode.workspace.onDidChangeNotebookDocument((event) => {
      updateIfQsharpNotebook(event.notebook);
    }),
  );

  subscriptions.push(
    vscode.workspace.onDidCloseNotebookDocument((notebook) => {
      closeIfKnownQsharpNotebook(notebook);
    }),
  );

  function updateIfQsharpNotebook(notebook: vscode.NotebookDocument) {
    if (notebook.notebookType === jupyterNotebookType) {
      const qsharpMetadata = getQSharpConfigMetadata(notebook);
      const qsharpCells = getQSharpCells(notebook);
      const notebookUri = notebook.uri.toString();
      if (qsharpCells.length > 0) {
        openQSharpNotebooks.add(notebookUri);
        languageService.updateNotebookDocument(
          notebookUri,
          notebook.version,
          qsharpMetadata,
          qsharpCells.map((cell) => {
            return {
              uri: cell.document.uri.toString(),
              version: cell.document.version,
              code: getQSharpText(cell.document),
            };
          }),
        );
      } else {
        // All Q# cells could have been deleted, check if we know this doc from previous calls
        closeIfKnownQsharpNotebook(notebook);
      }
    }
  }

  function closeIfKnownQsharpNotebook(notebook: vscode.NotebookDocument) {
    const notebookUri = notebook.uri.toString();
    if (openQSharpNotebooks.has(notebookUri)) {
      languageService.closeNotebookDocument(notebookUri);
      openQSharpNotebooks.delete(notebook.uri.toString());
    }
  }

  function getQSharpCells(notebook: vscode.NotebookDocument) {
    return notebook
      .getCells()
      .filter((cell) => isQsharpNotebookCell(cell.document));
  }

  function getQSharpText(document: vscode.TextDocument) {
    const magicRange = findQSharpCellMagic(document);
    if (magicRange) {
      const magicStartOffset = document.offsetAt(magicRange.start);
      const magicEndOffset = document.offsetAt(magicRange.end);
      // Erase the %%qsharp magic line if it's there.
      // Replace it with a comment so that document offsets remain the same.
      // This will save us from having to map offsets later when
      // communicating with the language service.
      const text = document.getText();
      return (
        text.substring(0, magicStartOffset) +
        "//qsharp" +
        text.substring(magicEndOffset)
      );
    } else {
      // No %%qsharp magic. This can happen if the user manually sets the
      // cell language to Q#. Python won't recognize the cell as a Q# cell,
      // so this will fail at runtime, but as the language service we respect
      // the manually set cell language, so we treat this as any other
      // Q# cell. We could consider raising a warning here to help the user.
      log.info(
        "found Q# cell without %%qsharp magic: " + document.uri.toString(),
      );
      return document.getText();
    }
  }

  return subscriptions;
}

/**
 * Finds an output cell that contains an item with the Q# config MIME type,
 * and returns the data from it. This data and is generated by the execution of a
 * `qsharp.init()` call. It's Q# configuration data to be passed
 * to the language service as "notebook metadata".
 */
function getQSharpConfigMetadata(notebook: vscode.NotebookDocument): object {
  const data = notebook
    .getCells()
    .flatMap((cell) => cell.outputs)
    .flatMap((output) => output.items)
    .find((item) => {
      return item.mime === qsharpConfigMimeType;
    })?.data;

  if (data) {
    const dataString = new TextDecoder().decode(data);
    log.trace("found Q# config metadata: " + dataString);
    return JSON.parse(dataString);
  } else {
    return {};
  }
}
