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

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

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

const openQdkNotebooks = new Set<string>();

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

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

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

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

  function updateIfQdkNotebook(notebook: vscode.NotebookDocument) {
    if (notebook.notebookType === jupyterNotebookType) {
      const qsharpMetadata = getQSharpConfigMetadata(notebook);
      const qdkCells = getQdkCells(notebook);
      const notebookUri = notebook.uri.toString();
      if (qdkCells.length > 0) {
        openQdkNotebooks.add(notebookUri);
        languageService.updateNotebookDocument(
          notebookUri,
          notebook.version,
          qsharpMetadata,
          qdkCells.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 (openQdkNotebooks.has(notebookUri)) {
      languageService.closeNotebookDocument(notebookUri);
      openQdkNotebooks.delete(notebook.uri.toString());
    }
  }

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

  function getQSharpText(document: vscode.TextDocument) {
    const qsharpMagicRange = findQSharpCellMagic(document);
    if (qsharpMagicRange) {
      const magicStartOffset = document.offsetAt(qsharpMagicRange.start);
      const magicEndOffset = document.offsetAt(qsharpMagicRange.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)
      );
    }

    const magicRange = findOpenQasmCellMagic(document);
    if (magicRange) {
      return extractQSharpText(document, magicRange);
    } else {
      // No %%qsharp/%%openqasm magic. This can happen if the user manually sets the
      // cell language to Q#/OpenQASM. Python won't recognize the cell as a QDK cell,
      // so this will fail at runtime, but as the language service we respect
      // the manually set cell language, so we treat this as a Q# cell.
      // We could consider raising a warning here to help the user.
      log.info(
        "found QDK cell without correct magic: " + document.uri.toString(),
      );
      return document.getText();
    }
  }

  // placeoholder for openqasm cell preprocessing
  function extractQSharpText(
    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    document: vscode.TextDocument,
    // eslint-disable-next-line @typescript-eslint/no-unused-vars
    magicRange: vscode.Range,
  ): string {
    return "";
  }

  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);
    const dataStringJSON = JSON.parse(dataString);
    // Re-normalize projectRoot URI if present.
    // Python and VS Code have different URI normalization logic,
    // so we need to parse and re-serialize it here.
    if (typeof dataStringJSON.projectRoot === "string") {
      const projectRootUri = vscode.Uri.parse(dataStringJSON.projectRoot, true);
      dataStringJSON.projectRoot = projectRootUri.toString();
    }

    return dataStringJSON;
  } else {
    // Default to Unrestricted profile for notebooks when no explicit configuration is provided
    // This aligns with the Python qsharp runtime behavior
    return { targetProfile: "unrestricted" };
  }
}
