import { EditorView, Decoration, DecorationSet, ViewPlugin, ViewUpdate } from '@codemirror/view'
import { RangeSetBuilder } from '@codemirror/state'
import { Shrimp } from '#/index'
import { type SyntaxNode } from '#parser/node'
import { log } from '#utils/utils'

const shrimp = new Shrimp()

export const shrimpHighlighter = ViewPlugin.fromClass(
  class {
    decorations: DecorationSet

    constructor(view: EditorView) {
      this.decorations = this.highlight(view)
    }

    update(update: ViewUpdate) {
      if (!update.docChanged) return

      this.decorations = this.highlight(update.view)
    }

    highlight(view: EditorView): DecorationSet {
      const builder = new RangeSetBuilder<Decoration>()
      const code = view.state.doc.toString()

      try {
        const tree = shrimp.parse(code)
        const decorations: { from: number; to: number; class: string }[] = []

        tree.iterate({
          enter: (node) => {
            const cls = tokenStyles[node.type.name]
            const isLeaf = node.children.length === 0
            if (cls && isLeaf) {
              decorations.push({ from: node.from, to: node.to, class: cls })
            }
          },
        })

        // Sort by position (required by RangeSetBuilder)
        decorations.sort((a, b) => a.from - b.from)

        for (const d of decorations) {
          builder.add(d.from, d.to, Decoration.mark({ class: d.class }))
        }
      } catch (error) {
        log('Parsing error in highlighter', error)
      }

      return builder.finish()
    }
  },
  {
    decorations: (v) => v.decorations,
  }
)

// Map node types to CSS classes
const tokenStyles: Record<string, string> = {
  keyword: 'tok-keyword',
  String: 'tok-string',
  StringFragment: 'tok-string',
  CurlyString: 'tok-string',
  Number: 'tok-number',
  Boolean: 'tok-bool',
  Null: 'tok-null',
  Identifier: 'tok-identifier',
  AssignableIdentifier: 'tok-variable-def',
  Comment: 'tok-comment',
  operator: 'tok-operator',
  Regex: 'tok-regex',
}