import { createHeadlessEditor } from '@lexical/headless'
import {
  $create,
  $getState,
  $getStateChange,
  $setState,
  createState,
  EditorConfig,
  LexicalNode,
  TextNode
} from 'lexical'
import { describe, expect, it } from 'vitest'

const DEFAULT_COLOR = 'inherit'
// This defines how the color property is parsed along with a default value
const colorState = createState('color', {
  parse: (value) => (typeof value === 'string' ? value : DEFAULT_COLOR)
})

export class ColoredNode extends TextNode {
  $config() {
    return this.config('colored', {
      extends: TextNode,
      // This defines the serialization of the color NodeState as
      // a flat property of the SerializedLexicalNode JSON instead of
      // nesting it in the '$' property
      stateConfigs: [{ flat: true, stateConfig: colorState }]
    })
  }

  createDOM(config: EditorConfig): HTMLElement {
    const element = super.createDOM(config)
    element.style.color = $getState(this, colorState)
    return element
  }

  updateDOM(prevNode: this, dom: HTMLElement, config: EditorConfig): boolean {
    if (super.updateDOM(prevNode, dom, config)) {
      return true
    }
    const colorChange = $getStateChange(this, prevNode, colorState)
    if (colorChange !== null) {
      dom.style.color = colorChange[0]
    }
    return false
  }

  // As per docs https://lexical.dev/docs/concepts/nodes#creating-custom-nodes-with-config-and-nodestate
  // Note that since these example use NodeState and $config, they will automatically get full and correct implementations of the following methods:
  // static clone
  // static importFromJSON
  // updateFromJSON
  // afterCloneFrom
  // exportJSON

  // Uncommenting this crude implementation will fix the problem

  //   updateFromJSON(serializedNode: LexicalUpdateJSON<SerializedLexicalNode>): this {
  //     const color = (serializedNode as any).color
  //     $setState(this, colorState, color)
  //     return this
  //   }
}

export function $createColoredNode(text: string, color: string): ColoredNode {
  // Since our constructor has 0 arguments, we set all of its properties
  // after construction.
  const node: ColoredNode = $create(ColoredNode)
  node.setTextContent(text)
  $setState(node, colorState, color)
  return node
}

export function $isColoredNode(node: LexicalNode | null | undefined): node is ColoredNode {
  return node instanceof ColoredNode
}

describe('Lexical bug', () => {
  it("new $config style custom node importJSON doesn't work as described", () => {
    const editor = createHeadlessEditor({
      namespace: 'test',
      nodes: [ColoredNode]
    })
    editor.update(
      () => {
        const coloredNode = $createColoredNode('Hello, world!', 'red')
        const color = $getState(coloredNode, colorState)
        expect(color).toBe('red')
        const json = coloredNode.exportJSON()
        console.log('ColoredNode exported', json)
        const coloredNode2 = ColoredNode.importJSON(json)
        console.log('ColoredNode imported', coloredNode2.exportJSON())
        const color2 = $getState(coloredNode2, colorState)
        expect(color2).toBe('red')
      },
      { discrete: true }
    )
  })
})