import { JupyterFrontEnd } from '@jupyterlab/application';
import { Menu } from '@lumino/widgets';
import { INotebookTracker } from '@jupyterlab/notebook';
import { postToWebView } from './webview';

export type GenericCommandVerb = 'Explain' | 'Refactor' | 'Fix' | 'Optimize';

/**
 * This function adds a context menu item for Amazon Q and the related commands to
 * pass content to Q chat. These context menus are available in notebook views.
 */
export const registerContextMenuActions = ({
  app,
  notebookTracker,
  showFlareWidget
}: {
  app: JupyterFrontEnd;
  notebookTracker: INotebookTracker;
  showFlareWidget: () => void;
}) => {
  const verbs: GenericCommandVerb[] = [
    'Explain',
    'Refactor',
    'Fix',
    'Optimize'
  ];

  // Helper function to get error message text
  const getErrorMessage = (): string => {
    const notebook = notebookTracker.currentWidget;
    if (notebook) {
      const activeCell = notebook.content.activeCell;
      if (activeCell && activeCell.model.type === 'code') {
        const outputs = (activeCell.model as any).outputs;
        if (outputs) {
          for (let i = 0; i < outputs.length; i++) {
            const output = outputs.get(i);
            if (output.type === 'error') {
              return `ErrorType: ${output._raw.ename}; ErrorValue: ${output._raw.evalue}; Trace: ${output._raw.traceback?.toString()}`;
            }
          }
        }
      }
    }
    return '';
  };

  // Helper function to get selection or cell content
  const getSelectionOrCellContent = (verb: GenericCommandVerb): string => {
    // For Fix verb, return error message + cell content
    if (verb === 'Fix' && hasError()) {
      const notebook = notebookTracker.currentWidget;
      const cellContent =
        notebook?.content.activeCell?.model.type === 'code'
          ? notebook.content.activeCell.model.sharedModel.getSource()
          : '';
      const errorMessage = getErrorMessage();
      return `Code:\n${cellContent}\n\nError:\n${errorMessage}`;
    }

    // For other verbs, use normal selection logic
    const notebook = notebookTracker.currentWidget;
    if (notebook) {
      const activeCell = notebook.content.activeCell;
      if (activeCell && activeCell.model.type === 'code') {
        const editor = activeCell.editor;
        if (editor) {
          const selection = editor.getSelection();
          if (
            selection.start.line !== selection.end.line ||
            selection.start.column !== selection.end.column
          ) {
            const selectedText = editor.model.sharedModel
              .getSource()
              .substring(
                editor.getOffsetAt(selection.start),
                editor.getOffsetAt(selection.end)
              );
            if (selectedText.trim()) {
              return selectedText;
            }
          }
        }
        return activeCell.model.sharedModel.getSource();
      }
    }

    return '';
  };

  // Helper function to check for errors in cell output
  const hasError = (): boolean => {
    const notebook = notebookTracker.currentWidget;
    if (notebook) {
      const activeCell = notebook.content.activeCell;
      if (activeCell && activeCell.model.type === 'code') {
        const outputs = (activeCell.model as any).outputs;
        if (outputs) {
          for (let i = 0; i < outputs.length; i++) {
            const output = outputs.get(i);
            if (output.type === 'error') {
              return true;
            }
          }
        }
      }
    }
    return false;
  };

  // Register individual commands
  verbs.forEach(verb => {
    const commandId = `sagemaker-gen-ai:${verb.toLowerCase()}`;

    app.commands.addCommand(commandId, {
      label: `${verb}`,
      isEnabled: () => (verb === 'Fix' ? hasError() : true),
      execute: () => {
        const selection = getSelectionOrCellContent(verb);
        if (selection) {
          console.log(`${verb}:`, selection);
          showFlareWidget();
          postToWebView({
            command: 'genericCommand',
            params: {
              genericCommand: verb,
              selection,
              triggerType: 'contextMenu'
            }
          });
        }
      }
    });
  });

  // Create submenu
  const submenu = new Menu({ commands: app.commands });
  submenu.title.label = 'Amazon Q';

  // Add commands to the submenu
  verbs.forEach(verb => {
    submenu.addItem({
      command: `sagemaker-gen-ai:${verb.toLowerCase()}`
    });
  });

  const selectorsWhereQMenuIsRendered = ['.jp-CodeCell'];

  selectorsWhereQMenuIsRendered.forEach(selector =>
    app.contextMenu.addItem({
      type: 'submenu',
      submenu: submenu,
      selector,
      rank: 0
    })
  );
};
