import { computed, defineComponent, h, onMounted, onUnmounted, ref, shallowRef } from 'vue'
import { $applyNodeReplacement } from 'lexical'
import VarLabel from '@/components/VarLabel/index.vue'
import type { NodeVar } from '@/types/var'
import { DecoratorBlockNode } from 'lexical-vue/LexicalDecoratorBlockNode'
import type { Component, PropType, ShallowRef } from 'vue'
import type {
	DOMConversionMap,
	DOMConversionOutput,
	DOMExportOutput,
	EditorConfig,
	ElementFormatType,
	LexicalEditor,
	LexicalNode,
	NodeKey,
	Spread
} from 'lexical'
import type { SerializedDecoratorBlockNode } from 'lexical-vue'

export type SerializedVarLabelNode = Spread<
	{
		className: string
		text: string
	},
	SerializedDecoratorBlockNode
>

type LegacySerializedVarLabelNode = {
	className?: unknown
	text?: unknown
	format?: unknown
}

const DATA_ATTR_CONTEXT_ID = 'data-prompt-editor-context-id'

const emptyNodeVars: NodeVar[] = []
const promptEditorNodeVarsRefs = new Map<string, ShallowRef<NodeVar[]>>()

const VarLabelDecorator = defineComponent({
	name: 'PromptEditorVarLabelDecorator',
	props: {
		label: {
			type: String,
			required: true
		},
		editor: {
			type: Object as PropType<LexicalEditor>,
			required: true
		}
	},
	setup(props) {
		const contextId = ref('')

		const syncContextId = (attempt = 0) => {
			const nextContextId =
				props.editor
					.getRootElement()
					?.closest<HTMLElement>(`[${DATA_ATTR_CONTEXT_ID}]`)
					?.getAttribute(DATA_ATTR_CONTEXT_ID) || ''

			if (contextId.value !== nextContextId) {
				contextId.value = nextContextId
			}

			if (!nextContextId && attempt < 6) {
				requestAnimationFrame(() => {
					syncContextId(attempt + 1)
				})
			}
		}

		let removeRootListener: (() => void) | null = null

		onMounted(() => {
			syncContextId()
			removeRootListener = props.editor.registerRootListener(() => {
				syncContextId()
			})
		})

		onUnmounted(() => {
			removeRootListener?.()
		})

		const nodeVars = computed(() => {
			if (!contextId.value) return emptyNodeVars
			return getNodeVarsRef(contextId.value).value
		})

		return () =>
			h(VarLabel, {
				label: props.label,
				nodeVars: nodeVars.value
			})
	}
})

function convertVarLabelElement(domNode: HTMLElement): null | DOMConversionOutput {
	const text = domNode.getAttribute('data-lexical-var-label')
	if (!text) return null

	const className = domNode.getAttribute('data-lexical-var-label-class') || 'custom-label'
	const node = $createVarLabelNode(className, text)
	return { node }
}

function getNodeVarsRef(contextId: string) {
	const existingRef = promptEditorNodeVarsRefs.get(contextId)
	if (existingRef) return existingRef

	const nextRef = shallowRef<NodeVar[]>([])
	promptEditorNodeVarsRefs.set(contextId, nextRef)
	return nextRef
}

export function syncPromptEditorVarLabelNodeVars(
	contextId: string,
	nodeVars: NodeVar[] | undefined
) {
	if (!contextId) return

	const nextNodeVars =
		nodeVars?.map((item) => ({
			...item,
			variableList: [...item.variableList]
		})) || emptyNodeVars

	getNodeVarsRef(contextId).value = nextNodeVars
}

export class VarLabelNode extends DecoratorBlockNode {
	__className: string
	__text: string

	static getType() {
		return 'var-label'
	}

	static clone(node: VarLabelNode) {
		return new VarLabelNode(node.__className, node.__text, node.__format, node.__key)
	}

	static importJSON(serializedNode: SerializedVarLabelNode): VarLabelNode {
		const legacyNode = serializedNode as SerializedVarLabelNode & LegacySerializedVarLabelNode
		const className =
			typeof legacyNode.className === 'string' && legacyNode.className.length
				? legacyNode.className
				: 'custom-label'
		const text = typeof legacyNode.text === 'string' ? legacyNode.text : ''
		const format: ElementFormatType = typeof legacyNode.format === 'string' ? legacyNode.format : ''

		return new VarLabelNode(className, text, format)
	}

	constructor(className: string, text: string, format?: ElementFormatType, key?: NodeKey) {
		super(format, key)
		this.__className = className
		this.__text = text
	}

	isInline() {
		return true
	}

	canInsertTextBefore() {
		return false
	}

	canInsertTextAfter() {
		return false
	}

	isKeyboardSelectable() {
		return false
	}

	createDOM(_config: EditorConfig) {
		const dom = document.createElement('span')
		dom.className = `${this.__className} var-label-token`
		dom.contentEditable = 'false'
		return dom
	}

	updateDOM(prevNode: VarLabelNode, dom: HTMLElement, _config: EditorConfig) {
		if (prevNode.__className !== this.__className) {
			dom.className = `${this.__className} var-label-token`
		}
		dom.contentEditable = 'false'

		return false
	}

	static importDOM(): DOMConversionMap | null {
		return {
			span: (domNode: HTMLElement) => {
				if (!domNode.hasAttribute('data-lexical-var-label')) return null

				return {
					conversion: convertVarLabelElement,
					priority: 1
				}
			}
		}
	}

	exportDOM(): DOMExportOutput {
		const element = document.createElement('span')
		element.setAttribute('data-lexical-var-label', this.__text)
		element.setAttribute('data-lexical-var-label-class', this.__className)
		element.textContent = this.__text
		return { element }
	}

	getTextContent() {
		return this.__text
	}

	decorate(_editor: LexicalEditor, _config: EditorConfig): Component {
		return h(VarLabelDecorator, {
			label: this.__text,
			editor: _editor
		})
	}

	exportJSON(): SerializedVarLabelNode {
		return {
			...super.exportJSON(),
			type: 'var-label',
			className: this.__className,
			text: this.__text,
			version: 1
		}
	}
}

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

export function $createVarLabelNode(className: string, text: string) {
	return $applyNodeReplacement(new VarLabelNode(className, text))
}
