import { StreamLanguage } from '@codemirror/language';

type LatexState = { inMath: boolean; mathDelimiter: string };

export const latexLanguage = StreamLanguage.define<LatexState>({
  startState() {
    return { inMath: false, mathDelimiter: '' };
  },
  token(stream, state): string | null {
    if (stream.match('%')) {
      stream.skipToEnd();
      return 'comment';
    }

    if (!state.inMath && stream.match('$$')) {
      state.inMath = true;
      state.mathDelimiter = '$$';
      return 'keyword';
    }
    if (
      state.inMath &&
      state.mathDelimiter === '$$' &&
      stream.match('$$')
    ) {
      state.inMath = false;
      state.mathDelimiter = '';
      return 'keyword';
    }

    if (!state.inMath && stream.peek() === '$' && !stream.match('$$', false)) {
      stream.next();
      state.inMath = true;
      state.mathDelimiter = '$';
      return 'keyword';
    }
    if (
      state.inMath &&
      state.mathDelimiter === '$' &&
      stream.peek() === '$'
    ) {
      stream.next();
      state.inMath = false;
      state.mathDelimiter = '';
      return 'keyword';
    }


    if (!state.inMath && stream.match('\\(')) {
      state.inMath = true;
      state.mathDelimiter = '\\(';
      return 'keyword';
    }
    if (
      state.inMath &&
      state.mathDelimiter === '\\(' &&
      stream.match('\\)')
    ) {
      state.inMath = false;
      state.mathDelimiter = '';
      return 'keyword';
    }

    if (!state.inMath && stream.match('\\[')) {
      state.inMath = true;
      state.mathDelimiter = '\\[';
      return 'keyword';
    }
    if (
      state.inMath &&
      state.mathDelimiter === '\\[' &&
      stream.match('\\]')
    ) {
      state.inMath = false;
      state.mathDelimiter = '';
      return 'keyword';
    }

    if (state.inMath) {
      if (stream.match(/^\\[a-zA-Z@]+/)) return 'atom';
      if (stream.match(/^\\./)) return 'atom';
      if (stream.match(/^[{}]/)) return 'bracket';
      if (stream.match(/^[0-9]+(\.[0-9]+)?/)) return 'number';
      stream.next();
      return 'string';
    }

    if (stream.match(/^\\(begin|end)\b/)) return 'keyword';
    if (stream.match(/^\{[a-zA-Z*]+\}/)) return 'typeName';
    if (stream.match(/^\\[a-zA-Z@]+\*?/)) return 'tagName';
    if (stream.match(/^\\./)) return 'escape';
    if (stream.match(/^[{}]/)) return 'bracket';
    if (stream.match(/^[\[\]]/)) return 'squareBracket';
    if (stream.match(/^[0-9]+(\.[0-9]+)?/)) return 'number';

    stream.next();
    return null;
  },
});