import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext";
import type { AssistantCitation } from "@t3tools/contracts";
import {
  serializeAssistantCitation,
  withAssistantCitationComment,
} from "@t3tools/shared/assistantCitations";
import {
  $applyNodeReplacement,
  $getNodeByKey,
  $getRoot,
  $isElementNode,
  DecoratorNode,
  HISTORY_PUSH_TAG,
  SKIP_DOM_SELECTION_TAG,
  type NodeKey,
  type SerializedLexicalNode,
  type Spread,
} from "lexical";
import { createContext, use, type ReactElement } from "react";
import type { AssistantCitationSourceAnchor } from "~/lib/assistantTextSelection";

import { AssistantCitationChip } from "./chat/AssistantCitationChip";
import { COMPOSER_INLINE_CHIP_DECORATOR_CLASS_NAME } from "./composerInlineChip";

type SerializedComposerCitationNode = Spread<
  {
    citation: AssistantCitation;
    source: string;
    type: "composer-citation";
    version: 1;
  },
  SerializedLexicalNode
>;

export type ComposerCitationCommentRequest = {
  previousValue: string;
  value: string;
  citationStart: number;
  sourceAnchor: AssistantCitationSourceAnchor;
};

export type ComposerCitationCommentTarget = {
  nodeKey: NodeKey;
  sourceAnchor?: AssistantCitationSourceAnchor;
  removeOnCancel?: boolean;
};

export const ComposerCitationCommentContext = createContext<{
  openComment: ComposerCitationCommentTarget | null;
  onOpenChange: (nodeKey: NodeKey, open: boolean) => void;
  onSubmitAndSend: () => void;
}>({ openComment: null, onOpenChange: () => {}, onSubmitAndSend: () => {} });

/** Consume a cite action once its controlled prompt has been committed to the editor. */
export function $consumeComposerCitationCommentRequest(requestRef: {
  current: ComposerCitationCommentRequest | null;
}): ComposerCitationCommentTarget | null {
  const request = requestRef.current;
  if (!request) return null;
  const root = $getRoot();
  const value = root.getTextContent();
  if (value === request.previousValue) return null;
  requestRef.current = null;
  if (value !== request.value) return null;

  // Controlled prompts use one paragraph with inline nodes and explicit line breaks.
  const paragraph = root.getFirstChild();
  if (!$isElementNode(paragraph)) return null;
  let offset = 0;
  for (const node of paragraph.getChildren()) {
    if (offset === request.citationStart && node instanceof ComposerCitationNode) {
      return {
        nodeKey: node.getKey(),
        sourceAnchor: request.sourceAnchor,
        removeOnCancel: true,
      };
    }
    offset += node.getTextContentSize();
  }
  return null;
}

function ComposerCitationDecorator(props: { citation: AssistantCitation; nodeKey: NodeKey }) {
  const [editor] = useLexicalComposerContext();
  const commentContext = use(ComposerCitationCommentContext);
  const commentTarget =
    commentContext.openComment?.nodeKey === props.nodeKey ? commentContext.openComment : null;
  const onSaveComment = (comment: string): boolean => {
    if (!editor.isEditable()) return false;
    let accepted = false;
    editor.update(
      () => {
        const node = $getNodeByKey(props.nodeKey);
        if (node instanceof ComposerCitationNode && node.isAttached()) {
          node.setComment(comment);
          accepted = true;
        }
      },
      { discrete: true, tag: [HISTORY_PUSH_TAG, SKIP_DOM_SELECTION_TAG] },
    );
    return accepted;
  };
  const onRemove = () => {
    if (!editor.isEditable()) return;
    editor.update(
      () => {
        const node = $getNodeByKey(props.nodeKey);
        if (node instanceof ComposerCitationNode) {
          node.selectPrevious();
          node.remove();
        }
      },
      { tag: HISTORY_PUSH_TAG },
    );
    editor.getRootElement()?.focus({ preventScroll: true });
  };

  return (
    <span
      className="inline-flex min-w-0 max-w-full"
      contentEditable={false}
      spellCheck={false}
      data-composer-citation-chip="true"
    >
      <AssistantCitationChip
        citation={props.citation}
        commentEditor={{
          open: commentTarget !== null,
          sourceAnchor: commentTarget?.sourceAnchor,
          onOpenChange: (open) => {
            if (open && !editor.isEditable()) return;
            commentContext.onOpenChange(props.nodeKey, open);
          },
          ...(commentTarget?.removeOnCancel ? { onCancel: onRemove } : {}),
          onSave: onSaveComment,
          onSaveAndSend: (comment) => {
            if (!onSaveComment(comment)) return false;
            commentContext.onSubmitAndSend();
            return true;
          },
        }}
        onRemove={onRemove}
      />
    </span>
  );
}

export class ComposerCitationNode extends DecoratorNode<ReactElement> {
  __citation: AssistantCitation;
  __source: string;

  static override getType(): string {
    return "composer-citation";
  }

  static override clone(node: ComposerCitationNode): ComposerCitationNode {
    return new ComposerCitationNode(node.__citation, node.__source, node.__key);
  }

  static override importJSON(serializedNode: SerializedComposerCitationNode): ComposerCitationNode {
    return $createComposerCitationNode(
      serializedNode.citation,
      serializedNode.source,
    ).updateFromJSON(serializedNode);
  }

  constructor(citation: AssistantCitation, source: string, key?: NodeKey) {
    super(key);
    this.__citation = citation;
    this.__source = source;
  }

  override exportJSON(): SerializedComposerCitationNode {
    const node = this.getLatest();
    return {
      ...super.exportJSON(),
      citation: node.__citation,
      source: node.__source,
      type: "composer-citation",
      version: 1,
    };
  }

  override createDOM(): HTMLElement {
    const dom = document.createElement("span");
    dom.className = `${COMPOSER_INLINE_CHIP_DECORATOR_CLASS_NAME} max-w-full`;
    return dom;
  }

  override updateDOM(): false {
    return false;
  }

  override getTextContent(): string {
    return this.getLatest().__source;
  }

  setComment(comment: string): this {
    const latest = this.getLatest();
    const citation = withAssistantCitationComment(latest.__citation, comment);
    if (citation.comment === latest.__citation.comment) return latest;
    const source = serializeAssistantCitation(citation);
    const writable = this.getWritable();
    writable.__citation = citation;
    writable.__source = source;
    return writable;
  }

  override isInline(): true {
    return true;
  }

  override decorate(): ReactElement {
    return <ComposerCitationDecorator citation={this.__citation} nodeKey={this.__key} />;
  }
}

export function $createComposerCitationNode(
  citation: AssistantCitation,
  source: string,
): ComposerCitationNode {
  return $applyNodeReplacement(new ComposerCitationNode(citation, source));
}