import { Plugin, PluginKey } from 'prosemirror-state';
import { Decoration, DecorationSet } from 'prosemirror-view';
import { Node as PMNode } from 'prosemirror-model';
import hljs from 'highlight.js/lib/core';
import javascript from 'highlight.js/lib/languages/javascript';
import typescript from 'highlight.js/lib/languages/typescript';
import json from 'highlight.js/lib/languages/json';
import xml from 'highlight.js/lib/languages/xml';
import css from 'highlight.js/lib/languages/css';
import scss from 'highlight.js/lib/languages/scss';
import sql from 'highlight.js/lib/languages/sql';
import python from 'highlight.js/lib/languages/python';
import java from 'highlight.js/lib/languages/java';
import bash from 'highlight.js/lib/languages/bash';
import shell from 'highlight.js/lib/languages/shell';
import markdown from 'highlight.js/lib/languages/markdown';
import plaintext from 'highlight.js/lib/languages/plaintext';
import c from 'highlight.js/lib/languages/c';
import cpp from 'highlight.js/lib/languages/cpp';
import csharp from 'highlight.js/lib/languages/csharp';
import go from 'highlight.js/lib/languages/go';
import rust from 'highlight.js/lib/languages/rust';
import php from 'highlight.js/lib/languages/php';
import ruby from 'highlight.js/lib/languages/ruby';
import kotlin from 'highlight.js/lib/languages/kotlin';
import swift from 'highlight.js/lib/languages/swift';
import yaml from 'highlight.js/lib/languages/yaml';
import dockerfile from 'highlight.js/lib/languages/dockerfile';
import diff from 'highlight.js/lib/languages/diff';
import graphql from 'highlight.js/lib/languages/graphql';

hljs.registerLanguage('javascript', javascript);
hljs.registerLanguage('typescript', typescript);
hljs.registerLanguage('json', json);
hljs.registerLanguage('xml', xml);
hljs.registerAliases(['html'], { languageName: 'xml' });
hljs.registerLanguage('css', css);
hljs.registerLanguage('scss', scss);
hljs.registerLanguage('sql', sql);
hljs.registerLanguage('python', python);
hljs.registerLanguage('java', java);
hljs.registerLanguage('bash', bash);
hljs.registerLanguage('shell', shell);
hljs.registerLanguage('markdown', markdown);
hljs.registerLanguage('plaintext', plaintext);
hljs.registerLanguage('c', c);
hljs.registerLanguage('cpp', cpp);
hljs.registerLanguage('csharp', csharp);
hljs.registerLanguage('go', go);
hljs.registerLanguage('rust', rust);
hljs.registerLanguage('php', php);
hljs.registerLanguage('ruby', ruby);
hljs.registerLanguage('kotlin', kotlin);
hljs.registerLanguage('swift', swift);
hljs.registerLanguage('yaml', yaml);
hljs.registerLanguage('dockerfile', dockerfile);
hljs.registerLanguage('diff', diff);
hljs.registerLanguage('graphql', graphql);

export const codeHighlightKey = new PluginKey<DecorationSet>('nileCodeHighlight');

/** Languages offered in the code block language picker. '' = auto-detect. */
export const CODE_LANGUAGES: { value: string; label: string }[] = [
  { value: '', label: 'Auto' },
  { value: 'plaintext', label: 'Plain text' },
  { value: 'javascript', label: 'JavaScript' },
  { value: 'typescript', label: 'TypeScript' },
  { value: 'json', label: 'JSON' },
  { value: 'html', label: 'HTML' },
  { value: 'css', label: 'CSS' },
  { value: 'scss', label: 'SCSS' },
  { value: 'sql', label: 'SQL' },
  { value: 'python', label: 'Python' },
  { value: 'java', label: 'Java' },
  { value: 'c', label: 'C' },
  { value: 'cpp', label: 'C++' },
  { value: 'csharp', label: 'C#' },
  { value: 'go', label: 'Go' },
  { value: 'rust', label: 'Rust' },
  { value: 'php', label: 'PHP' },
  { value: 'ruby', label: 'Ruby' },
  { value: 'kotlin', label: 'Kotlin' },
  { value: 'swift', label: 'Swift' },
  { value: 'bash', label: 'Bash' },
  { value: 'shell', label: 'Shell' },
  { value: 'yaml', label: 'YAML' },
  { value: 'dockerfile', label: 'Dockerfile' },
  { value: 'graphql', label: 'GraphQL' },
  { value: 'diff', label: 'Diff' },
  { value: 'markdown', label: 'Markdown' },
];

const AUTO_DETECT_SUBSET = [
  'javascript', 'typescript', 'json', 'xml', 'css', 'sql', 'python', 'java',
  'bash', 'go', 'rust', 'ruby',
];

const MAX_HIGHLIGHT_LENGTH = 20000;

/**
 * Converts highlight.js' emitted HTML into inline decorations. hljs escapes
 * but never reorders text, so parsed text offsets map 1:1 onto the node.
 */
function decorationsForCode(text: string, language: string | null, base: number): Decoration[] {
  if (!text || text.length > MAX_HIGHLIGHT_LENGTH) return [];
  let value: string;
  try {
    if (language && hljs.getLanguage(language)) {
      value = hljs.highlight(text, { language }).value;
    } else if (language === 'plaintext') {
      return [];
    } else {
      value = hljs.highlightAuto(text, AUTO_DETECT_SUBSET).value;
    }
  } catch {
    return [];
  }

  const template = document.createElement('template');
  template.innerHTML = value;
  const decorations: Decoration[] = [];
  let offset = 0;
  const walk = (node: globalThis.Node, classes: string[]): void => {
    node.childNodes.forEach(child => {
      if (child.nodeType === 3) {
        const length = child.nodeValue?.length ?? 0;
        if (classes.length && length) {
          decorations.push(
            Decoration.inline(base + offset, base + offset + length, {
              class: classes.join(' '),
            })
          );
        }
        offset += length;
      } else if (child.nodeType === 1) {
        const el = child as Element;
        walk(el, [...classes, ...el.className.split(/\s+/).filter(Boolean)]);
      }
    });
  };
  walk(template.content, []);
  return decorations;
}

function buildDecorations(doc: PMNode): DecorationSet {
  const decorations: Decoration[] = [];
  doc.descendants((node, pos) => {
    if (node.type.name === 'code_block') {
      decorations.push(
        ...decorationsForCode(node.textContent, node.attrs.language, pos + 1)
      );
      return false;
    }
    return true;
  });
  return decorations.length ? DecorationSet.create(doc, decorations) : DecorationSet.empty;
}

/** Syntax highlighting for code blocks, recomputed on document changes. */
export function codeHighlightPlugin(): Plugin<DecorationSet> {
  return new Plugin<DecorationSet>({
    key: codeHighlightKey,
    state: {
      init: (_config, state) => buildDecorations(state.doc),
      apply: (tr, old) => (tr.docChanged ? buildDecorations(tr.doc) : old),
    },
    props: {
      decorations(state) {
        return codeHighlightKey.getState(state);
      },
    },
  });
}