import {
  JupyterFrontEnd,
  JupyterFrontEndPlugin
} from '@jupyterlab/application';
import { ICommandPalette } from '@jupyterlab/apputils';
import { INotebookTracker } from '@jupyterlab/notebook';
import { IIOPubMessage } from '@jupyterlab/services/lib/kernel/messages';
import * as monaco from 'monaco-editor';
import OpenAI from 'openai';
import { ISettingRegistry } from '@jupyterlab/settingregistry';
import { OpenAIClient, AzureKeyCredential } from '@azure/openai';
import { calculateHash, cosineSimilarity } from './utils';
import { generatePrompt, systemPrompt } from './prompt';

const PLUGIN_ID = 'pretzelai:plugin';

const NUMBER_OF_SIMILAR_CELLS = 3;
const EMBEDDING_MODEL = 'text-embedding-3-large';

const extension: JupyterFrontEndPlugin<void> = {
  id: PLUGIN_ID,
  autoStart: true,
  requires: [ICommandPalette, INotebookTracker, ISettingRegistry],
  activate: async (
    app: JupyterFrontEnd,
    palette: ICommandPalette,
    notebookTracker: INotebookTracker,
    settingRegistry: ISettingRegistry
  ) => {
    const { commands } = app;
    const command = 'pretzelai:replace-code';
    const placeholderDisabled =
      'To use AI features, please set your OpenAI API key or Azure API details in the Pretzel AI Settings.\n' +
      'You can also use the free Pretzel AI server.\n' +
      'Go To: Settings > Settings Editor > Pretzel AI Settings to configure';

    const placeHolderEnabled =
      'Ask AI. Use @variableName to refernce defined variables and dataframes';
    let openAiApiKey = '';
    let openAiBaseUrl = '';
    let aiService = 'OpenAI API key';
    let azureBaseUrl = '';
    let azureDeploymentName = '';
    let azureApiKey = '';
    let openai: OpenAI;

    // Function to load settings
    function loadSettings(updateFunc?: () => void) {
      settingRegistry
        .load(extension.id)
        .then(settings => {
          openAiApiKey = (settings.get('openAiSettings').composite as any)
            .openAiApiKey as string;
          aiService = settings.get('aiService').composite as string;
          openAiBaseUrl = (settings.get('openAiSettings').composite as any)
            .openAiBaseUrl as string;
          azureBaseUrl = (settings.get('azureSettings').composite as any)
            .azureBaseUrl as string;
          azureDeploymentName = (settings.get('azureSettings').composite as any)
            .azureDeploymentName as string;
          azureApiKey = (settings.get('azureSettings').composite as any)
            .azureApiKey as string;
          updateFunc?.();
        })
        .catch(reason => {
          console.error('Failed to load settings for Pretzel', reason);
        });
    }
    loadSettings();

    function loadOpenai() {
      if (openAiApiKey) {
        openai = new OpenAI({
          apiKey: openAiApiKey,
          dangerouslyAllowBrowser: true
        });
      } else {
        setTimeout(loadOpenai, 1000);
      }
    }
    loadOpenai();

    // Listen for future changes in settings
    settingRegistry.pluginChanged.connect((sender, plugin) => {
      if (plugin === extension.id) {
        const updateFunc = async () => {
          const submitButton = document.querySelector(
            '.pretzelInputSubmitButton'
          );
          const inputField = document.querySelector('.pretzelInputField');

          if (submitButton) {
            if (
              (aiService === 'OpenAI API key' && openAiApiKey) ||
              aiService === 'Use Pretzel AI Server' ||
              (aiService === 'Use Azure API' &&
                azureBaseUrl &&
                azureDeploymentName &&
                azureApiKey)
            ) {
              (submitButton as HTMLInputElement).disabled = false;
              (inputField as HTMLInputElement).placeholder = placeHolderEnabled;
            } else {
              (submitButton as HTMLInputElement).disabled = true;
              (inputField as HTMLInputElement).placeholder =
                placeholderDisabled;
            }
          }
        };
        loadSettings(updateFunc);
      }
    });

    type Embedding = {
      id: string;
      source: string;
      hash: string;
      embedding: number[];
    };

    let embeddings: Embedding[];

    async function createEmbeddings(
      embeddingsJSON: Embedding[],
      cells: any[],
      path: string
    ) {
      let newEmbeddings = false;
      embeddings = embeddingsJSON;
      const promises = cells.map(cell => {
        return (async () => {
          const index = embeddings.findIndex(e => e.id === cell.id);
          if (index !== -1) {
            const hash = await calculateHash(cell.source);
            if (hash !== embeddings[index].hash) {
              try {
                console.log(
                  'Hash changed: creating embedding for cell ',
                  cell.id
                );
                newEmbeddings = true;
                const response = await openai.embeddings.create({
                  model: EMBEDDING_MODEL,
                  input: cell.source
                });
                embeddings[index] = {
                  id: cell.id,
                  source: cell.source,
                  hash,
                  embedding: response.data[0].embedding
                };
              } catch (error) {
                console.error('Error generating embedding:', error);
              }
            }
          } else {
            console.log('New cell: creating embedding for cell ', cell.id);
            newEmbeddings = true;
            try {
              const response = await openai.embeddings.create({
                model: EMBEDDING_MODEL,
                input: cell.source
              });
              const hash = await calculateHash(cell.source);
              embeddings.push({
                id: cell.id,
                source: cell.source,
                hash,
                embedding: response.data[0].embedding
              });
            } catch (error) {
              console.error('Error generating embedding:', error);
            }
          }
        })();
      });
      await Promise.allSettled(promises);
      if (newEmbeddings) {
        app.serviceManager.contents.save(path, {
          type: 'file',
          format: 'text',
          content: JSON.stringify(embeddings)
        });
      }
    }

    // Function to print the source of all cells once the notebook is defined
    function getEmbeddings() {
      const notebook = notebookTracker.currentWidget;
      if (notebook?.model && openai) {
        const currentNotebookPath = notebook.context.path;
        const embeddingsPath =
          currentNotebookPath.replace('.ipynb', '') + '_embeddings.json';

        app.serviceManager.contents
          .get(embeddingsPath)
          .then(file => {
            try {
              const embJSON = JSON.parse(file.content);
              createEmbeddings(
                embJSON,
                notebook!.model!.sharedModel.cells,
                embeddingsPath
              );
            } catch (error) {
              console.error('Error parsing embeddings JSON:', error);
            }
          })
          .catch(async error => {
            createEmbeddings(
              [],
              notebook!.model!.sharedModel.cells,
              embeddingsPath
            );
          });
        // Temporary solution to keep refreshing hashes in non blocking thread
        setTimeout(getEmbeddings, 1000);
      } else {
        setTimeout(getEmbeddings, 1000);
      }
    }
    getEmbeddings();

    async function getVariableValue(
      variableName: string
    ): Promise<string | null> {
      const notebook = notebookTracker.currentWidget;
      if (notebook && notebook.sessionContext.session?.kernel) {
        const kernel = notebook.sessionContext.session.kernel;
        try {
          // get the type - if dataframe, we get columns
          // if other, we get the string representation
          const executeRequest = kernel.requestExecute({
            code: `print(${variableName})`
          });
          let variableValue: string | null = null;

          // Registering a message hook to intercept messages
          kernel.registerMessageHook(
            executeRequest.msg.header.msg_id,
            (msg: IIOPubMessage) => {
              if (
                msg.header.msg_type === 'stream' &&
                // @ts-expect-error tserror
                msg.content.name === 'stdout'
              ) {
                // @ts-expect-error tserror
                variableValue = msg.content.text.trim();
              }
              return true;
            }
          );

          // Await the completion of the execute request
          const reply = await executeRequest.done;
          if (reply && reply.content.status === 'ok') {
            return variableValue;
          } else {
            console.error('Failed to retrieve variable value');
            return null;
          }
        } catch (error) {
          console.error('Error retrieving variable value:', error);
          return null;
        }
      } else {
        console.error('No active kernel found');
        return null;
      }
    }

    commands.addCommand(command, {
      label: 'Replace Cell Code',
      execute: () => {
        const activeCell = notebookTracker.activeCell;

        if (activeCell) {
          // Cmd K twice should toggle the box
          // Check if an existing div with ID pretzelParentContainerAI exists on activeCell.node
          const existingDiv = activeCell.node.querySelector(
            '.pretzelParentContainerAI'
          );
          if (existingDiv) {
            // If so, delete that div
            existingDiv.remove();
            // Switch focus back to the Jupyter cell
            activeCell!.editor!.focus();
            return;
          }

          const oldCode = activeCell.model.sharedModel.source;

          // Create a parent container for all dynamically created elements
          const parentContainer = document.createElement('div');
          parentContainer.classList.add('pretzelParentContainerAI');
          activeCell.node.appendChild(parentContainer);
          // Create an input field and append it below the cell
          const inputContainer = document.createElement('div');
          inputContainer.style.marginTop = '10px';
          inputContainer.style.marginLeft = '70px';
          inputContainer.style.display = 'flex';
          inputContainer.style.flexDirection = 'column';
          parentContainer.appendChild(inputContainer);

          const inputField = document.createElement('textarea');
          inputField.classList.add('pretzelInputField');
          inputField.placeholder = placeHolderEnabled;
          inputField.style.width = '100%';
          inputField.style.height = '100px';
          inputContainer.appendChild(inputField);
          inputField.addEventListener('keydown', event => {
            if (event.key === 'Escape') {
              // TODO: this doesn't work - the Escape key isn't being captured
              // but every other key press is being captured
              event.preventDefault(); // Prevent any default behavior
              // Shift focus back to the editor of the active cell
              const activeCell = notebookTracker.activeCell;
              if (activeCell && activeCell.editor) {
                activeCell.editor.focus(); // Focus the editor of the active cell
              }
            }
            // handle enter key press to trigger submit
            if (event.key === 'Enter') {
              event.preventDefault();
              if (!submitButton.disabled) {
                handleAccept();
              }
            }
          });

          const inputFieldButtonsContainer = document.createElement('div');
          inputFieldButtonsContainer.style.marginTop = '10px';
          inputFieldButtonsContainer.style.display = 'flex';
          inputFieldButtonsContainer.style.flexDirection = 'row';
          inputContainer.appendChild(inputFieldButtonsContainer);
          inputField.focus();

          const submitButton = document.createElement('button');
          submitButton.classList.add('pretzelInputSubmitButton');
          submitButton.textContent = 'Submit';
          submitButton.style.backgroundColor = 'lightblue';
          submitButton.style.borderRadius = '5px';
          submitButton.style.border = '1px solid darkblue';
          submitButton.style.maxWidth = '100px';
          submitButton.style.minHeight = '25px';
          submitButton.style.marginTop = '10px';
          submitButton.style.marginRight = '10px';
          inputFieldButtonsContainer.appendChild(submitButton);

          // write code to add a button the removed the inputField and submitButton
          const removeButton = document.createElement('button');
          removeButton.textContent = 'Remove';
          removeButton.style.backgroundColor = 'lightcoral';
          removeButton.style.borderRadius = '5px';
          removeButton.style.border = '1px solid darkred';
          removeButton.style.maxWidth = '100px';
          removeButton.style.minHeight = '25px';
          removeButton.style.marginTop = '10px';
          inputFieldButtonsContainer.appendChild(removeButton);
          removeButton.addEventListener('click', () => {
            activeCell.node.removeChild(parentContainer);
          });

          if (
            (aiService === 'OpenAI API key' && !openAiApiKey) ||
            (aiService === 'Use Azure API' &&
              (!azureBaseUrl || !azureDeploymentName || !azureApiKey))
          ) {
            inputField.placeholder = placeholderDisabled;
            submitButton.disabled = true;
          }

          const handleAccept = async () => {
            let userInput = inputField.value;
            if (userInput !== '') {
              parentContainer.removeChild(inputContainer);
              let diffEditor: monaco.editor.IStandaloneDiffEditor | null = null;
              const renderEditor = (gen: string) => {
                try {
                  if (!diffEditor) {
                    createEditorComponents();
                  }
                  const modifiedModel = diffEditor!.getModel()!.modified;
                  const endLineNumber = modifiedModel.getLineCount();
                  const endColumn =
                    modifiedModel.getLineMaxColumn(endLineNumber);
                  modifiedModel.applyEdits([
                    {
                      range: new monaco.Range(
                        endLineNumber,
                        endColumn,
                        endLineNumber,
                        endColumn
                      ),
                      text: gen,
                      forceMoveMarkers: true
                    }
                  ]);
                } catch (error) {
                  console.log('Error rendering editor:', error);
                }
              };

              const createEditorComponents = () => {
                // generate the editor components
                // first, top level container to hold all diff related items
                const diffContainer = document.createElement('div');
                diffContainer.style.marginTop = '10px';
                diffContainer.style.display = 'flex';
                diffContainer.style.flexDirection = 'column';
                parentContainer.appendChild(diffContainer);

                // next, container to hold the diff editor
                const diffEditorContainer = document.createElement('div');
                diffEditorContainer.style.height = '200px';
                diffContainer.appendChild(diffEditorContainer);

                // finally, the diff editor itself
                const currentTheme =
                  document.body.getAttribute('data-jp-theme-light') === 'true'
                    ? 'vs'
                    : 'vs-dark';
                diffEditor = monaco.editor.createDiffEditor(
                  diffEditorContainer,
                  {
                    readOnly: true,
                    theme: currentTheme
                  }
                );
                diffEditor.setModel({
                  original: monaco.editor.createModel(oldCode, 'python'),
                  modified: monaco.editor.createModel('', 'python')
                });

                const diffButtonsContainer = document.createElement('div');
                diffButtonsContainer.style.marginTop = '10px';
                diffButtonsContainer.style.marginLeft = '70px';
                diffButtonsContainer.style.display = 'flex';
                diffButtonsContainer.style.flexDirection = 'row';
                diffContainer.appendChild(diffButtonsContainer);

                // Create "Accept" and "Reject" buttons
                const acceptButton = document.createElement('button');
                acceptButton.textContent = 'Accept';
                acceptButton.style.backgroundColor = 'lightblue';
                acceptButton.style.borderRadius = '5px';
                acceptButton.style.border = '1px solid darkblue';
                acceptButton.style.maxWidth = '100px';
                acceptButton.style.minHeight = '25px';
                acceptButton.style.marginRight = '10px';
                acceptButton.addEventListener('click', () => {
                  const modifiedCode = diffEditor!
                    .getModel()!
                    .modified.getValue();
                  activeCell.model.sharedModel.source = modifiedCode;
                  commands.execute('notebook:run-cell');
                  activeCell.node.removeChild(parentContainer);
                });
                diffButtonsContainer.appendChild(acceptButton);

                const rejectButton = document.createElement('button');
                rejectButton.textContent = 'Reject';
                rejectButton.style.backgroundColor = 'lightblue';
                rejectButton.style.borderRadius = '5px';
                rejectButton.style.border = '1px solid darkblue';
                rejectButton.style.maxWidth = '100px';
                rejectButton.style.minHeight = '25px';
                rejectButton.style.marginRight = '10px';
                rejectButton.addEventListener('click', () => {
                  activeCell.node.removeChild(parentContainer);
                  activeCell.model.sharedModel.source = oldCode;
                });
                diffButtonsContainer.appendChild(rejectButton);

                const editPromptButton = document.createElement('button');
                editPromptButton.textContent = 'Edit Prompt';
                editPromptButton.style.backgroundColor = 'lightblue';
                editPromptButton.style.borderRadius = '5px';
                editPromptButton.style.border = '1px solid darkblue';
                editPromptButton.style.maxWidth = '100px';
                editPromptButton.style.minHeight = '25px';
                editPromptButton.style.marginRight = '10px';
                editPromptButton.addEventListener('click', () => {
                  parentContainer.removeChild(diffContainer);
                  parentContainer.appendChild(inputContainer);
                  diffEditor = null;
                });
                diffButtonsContainer.appendChild(editPromptButton);

                // Handle Enter key press to trigger accept on accept/reject buttons
                diffButtonsContainer.addEventListener('keydown', event => {
                  if (event.key === 'Enter') {
                    event.preventDefault();
                    const activeElement = document.activeElement;
                    if (activeElement === acceptButton) {
                      acceptButton.click();
                    } else if (activeElement === rejectButton) {
                      rejectButton.click();
                    }
                  }
                });

                // Handle Escape key press to trigger reject on accept/reject buttons
                diffButtonsContainer.addEventListener('keydown', event => {
                  if (event.key === 'Escape') {
                    event.preventDefault();
                    rejectButton.click();
                  }
                });
              };

              const variablePattern = /@(\w+)/g;
              let match;
              let modifiedUserInput = userInput;
              while ((match = variablePattern.exec(userInput)) !== null) {
                try {
                  const variableName = match[1];
                  // get value of var using the getVariableValue function
                  const variableType = await getVariableValue(
                    `type(${variableName})`
                  );

                  // check if variableType is dataframe
                  // if it is, get columns and add to modifiedUserInput
                  if (variableType?.includes('DataFrame')) {
                    const variableColumns = await getVariableValue(
                      `${variableName}.columns`
                    );
                    modifiedUserInput += `\n${variableName} is a dataframe with the following columns: ${variableColumns}\n`;
                  } else if (variableType) {
                    const variableValue = await getVariableValue(variableName);
                    modifiedUserInput += `\nPrinting ${variableName} in Python returns the string ${variableValue}\n`;
                  }
                } catch (error) {
                  console.error(`Error accessing variable ${match[1]}:`, error);
                }
              }
              userInput = modifiedUserInput;

              if (aiService === 'OpenAI API key' && openAiApiKey) {
                try {
                  const openai = new OpenAI({
                    apiKey: openAiApiKey,
                    dangerouslyAllowBrowser: true,
                    baseURL: openAiBaseUrl ? openAiBaseUrl : undefined
                  });
                  const complete = async () => {
                    const response = await openai.embeddings.create({
                      model: EMBEDDING_MODEL,
                      input: userInput
                    });
                    const userInputEmbedding = response.data[0].embedding;
                    const similarities = embeddings.map(embedding =>
                      cosineSimilarity(embedding.embedding, userInputEmbedding)
                    );
                    const topSimilarities = similarities
                      .map((value, index) => ({ value, index }))
                      .sort((a, b) => b.value - a.value)
                      .slice(0, NUMBER_OF_SIMILAR_CELLS)
                      .map(e => embeddings[e.index].source);
                    const prompt = generatePrompt(
                      userInput,
                      oldCode,
                      topSimilarities
                    );
                    console.log(prompt);
                    const stream = await openai.chat.completions.create({
                      model: 'gpt-4-turbo-preview',
                      messages: [
                        {
                          role: 'system',
                          content: systemPrompt
                        },
                        {
                          role: 'user',
                          content: prompt
                        }
                      ],
                      stream: true
                    });
                    for await (const chunk of stream) {
                      renderEditor(chunk.choices[0]?.delta?.content || '');
                    }
                  };
                  complete();
                } catch (error) {
                  activeCell.node.removeChild(parentContainer);
                }
              } else if (aiService === 'Use Pretzel AI Server') {
                const options: any = {
                  method: 'POST',
                  headers: {
                    'Content-Type': 'application/json'
                  },
                  body: JSON.stringify({
                    oldCode,
                    userInput
                  })
                };

                try {
                  const response = await fetch(
                    'https://q8qeei2tn4.execute-api.us-west-1.amazonaws.com/default/pretzel_notebook',
                    options
                  );
                  const data = await response.json();
                  const gen = data.message;
                  renderEditor(gen);
                } catch (error) {
                  activeCell.model.sharedModel.source = `# Error: ${error}\n${oldCode}`;
                  activeCell.node.removeChild(parentContainer);
                }
              } else if (
                aiService === 'Use Azure API' &&
                azureBaseUrl &&
                azureDeploymentName &&
                azureApiKey
              ) {
                try {
                  const client = new OpenAIClient(
                    azureBaseUrl,
                    new AzureKeyCredential(azureApiKey)
                  );
                  const deploymentId = azureDeploymentName;
                  const prompt = [
                    `Write python code to do \n"""\n${userInput}\n"""\nThe previous code is\n"""\n${oldCode}\n"""\nReturn ONLY executable python code, no backticks`
                  ];

                  const result = await client.getCompletions(
                    deploymentId,
                    prompt
                  );

                  for (const choice of result.choices) {
                    renderEditor(choice.text);
                  }
                } catch (error) {
                  activeCell.model.sharedModel.source = `# Error: ${error}\n${oldCode}`;
                  activeCell.node.removeChild(parentContainer);
                }
              }
            }
          };

          // Handle submit button click to trigger accept
          submitButton.addEventListener('click', handleAccept);
        }
      }
    });

    const category = 'Cell Operations';
    palette.addItem({ command, category });

    app.commands.addKeyBinding({
      command,
      keys: ['Accel K'],
      selector: '.jp-Notebook'
    });
  }
};

export default extension;
