import {SchemaExport} from "@fiduswriter/document/schema/export"
import type {BibDB, ImageDB} from "@fiduswriter/document"
import {Dialog, activateWait, deactivateWait, makeWorker} from "fwtoolkit"
import {E2EEEncryptor} from "fwtoolkit/e2ee/encryptor"
import {E2EEKeyManager} from "fwtoolkit/e2ee/key-manager"
import {enterPassphraseDialog} from "fwtoolkit/e2ee/passphrase-dialog"
import {PassphraseManager} from "fwtoolkit/e2ee/passphrase-manager"
import {enterPasswordDialog} from "fwtoolkit/e2ee/password-dialog"
import {receiveTransaction, sendableSteps} from "prosemirror-collab"
import type {Node} from "prosemirror-model"
import {EditorState} from "prosemirror-state"
import type {Transaction} from "prosemirror-state"
import {Step} from "prosemirror-transform"

import {
    getSelectionUpdate,
    removeCollaboratorSelection,
    updateCollaboratorSelection
} from "../state_plugins/index.js"
import {Merge} from "./merge/index.js"
import type {ModCollab} from "./index.js"

interface IncomingDoc {
    v: number
    content: unknown
    comments: unknown
    bibliography: unknown
    images?: Record<string, {copyright?: unknown}>
    template?: {content: unknown}
}

interface DiffData {
    type?: string
    v?: number
    rid?: number
    cid?: string | number
    ds?: unknown[]
    fs?: unknown[]
    cu?: unknown[]
    bu?: unknown[]
    iu?: unknown[]
    footnoterender?: boolean
    reject_request_id?: number
    ep?: string
    e2ee_salt?: string
}

export class ModCollabDoc {
    mod: ModCollab
    merge: Merge
    unconfirmedDiffs: Record<number, any>
    confirmStepsRequestCounter: number
    awaitingDiffResponse: boolean
    receiving: boolean
    currentlyCheckingVersion: boolean
    enableCheckVersion?: number
    footnoteRender: boolean
    initialVersionConfirmed: boolean
    initialDocLoaded: boolean
    templateAdjustmentPending: boolean
    sendNextDiffTimer?: number
    lastSelectionUpdateState?: EditorState

    constructor(mod: ModCollab) {
        mod.doc = this
        this.mod = mod
        this.merge = new Merge(mod)
        this.unconfirmedDiffs = {}
        this.confirmStepsRequestCounter = 0
        this.awaitingDiffResponse = false
        this.receiving = false
        this.currentlyCheckingVersion = false
        this.footnoteRender = false // If the offline user edited a footnote , it needs to be rendered properly to connected users too!
        this.initialVersionConfirmed = false
        this.initialDocLoaded = false
        this.templateAdjustmentPending = false
    }

    cancelCurrentlyCheckingVersion(): void {
        this.currentlyCheckingVersion = false
        window.clearTimeout(this.enableCheckVersion)
    }

    checkVersion(offline = false): void {
        // Guard: WebSocket may not be initialized yet (e.g. during
        // loadDocument before the WS connector is created).
        if (!this.mod.editor.ws) {
            return
        }
        const ws = this.mod.editor.ws
        ws.send(() => {
            if (
                this.currentlyCheckingVersion ||
                this.mod.editor.docInfo.version === undefined
            ) {
                return
            }
            this.currentlyCheckingVersion = true
            this.enableCheckVersion = window.setTimeout(() => {
                this.currentlyCheckingVersion = false
            }, 1000)
            if (ws.connected) {
                this.disableDiffSending()
            }

            const msg: Record<string, unknown> = {
                type: "check_version",
                v: this.mod.editor.docInfo.version
            }
            if (offline) {
                msg.offline = true
            }
            return msg
        })
    }

    disableDiffSending(): void {
        this.awaitingDiffResponse = true
        // If no answer has been received from the server within 2 seconds,
        // check the version
        this.sendNextDiffTimer = window.setTimeout(() => {
            this.awaitingDiffResponse = false
            this.sendToCollaborators()
        }, 8000)
    }

    enableDiffSending(): void {
        window.clearTimeout(this.sendNextDiffTimer)
        this.awaitingDiffResponse = false
        this.sendToCollaborators()
    }

    receiveDocument(data: {doc: IncomingDoc}): void {
        this.cancelCurrentlyCheckingVersion()
        if (
            this.mod.editor.docInfo.confirmedDoc &&
            !this.mod.editor.e2ee?.encrypted
        ) {
            this.merge.adjustDocument(data as any)
        } else {
            this.loadDocument(data as any)
        }
    }

    finishInitialLoad(): void {
        if (
            this.initialDocLoaded &&
            this.initialVersionConfirmed &&
            !this.templateAdjustmentPending
        ) {
            deactivateWait()
            this.mod.editor.waitingForDocument = false
        }
    }

    confirmVersion(version: number): void {
        if (!this.mod.editor.docInfo.confirmedDoc) {
            return
        }
        if (!this.initialVersionConfirmed) {
            this.initialVersionConfirmed = true
            if (version === this.mod.editor.docInfo.version) {
                this.finishInitialLoad()
            }
        }
    }

    loadDocument(data: {
        doc: IncomingDoc
        time: number
        doc_info: Record<string, any>
    }): void {
        const isInitialLoad = !this.mod.editor.docInfo.confirmedDoc
        // Reset collaboration
        this.unconfirmedDiffs = {}
        if (this.awaitingDiffResponse) {
            this.enableDiffSending()
        }

        this.mod.editor.clientTimeAdjustment = Date.now() - data.time

        const token = this.mod.editor.docInfo.token
        this.mod.editor.docInfo = data.doc_info as any
        this.mod.editor.docInfo.token = token
        // For guests, update user object with the token UUID (stable identity)
        // and session_id (for display). Token UUID persists across reconnections.
        // session_id may not be present when loading from REST (it comes via
        // the session_info WebSocket message). Only update if available.
        if (
            !this.mod.editor.user.is_authenticated &&
            data.doc_info.session_id !== undefined
        ) {
            const sessionId = data.doc_info.session_id as string
            this.mod.editor.user = {
                id: token as unknown as number, // stable — same token UUID even after reconnect
                username: `guest${sessionId}`, // display name per session
                name: `Guest ${sessionId}`,
                is_authenticated: false
            } as any
        }
        this.mod.editor.docInfo.token = token
        this.mod.editor.docInfo.version = data.doc.v
        this.mod.editor.docInfo.updated = new Date()

        // Check if this is an E2EE document. If so, we need to decrypt
        // the content before loading it into ProseMirror.
        const isE2EE = data.doc_info.e2ee === true
        if (isE2EE) {
            this._loadE2EEDocument(data.doc, data.doc_info, isInitialLoad)
        } else {
            this.mod.editor.e2ee = {encrypted: false}
            this._loadUnencryptedDocument(data.doc, isInitialLoad)
        }
    }

    /**
     * Load an E2EE document: try passphrase/DEK first, then fall back to
     * per-document password prompt, decrypt, then load.
     * @private
     */
    async _loadE2EEDocument(
        doc: IncomingDoc,
        doc_info: Record<string, any>,
        isInitialLoad: boolean,
        urlFragmentPassword = ""
    ): Promise<void> {
        // Extract password from URL fragment if present (share link format:
        // https://example.com/share/TOKEN/#target?password=PASSWORD). The fragment is never
        // sent to the server, so this is safe.
        const locationHash = window.location.hash
        if (locationHash && locationHash.includes("?")) {
            urlFragmentPassword = decodeURIComponent(
                new URLSearchParams(locationHash.split("?")[1]).get(
                    "password"
                ) || ""
            )
        }

        // Get the salt and iterations. These may come from the REST response
        // (doc_info) or from the WebSocket session_info message (already
        // stored in this.mod.editor.e2ee).
        let salt: string | undefined, iterations: number | undefined
        if (this.mod.editor.e2ee) {
            salt = this.mod.editor.e2ee.encryptionSalt
            iterations = this.mod.editor.e2ee.encryptionIterations
        }
        if (doc_info.e2ee_salt) {
            salt = doc_info.e2ee_salt
        }
        if (doc_info.e2ee_iterations) {
            iterations = doc_info.e2ee_iterations
        }

        const docId = this.mod.editor.docInfo.id as number as number

        // Helper to save password for passphrase users after successful decryption
        // (migrates per-document-password docs into the passphrase system)
        const maybeSavePasswordForPassphrase = async (password: string) => {
            if (!PassphraseManager.hasKeysInSession()) {
                return
            }
            try {
                const existing =
                    await PassphraseManager.getDocumentPassword(docId)
                if (!existing) {
                    await PassphraseManager.saveDocumentPassword(
                        docId,
                        password,
                        null as any,
                        "user",
                        true
                    )
                }
            } catch (_e) {
                console.error(
                    "E2EE: Failed to save document password for passphrase:",
                    _e
                )
            }
        }

        // Helper to resolve a password to a key and attempt decryption
        const tryPassword = async (password: string) => {
            const key = await E2EEKeyManager.resolvePasswordToKey(
                password,
                salt
                    ? new Uint8Array(
                          atob(salt)
                              .split("")
                              .map(c => c.charCodeAt(0))
                      )
                    : (null as any),
                iterations || 600000
            )
            await this._decryptAndLoadDoc(
                doc,
                key,
                salt,
                iterations,
                isInitialLoad,
                urlFragmentPassword
            )
            // Cache password and key for this session
            E2EEKeyManager.storePasswordInSession(docId, password)
            await E2EEKeyManager.storeKeyInSession(docId, key)
            if (this.mod.editor.e2ee) {
                this.mod.editor.e2ee.key = key
                this.mod.editor.e2ee.password = password
            }
            // Migrate per-document-password docs to passphrase system
            await maybeSavePasswordForPassphrase(password)
        }

        // --- Try passphrase path first ---
        if (PassphraseManager.hasKeysInSession()) {
            const password = await PassphraseManager.getDocumentPassword(docId)
            if (password) {
                try {
                    await tryPassword(password)
                    if (this.mod.editor.e2ee) {
                        this.mod.editor.e2ee.usesPassphrase = true
                    }
                    return
                } catch (_error) {
                    // Password didn't work — fall through
                }
            }
        }

        // If passphrase keys exist but are not in session, try to unlock
        if (!PassphraseManager.hasKeysInSession()) {
            const hasKeys = await PassphraseManager.hasEncryptionKeys()
            if (hasKeys) {
                let errorMessage = ""
                let unlocked = false
                while (!unlocked) {
                    const result = await new Promise<{
                        action: string
                        passphrase?: string
                    }>(resolve => {
                        enterPassphraseDialog(
                            (pwd: string) =>
                                resolve({action: "unlock", passphrase: pwd}),
                            () => resolve({action: "recover"}),
                            {errorMessage}
                        )
                    })
                    if (result.action === "unlock" && result.passphrase) {
                        try {
                            await PassphraseManager.unlockWithPassphrase(
                                result.passphrase
                            )
                            const password =
                                await PassphraseManager.getDocumentPassword(
                                    docId
                                )
                            if (password) {
                                await tryPassword(password)
                                if (this.mod.editor.e2ee) {
                                    this.mod.editor.e2ee.usesPassphrase = true
                                }
                                unlocked = true
                                return
                            }
                        } catch (_e) {
                            errorMessage = gettext(
                                "Incorrect passphrase. Please try again."
                            )
                        }
                    } else if (result?.action === "recover") {
                        // Recovery flow
                        const {recoverWithKeyDialog} = await import(
                            "fwtoolkit/e2ee/passphrase-dialog"
                        )
                        const recoverResult = await new Promise<{
                            recoveryKey: string
                            newPassphrase: string
                        } | null>(resolve => {
                            recoverWithKeyDialog(resolve)
                        })
                        if (recoverResult) {
                            try {
                                const {newRecoveryKey} =
                                    (await PassphraseManager.recoverWithRecoveryKey(
                                        recoverResult.recoveryKey,
                                        recoverResult.newPassphrase
                                    )) as {newRecoveryKey: string}
                                const {showRecoveryKeyDialog} = await import(
                                    "fwtoolkit/e2ee/passphrase-dialog"
                                )
                                await new Promise<void>(resolve => {
                                    showRecoveryKeyDialog(
                                        newRecoveryKey,
                                        resolve
                                    )
                                })
                                // After recovery, try to get document password again
                                const password =
                                    await PassphraseManager.getDocumentPassword(
                                        docId
                                    )
                                if (password) {
                                    await tryPassword(password)
                                    if (this.mod.editor.e2ee) {
                                        this.mod.editor.e2ee.usesPassphrase = true
                                    }
                                    return
                                }
                            } catch (e) {
                                console.error("E2EE: Recovery failed:", e)
                                const errorDialog = new Dialog({
                                    title: gettext("Recovery Failed"),
                                    id: "e2ee-recovery-failed",
                                    body: gettext(
                                        "The recovery key you entered is incorrect, or the recovery process failed. Please try again."
                                    ),
                                    buttons: [
                                        {
                                            text: gettext("Retry"),
                                            classes: "fw-dark",
                                            click: () => {
                                                errorDialog.close()
                                                this._loadE2EEDocument(
                                                    doc,
                                                    doc_info,
                                                    isInitialLoad,
                                                    urlFragmentPassword
                                                )
                                            }
                                        },
                                        {
                                            text: gettext("Cancel"),
                                            classes: "fw-light",
                                            click: () => errorDialog.close()
                                        }
                                    ],
                                    canClose: false
                                })
                                errorDialog.open()
                                return
                            }
                        }
                        // If recovery was cancelled, continue the loop to show
                        // the passphrase dialog again.
                        continue
                    } else {
                        // User cancelled the passphrase dialog — break out of
                        // the loop and fall through to other password sources.
                        break
                    }
                }
            }
        }

        // --- Try sessionStorage password ---
        const sessionPassword = E2EEKeyManager.getPasswordFromSession(docId)
        if (sessionPassword) {
            try {
                await tryPassword(sessionPassword)
                return
            } catch (_error) {
                E2EEKeyManager.clearPasswordFromSession(docId)
                E2EEKeyManager.clearKeyFromSession(docId)
            }
        }

        // --- Try URL fragment password ---
        if (urlFragmentPassword) {
            try {
                await tryPassword(urlFragmentPassword)
                return
            } catch (_error) {
                // URL password didn't work — fall through to prompt
            }
        }

        // --- Try existing key (e.g. from _createE2EEDocument) ---
        const existingKey = this.mod.editor.e2ee?.key
        if (existingKey) {
            try {
                await this._decryptAndLoadDoc(
                    doc,
                    existingKey,
                    salt,
                    iterations,
                    isInitialLoad,
                    urlFragmentPassword
                )
            } catch (error) {
                console.error(
                    "E2EE: Decryption failed with existing key",
                    error
                )
                if (this.mod.editor.e2ee) {
                    this.mod.editor.e2ee.key = null as any
                }
                await this._promptPasswordAndDecrypt(
                    doc,
                    doc_info,
                    salt,
                    iterations,
                    isInitialLoad,
                    urlFragmentPassword
                )
            }
            return
        }

        // No key available — prompt for the password
        await this._promptPasswordAndDecrypt(
            doc,
            doc_info,
            salt,
            iterations,
            isInitialLoad,
            urlFragmentPassword
        )
    }

    /**
     * Prompt for a password, derive the key, and decrypt/load the document.
     * @private
     */
    async _promptPasswordAndDecrypt(
        doc: IncomingDoc,
        doc_info: Record<string, any>,
        salt: string | undefined,
        iterations: number | undefined,
        isInitialLoad: boolean,
        urlFragmentPassword = ""
    ): Promise<void> {
        // Decode the salt from Base64 to Uint8Array
        let saltBytes: Uint8Array | null = null
        if (salt) {
            const binary = atob(salt)
            saltBytes = new Uint8Array(binary.length)
            for (let i = 0; i < binary.length; i++) {
                saltBytes[i] = binary.charCodeAt(i)
            }
        }

        const goBack = () => {
            const folderPath = this.mod.editor.docInfo.path.slice(
                0,
                this.mod.editor.docInfo.path.lastIndexOf("/")
            )
            if (this.mod.editor.app && this.mod.editor.app.goTo) {
                if (!folderPath.length) {
                    this.mod.editor.app.goTo("/")
                } else {
                    this.mod.editor.app.goTo(`/documents${folderPath}/`)
                }
            } else {
                window.location.href = "/"
            }
        }

        await enterPasswordDialog(
            async (password: string) => {
                try {
                    const key = await E2EEKeyManager.resolvePasswordToKey(
                        password,
                        saltBytes as any,
                        iterations || 600000
                    )
                    await this._decryptAndLoadDoc(
                        doc,
                        key,
                        salt,
                        iterations,
                        isInitialLoad,
                        urlFragmentPassword
                    )
                    // Cache password and key for this session
                    E2EEKeyManager.storePasswordInSession(
                        this.mod.editor.docInfo.id as number,
                        password
                    )
                    await E2EEKeyManager.storeKeyInSession(
                        this.mod.editor.docInfo.id as number,
                        key
                    )
                    if (this.mod.editor.e2ee) {
                        this.mod.editor.e2ee.key = key
                        this.mod.editor.e2ee.password = password
                    }
                    // Migrate per-document-password docs to passphrase system
                    if (PassphraseManager.hasKeysInSession()) {
                        try {
                            const existing =
                                await PassphraseManager.getDocumentPassword(
                                    this.mod.editor.docInfo.id as number
                                )
                            if (!existing) {
                                await PassphraseManager.saveDocumentPassword(
                                    this.mod.editor.docInfo.id as number,
                                    password,
                                    null as any,
                                    "user",
                                    true
                                )
                            }
                        } catch (_e) {
                            console.error(
                                "E2EE: Failed to save document password for passphrase:",
                                _e
                            )
                        }
                    }
                } catch (_error) {
                    console.error("E2EE DECRYPT ERROR:", _error)
                    const errorDialog = new Dialog({
                        title: gettext("Decryption Failed"),
                        id: "e2ee-decryption-failed",
                        body: gettext(
                            "That password didn't work. Please try again."
                        ),
                        buttons: [
                            {
                                text: gettext("Retry"),
                                classes: "fw-dark",
                                click: () => {
                                    this._promptPasswordAndDecrypt(
                                        doc,
                                        doc_info,
                                        salt,
                                        iterations,
                                        isInitialLoad,
                                        ""
                                    )
                                }
                            },
                            {
                                text: gettext("Cancel"),
                                classes: "fw-light",
                                click: () => {
                                    errorDialog.close()
                                    goBack()
                                }
                            }
                        ],
                        canClose: false
                    })
                    errorDialog.open()
                }
            },
            urlFragmentPassword,
            goBack as any
        )
    }

    /**
     * Decrypt document content/comments/bibliography with the given key
     * and load into ProseMirror.
     *
     * For a newly created E2EE document, no encrypted snapshot has been
     * saved yet, so the content/comments/bibliography are still plaintext
     * JSON objects (the template content). We detect this by checking
     * whether the value is a string (encrypted = Base64 ciphertext) or
     * an object (plaintext).
     *
     * @private
     */
    async _decryptAndLoadDoc(
        doc: IncomingDoc,
        key: CryptoKey,
        salt: string | undefined,
        iterations: number | undefined,
        isInitialLoad: boolean,
        urlFragmentPassword: string
    ): Promise<void> {
        let decryptedContent = doc.content
        if (typeof doc.content === "string") {
            decryptedContent = await E2EEEncryptor.decryptObject(
                doc.content,
                key
            )
        }

        let decryptedComments = doc.comments
        if (typeof doc.comments === "string") {
            decryptedComments = await E2EEEncryptor.decryptObject(
                doc.comments,
                key
            )
        }

        let decryptedBibliography = doc.bibliography
        if (typeof doc.bibliography === "string") {
            decryptedBibliography = await E2EEEncryptor.decryptObject(
                doc.bibliography,
                key
            )
        }

        // Store the E2EE state on the editor
        const e2ee = this.mod.editor.e2ee = {
            encrypted: true,
            encryptionSalt: salt,
            encryptionIterations: iterations || 600000,
            key: key,
            snapshotManager: this.mod.editor.e2ee?.snapshotManager || undefined
        }
        // Initialize snapshot manager now that we have the key
        if (!e2ee.snapshotManager) {
            const {E2EESnapshotManager} = await import(
                "../e2ee/snapshot-manager.js"
            )
            e2ee.snapshotManager = new E2EESnapshotManager(
                this.mod.editor
            )
        }
        ;(e2ee.snapshotManager as any).setKey(key)

        // Cache the key in sessionStorage so the user doesn't have to
        // re-enter the password when reopening the document this session.
        try {
            await E2EEKeyManager.storeKeyInSession(
                this.mod.editor.docInfo.id as number,
                key
            )
        } catch (_e) {
            // If the key is non-extractable, we can't store it.
        }

        // Remove the password from the URL fragment to avoid
        // leaving it in the address bar or browser history.
        if (urlFragmentPassword && window.history.replaceState) {
            const hash = window.location.hash // "#title?password=fish&targetting=true"
            const [anchor, queryString] = hash.split("?") // ["#title", "password=fish&targetting=true"]

            const params = new URLSearchParams(queryString || "")
            params.delete("password") // Removes "password" if it exists

            const newQueryString = params.toString()
            const newHash = newQueryString
                ? `${anchor}?${newQueryString}`
                : anchor

            window.history.replaceState(null, "", newHash)
        }

        // Decrypt image copyright metadata for E2EE documents.
        // The server stores it as an opaque encrypted string; we decrypt
        // it here so the image DB receives plaintext objects.
        let decryptedImages = doc.images
        if (doc.images) {
            decryptedImages = {}
            const {E2EEEncryptor} = await import("fwtoolkit/e2ee/encryptor")
            for (const [id, image] of Object.entries(doc.images)) {
                ;(decryptedImages as Record<string, any>)[id] = image
                if (typeof image.copyright === "string") {
                    try {
                        ;(decryptedImages as Record<string, any>)[
                            id
                        ].copyright = await E2EEEncryptor.decryptObject(
                            image.copyright,
                            key
                        )
                    } catch (_e) {
                        // If decryption fails, leave as-is (legacy plaintext)
                    }
                }
            }
        }

        // Now load the decrypted document using the standard path
        const decryptedDoc = {
            ...doc,
            content: decryptedContent,
            comments: decryptedComments,
            bibliography: decryptedBibliography,
            images: decryptedImages
        }
        this._loadUnencryptedDocument(decryptedDoc as IncomingDoc, isInitialLoad)

        // Cache the decrypted title in sessionStorage so the document overview
        // can display the real title without prompting for the password again.
        let title = ""
        this.mod.editor.view.state.doc.firstChild?.forEach(child => {
            if (!child.marks.find(mark => mark.type.name === "deletion")) {
                title += child.textContent
            }
        })
        sessionStorage.setItem(
            `e2ee_title_${this.mod.editor.docInfo.id as number}`,
            title
        )

        // For newly created E2EE documents, the initial content is still
        // plaintext (it came from the template). We need to send an initial
        // encrypted snapshot so the server stores encrypted content.
        // We detect a new document by checking if the content was plaintext
        // (not a Base64 string) and this is the initial load.
        const isNewE2EEDocument =
            isInitialLoad && typeof doc.content !== "string"
        if (isNewE2EEDocument && e2ee.snapshotManager) {
            ;(e2ee.snapshotManager as any).sendInitialSnapshot(
                decryptedContent,
                decryptedComments,
                decryptedBibliography,
                this.mod.editor.docInfo.version
            )
        }

        // Now that the key is available and the document is loaded, open the
        // WebSocket connection. This ensures encrypted catch-up diffs (ep)
        // from the server can be decrypted as soon as they arrive.
        if (this.mod.editor.ws && !this.mod.editor.ws.connected) {
            this.mod.editor.startWebSocket()
        } else if (!this.mod.editor.ws) {
            // Non-collaborative mode: mark version as confirmed so the editor
            // finishes loading and periodic saves can begin.
            ;(this.mod.editor.mod.collab as any).doc.confirmVersion(
                this.mod.editor.docInfo.version
            )
        }
    }

    /**
     * Load an unencrypted (or already decrypted) document into ProseMirror.
     * This is the original loadDocument logic, extracted for reuse.
     * @private
     */
    _loadUnencryptedDocument(doc: IncomingDoc, isInitialLoad: boolean): void {
        // Remember location hash to scroll there subsequently.
        const [locationHash, _queryString] = window.location.hash.split("?")
        this.mod.editor.mod.db!.bibDB.setDB(doc.bibliography as unknown as BibDB)
        this.mod.editor.mod.db!.imageDB.setDB(doc.images as unknown as ImageDB)
        const stateDoc = this.mod.editor.schema.nodeFromJSON(
            doc.content as Record<string, unknown>
        )
        const plugins = this.mod.editor.statePlugins.map(plugin => {
            if (plugin[1]) {
                return (plugin[0] as (...args: any[]) => any)(plugin[1]())
            } else {
                return (plugin[0] as (...args: any[]) => any)()
            }
        })

        const stateConfig = {
            schema: this.mod.editor.schema,
            doc: stateDoc,
            plugins
        }

        // Set document in prosemirror
        this.mod.editor.view.setProps({
            state: EditorState.create(stateConfig)
        })
        this.mod.editor.view.setProps({nodeViews: {}}) // Needed to initialize nodeViews in plugins
        // Set initial confirmed doc
        this.mod.editor.docInfo.confirmedDoc = this.mod.editor.view.state.doc

        // Render footnotes based on main doc
        ;(this.mod.editor.mod.footnotes as any).fnEditor.renderAllFootnotes()

        //  Setup comment handling
        ;(this.mod.editor.mod.comments as any).store.reset()
        ;(this.mod.editor.mod.comments as any).store.loadComments(doc.comments)
        ;(this.mod.editor.mod.marginboxes as any).view(this.mod.editor.view)
        if (locationHash && locationHash.length) {
            this.mod.editor.scrollIdIntoView(locationHash.slice(1))
        }
        // Update the header bar to reflect the loaded document's title.
        // This is needed because setStyles() (which calls headerView.update())
        // runs before loadDocument(), so the header still shows "Untitled"
        // from the empty initial document.
        if (this.mod.editor.menu.headerView) {
            ;(
                this.mod.editor.menu.headerView as {update: () => void}
            ).update()
        }
        if (isInitialLoad) {
            this.initialDocLoaded = true
        } else {
            deactivateWait()
            this.mod.editor.waitingForDocument = false
        }
        if (doc.template) {
            // We received the template. That means we are the first user present with write access.
            // We will adjust the document to the template if necessary.
            // For E2EE documents, the content has already been decrypted
            // by this point, so template adjustment works normally.
            if (isInitialLoad) {
                this.templateAdjustmentPending = true
            }
            activateWait(true, gettext("Updating document. Please wait..."))
            const activateWaitTimer = setTimeout(() => {
                activateWait(
                    true,
                    gettext(
                        "It's taking a bit longer than usual, but it should be ready soon. Please wait..."
                    )
                )
            }, 60000)
            const adjustWorker = makeWorker(
                staticUrl("js/adjust_doc_to_template_worker.js")
            )
            adjustWorker.onmessage = (message: MessageEvent) => {
                if (message.data.type === "result") {
                    if (message.data.steps.length) {
                        const tr = this.mod.editor.view.state.tr
                        message.data.steps.forEach((step: unknown) =>
                            tr.step(Step.fromJSON(this.mod.editor.schema, step))
                        )
                        tr.setMeta("remote", true)
                        this.mod.editor.view.dispatch(tr)
                    }
                    // clearing timer for updating message since operation is completed
                    clearTimeout(activateWaitTimer)
                    if (isInitialLoad) {
                        this.templateAdjustmentPending = false
                        this.finishInitialLoad()
                    } else {
                        deactivateWait()
                    }
                    this.setDocSettings()
                }
            }
            adjustWorker.onerror = (error: ErrorEvent) => console.error(error)
            const schemaExporter = new SchemaExport()
            adjustWorker.postMessage({
                schemaSpec: JSON.parse(schemaExporter.init()),
                doc: doc.content,
                template: doc.template.content,
                documentStyleSlugs: (
                    this.mod.editor.mod.documentTemplate as any
                ).documentStyles.map((style: {slug: string}) => style.slug)
            })
        } else {
            this.setDocSettings()
            if (isInitialLoad) {
                this.finishInitialLoad()
            }
        }
    }

    setDocSettings(): void {
        // Set part specific settings
        ;(this.mod.editor.mod.documentTemplate as any).addDocPartSettings()
        ;(this.mod.editor.mod.documentTemplate as any).addCitationStylesMenuEntries()
    }

    sendToCollaborators(): void {
        // Guard: WebSocket may not be initialized yet (e.g. during
        // loadDocument before the WS connector is created).
        if (!this.mod.editor.ws) {
            return
        }

        // For E2EE documents, we need to encrypt diffs before sending.
        // Since encryption is async, we use a separate code path.
        const isE2EE =
            this.mod.editor.e2ee &&
            this.mod.editor.e2ee.encrypted &&
            this.mod.editor.e2ee.key

        if (isE2EE) {
            this._sendE2EEDiff()
            return
        }

        // Handle either doc change and comment updates OR caret update. Priority
        // for doc change/comment update.
        this.mod.editor.ws.send(() => {
            if (
                this.awaitingDiffResponse ||
                this.mod.editor.waitingForDocument ||
                this.receiving
            ) {
                return false
            } else if (
                sendableSteps(this.mod.editor.view.state) ||
                (this.mod.editor.mod.comments as any).store
                    .unsentEvents()
                    .length ||
                (this.mod.editor.mod.db!.bibDB as any).unsentEvents().length ||
                (this.mod.editor.mod.db!.imageDB as any).unsentEvents().length
            ) {
                this.disableDiffSending()
                const stepsToSend = sendableSteps(this.mod.editor.view.state),
                    fnStepsToSend = sendableSteps(
                        (this.mod.editor.mod.footnotes as any).fnEditor.view
                            .state
                    ),
                    commentUpdates = (this.mod.editor.mod.comments as any).store
                        .unsentEvents(),
                    bibliographyUpdates = (
                        this.mod.editor.mod.db!.bibDB as any
                    ).unsentEvents(),
                    imageUpdates = (
                        this.mod.editor.mod.db!.imageDB as any
                    ).unsentEvents()

                if (
                    !stepsToSend &&
                    !fnStepsToSend &&
                    !commentUpdates.length &&
                    !bibliographyUpdates.length &&
                    !imageUpdates.length
                ) {
                    // no diff. abandon operation
                    return
                }
                const rid = this.confirmStepsRequestCounter++,
                    unconfirmedDiff: Record<string, any> = {
                        type: "diff",
                        v: this.mod.editor.docInfo.version,
                        rid
                    }

                unconfirmedDiff["cid"] = this.mod.editor.client_id

                if (stepsToSend) {
                    unconfirmedDiff["ds"] = stepsToSend.steps.map(s =>
                        s.toJSON()
                    )
                    // In case the title changed, we also add a title field to
                    // update the title field instantly - important for the
                    // document overview page.
                    let newTitle = ""
                    this.mod.editor.view.state.doc.firstChild?.forEach(
                        child => {
                            if (
                                !child.marks.find(
                                    mark => mark.type.name === "deletion"
                                )
                            ) {
                                newTitle += child.textContent
                            }
                        }
                    )
                    newTitle = newTitle.slice(0, 255)
                    let oldTitle = ""
                    ;(
                        this.mod.editor.docInfo.confirmedDoc as Node
                    ).firstChild?.forEach(child => {
                        if (
                            !child.marks.find(
                                mark => mark.type.name === "deletion"
                            )
                        ) {
                            oldTitle += child.textContent
                        }
                    })
                    oldTitle = oldTitle.slice(0, 255)
                    if (newTitle !== oldTitle) {
                        unconfirmedDiff["ti"] = newTitle
                    }
                }

                if (fnStepsToSend) {
                    // We add the client ID to every single step
                    unconfirmedDiff["fs"] = fnStepsToSend.steps.map(s =>
                        s.toJSON()
                    )
                }
                if (this.footnoteRender) {
                    unconfirmedDiff["footnoterender"] = true
                    this.footnoteRender = false
                }
                if (commentUpdates.length) {
                    unconfirmedDiff["cu"] = commentUpdates
                }
                if (bibliographyUpdates.length) {
                    unconfirmedDiff["bu"] = bibliographyUpdates
                }
                if (imageUpdates.length) {
                    unconfirmedDiff["iu"] = imageUpdates
                }

                this.unconfirmedDiffs[rid] = Object.assign(
                    {doc: this.mod.editor.view.state.doc},
                    unconfirmedDiff
                )
                return unconfirmedDiff
            } else if (
                this.mod.editor.currentView?.state &&
                getSelectionUpdate(this.mod.editor.currentView.state)
            ) {
                const currentView = this.mod.editor.currentView

                if (this.lastSelectionUpdateState === currentView.state) {
                    // Selection update has been sent for this state already. Skip
                    return false
                }
                this.lastSelectionUpdateState = currentView.state
                // Create a new caret as the current user
                const selectionUpdate = getSelectionUpdate(
                    currentView.state
                ) as {anchor: number; head: number}
                return {
                    type: "selection_change",
                    id: this.mod.editor.user.id,
                    v: this.mod.editor.docInfo.version,
                    session_id: this.mod.editor.docInfo.session_id,
                    anchor: selectionUpdate.anchor,
                    head: selectionUpdate.head,
                    // Whether the selection is in the footnote or the main editor
                    editor:
                        currentView === this.mod.editor.view
                            ? "main"
                            : "footnotes"
                }
            } else {
                return false
            }
        })
    }

    /**
     * Send an encrypted diff for an E2EE document.
     *
     * For E2EE documents, the diff payload (steps, comments, bibliography,
     * image updates) is encrypted as a single blob before sending. The
     * server relays the encrypted diff to other clients without being able
     * to read the content.
     *
     * The unencrypted diff is stored locally in `unconfirmedDiffs` for
     * confirmation/rejection handling. Only the wire format is encrypted.
     *
     * @private
     */
    async _sendE2EEDiff(): Promise<void> {
        // Check if there's anything to send
        if (
            this.awaitingDiffResponse ||
            this.mod.editor.waitingForDocument ||
            this.receiving
        ) {
            // No diff to send, but check for selection update.
            // Guard: the editor view may not be initialized yet
            // (e.g. during loadDocument before ProseMirror is set up).
            if (this.mod.editor.view) {
                this._sendE2EESelectionChange()
            }
            return
        }

        if (
            !sendableSteps(this.mod.editor.view.state) &&
            !(this.mod.editor.mod.comments as any).store.unsentEvents().length &&
            !(this.mod.editor.mod.db!.bibDB as any).unsentEvents().length &&
            !(this.mod.editor.mod.db!.imageDB as any).unsentEvents().length
        ) {
            // No diff, check for selection update
            if (this.mod.editor.view) {
                this._sendE2EESelectionChange()
            }
            return
        }

        this.disableDiffSending()

        const stepsToSend = sendableSteps(this.mod.editor.view.state),
            fnStepsToSend = sendableSteps(
                (this.mod.editor.mod.footnotes as any).fnEditor.view.state
            ),
            commentUpdates = (this.mod.editor.mod.comments as any).store
                .unsentEvents(),
            bibliographyUpdates = (this.mod.editor.mod.db!.bibDB as any)
                .unsentEvents(),
            imageUpdates = (this.mod.editor.mod.db!.imageDB as any)
                .unsentEvents()

        if (
            !stepsToSend &&
            !fnStepsToSend &&
            !commentUpdates.length &&
            !bibliographyUpdates.length &&
            !imageUpdates.length
        ) {
            // no diff. abandon operation
            this.enableDiffSending()
            return
        }

        const rid = this.confirmStepsRequestCounter++,
            unconfirmedDiff: Record<string, any> = {
                type: "diff",
                v: this.mod.editor.docInfo.version,
                rid
            }

        unconfirmedDiff["cid"] = this.mod.editor.client_id

        // Collect the payload fields that need encryption
        const encryptedPayload: Record<string, any> = {}

        if (stepsToSend) {
            encryptedPayload["ds"] = stepsToSend.steps.map(s => s.toJSON())
            // Track title changes for the document overview
            let newTitle = ""
            this.mod.editor.view.state.doc.firstChild?.forEach(child => {
                if (!child.marks.find(mark => mark.type.name === "deletion")) {
                    newTitle += child.textContent
                }
            })
            newTitle = newTitle.slice(0, 255)
            let oldTitle = ""
            ;(this.mod.editor.docInfo.confirmedDoc as Node).firstChild?.forEach(
                child => {
                    if (
                        !child.marks.find(
                            mark => mark.type.name === "deletion"
                        )
                    ) {
                        oldTitle += child.textContent
                    }
                }
            )
            oldTitle = oldTitle.slice(0, 255)
            if (newTitle !== oldTitle) {
                // For E2EE documents, the title is encrypted too.
                // We don't send "ti" in plaintext; it goes in the
                // encrypted payload.
                encryptedPayload["ti"] = newTitle
            }
        }

        if (fnStepsToSend) {
            encryptedPayload["fs"] = fnStepsToSend.steps.map(s => s.toJSON())
        }
        if (this.footnoteRender) {
            unconfirmedDiff["footnoterender"] = true
            this.footnoteRender = false
        }
        if (commentUpdates.length) {
            encryptedPayload["cu"] = commentUpdates
        }
        if (bibliographyUpdates.length) {
            encryptedPayload["bu"] = bibliographyUpdates
        }
        if (imageUpdates.length) {
            encryptedPayload["iu"] = imageUpdates
        }

        // Store the unencrypted diff locally for confirmation handling.
        // We keep the plaintext steps so confirmDiff can apply them.
        this.unconfirmedDiffs[rid] = Object.assign(
            {doc: this.mod.editor.view.state.doc},
            unconfirmedDiff,
            encryptedPayload
        )

        try {
            // Encrypt the payload fields as a single blob
            const key = this.mod.editor.e2ee!.key as CryptoKey
            const ep = await E2EEEncryptor.encryptObject(encryptedPayload, key)

            // Build the wire-format diff: metadata in plaintext,
            // payload encrypted
            const wireDiff: Record<string, any> = {
                type: "diff",
                v: unconfirmedDiff.v,
                rid: unconfirmedDiff.rid,
                cid: unconfirmedDiff.cid,
                ep, // encrypted payload
                e2ee_salt: this.mod.editor.e2ee!.encryptionSalt
            }
            if (unconfirmedDiff.footnoterender) {
                wireDiff.footnoterender = true
            }

            // Send the encrypted diff
            this.mod.editor.ws!.send(() => wireDiff)
        } catch (error) {
            console.error("E2EE: Failed to encrypt diff", error)
            // Re-enable diff sending so we can try again
            this.enableDiffSending()
        }
    }

    /**
     * Send a selection change for an E2EE document.
     * Selection changes are not encrypted (they only contain cursor position).
     *
     * @private
     */
    _sendE2EESelectionChange(): void {
        if (
            !this.mod.editor.currentView ||
            !this.mod.editor.currentView.state
        ) {
            return
        }
        if (!getSelectionUpdate(this.mod.editor.currentView.state)) {
            return
        }
        const currentView = this.mod.editor.currentView
        if (this.lastSelectionUpdateState === currentView.state) {
            return
        }
        this.lastSelectionUpdateState = currentView.state
        const selectionUpdate = getSelectionUpdate(currentView.state) as {
            anchor: number
            head: number
        }
        this.mod.editor.ws!.send(() => ({
            type: "selection_change",
            id: this.mod.editor.user.id,
            v: this.mod.editor.docInfo.version,
            session_id: this.mod.editor.docInfo.session_id,
            anchor: selectionUpdate.anchor,
            head: selectionUpdate.head,
            editor: currentView === this.mod.editor.view ? "main" : "footnotes"
        }))
    }

    receiveSelectionChange(data: {
        id: number
        editor: string
        session_id: string
        anchor: number
        head: number
    }): void {
        const participant = this.mod.participants.find(
            par => par.id === data.id
        )
        let tr: Transaction | false, fnTr: Transaction | false
        if (!participant) {
            // participant is still unknown to us. Ignore
            return
        }
        if (data.editor === "footnotes") {
            fnTr = updateCollaboratorSelection(
                (this.mod.editor.mod.footnotes as any).fnEditor.view.state,
                participant as {id: number; name: string},
                data
            )
            tr = removeCollaboratorSelection(this.mod.editor.view.state, data)
        } else {
            tr = updateCollaboratorSelection(
                this.mod.editor.view.state,
                participant as {id: number; name: string},
                data
            )
            fnTr = removeCollaboratorSelection(
                (this.mod.editor.mod.footnotes as any).fnEditor.view.state,
                data
            )
        }
        if (tr) {
            this.mod.editor.view.dispatch(tr)
        }
        if (fnTr) {
            ;(this.mod.editor.mod.footnotes as any).fnEditor.view.dispatch(fnTr)
        }
    }

    receiveDiff(data: DiffData, serverFix = false): void {
        // Check if this is an encrypted diff for an E2EE document.
        // Encrypted diffs have an "ep" (encrypted payload) field instead
        // of plaintext ds/fs/cu/bu/iu fields. We decrypt the payload
        // first, then process the diff normally.
        if (data["ep"] && this.mod.editor.e2ee && this.mod.editor.e2ee.key) {
            this._receiveE2EEDiff(data, serverFix)
            return
        }
        this._processDiff(data, serverFix)
    }

    /**
     * Process a diff (unencrypted or already decrypted).
     * @private
     */
    _processDiff(data: DiffData, serverFix = false): void {
        this.mod.editor.docInfo.version!++
        if (data["bu"]) {
            // bibliography updates
            (this.mod.editor.mod.db!.bibDB as any).receive(data["bu"])
        }
        if (data["iu"]) {
            // images updates
            (this.mod.editor.mod.db!.imageDB as any).receive(data["iu"])
        }
        if (data["cu"]) {
            // comment updates
            (this.mod.editor.mod.comments as any).store.receive(data["cu"])
        }
        if (data["ds"]) {
            // document steps
            this.applyDiffs(data["ds"], data["cid"])
        }
        if (data["fs"]) {
            // footnote steps
            ;(this.mod.editor.mod.footnotes as any).fnEditor.applyDiffs(
                data["fs"],
                data["cid"]
            )
        }
        if (data["footnoterender"]) {
            // re-render footnotes properly
            ;(this.mod.editor.mod.footnotes as any).fnEditor.renderAllFootnotes()
        }

        if (serverFix) {
            // Diff is a fix created by server due to missing diffs.
            if ("reject_request_id" in data) {
                delete this.unconfirmedDiffs[data.reject_request_id as number]
            }
            this.cancelCurrentlyCheckingVersion()

            // There may be unsent local changes. Send them now after .5 seconds,
            // in case collaborators want to send something first.
            this.enableDiffSending()
            window.setTimeout(() => this.sendToCollaborators(), 500)
        }
    }

    /**
     * Receive and decrypt an E2EE diff from another client.
     *
     * For E2EE documents, diffs arrive with an "ep" (encrypted payload)
     * field containing the encrypted steps, comments, bibliography, and
     * image updates. This method decrypts the payload and then processes
     * the diff using the standard _processDiff path.
     *
     * @param data - The diff message with "ep" field
     * @param serverFix - Whether this is a server-generated fix
     * @private
     */
    async _receiveE2EEDiff(
        data: DiffData,
        serverFix = false
    ): Promise<void> {
        try {
            const key = this.mod.editor.e2ee!.key as CryptoKey
            const decryptedPayload = await E2EEEncryptor.decryptObject(
                data["ep"] as string,
                key
            )

            // Merge the decrypted payload with the metadata fields
            // that were sent in plaintext (cid, rid, footnoterender)
            const mergedData: DiffData = {
                cid: data["cid"],
                rid: data["rid"],
                footnoterender: data["footnoterender"],
                reject_request_id: data["reject_request_id"],
                ...(decryptedPayload as Record<string, unknown>)
            }

            // Process the decrypted diff normally
            this._processDiff(mergedData, serverFix)
        } catch (error) {
            console.error("E2EE: Failed to decrypt incoming diff", error)
            // If decryption fails, we still need to increment the version
            // to stay in sync with the server, but we skip applying the
            // diff content.
            this.mod.editor.docInfo.version!++
        }
    }

    setConfirmedDoc(tr: Transaction, stepsLength: number): void {
        // Find the latest version of the doc without any unconfirmed local changes

        const rebased = tr.getMeta("rebased") as number,
            docNumber = rebased + stepsLength

        this.mod.editor.docInfo.confirmedDoc =
            docNumber === tr.docs.length ? tr.doc : tr.docs[docNumber]
    }

    confirmDiff(request_id: number): void {
        const unconfirmedDiffs = this.unconfirmedDiffs[request_id]
        if (!unconfirmedDiffs) {
            return
        }
        this.mod.editor.docInfo.version!++

        const sentSteps = unconfirmedDiffs["ds"] // document steps
        if (sentSteps) {
            const ourIds = sentSteps.map((_step: any) => this.mod.editor.client_id)
            const tr = receiveTransaction(
                this.mod.editor.view.state,
                sentSteps,
                ourIds
            )
            this.mod.editor.view.dispatch(tr)
            this.mod.editor.docInfo.confirmedDoc = unconfirmedDiffs["doc"]
        }

        const sentFnSteps = unconfirmedDiffs["fs"] // footnote steps
        if (sentFnSteps) {
            const fnTr = receiveTransaction(
                (this.mod.editor.mod.footnotes as any).fnEditor.view.state,
                sentFnSteps,
                sentFnSteps.map((_step: any) => this.mod.editor.client_id)
            )
            ;(this.mod.editor.mod.footnotes as any).fnEditor.view.dispatch(fnTr)
        }

        const sentComments = unconfirmedDiffs["cu"] // comment updates
        if (sentComments) {
            ;(this.mod.editor.mod.comments as any).store.eventsSent(
                sentComments
            )
        }

        const sentBibliographyUpdates = unconfirmedDiffs["bu"] // bibliography updates
        if (sentBibliographyUpdates) {
            ;(this.mod.editor.mod.db!.bibDB as any).eventsSent(
                sentBibliographyUpdates
            )
        }

        const sentImageUpdates = unconfirmedDiffs["iu"] // image updates
        if (sentImageUpdates) {
            ;(this.mod.editor.mod.db!.imageDB as any).eventsSent(
                sentImageUpdates
            )
        }

        delete this.unconfirmedDiffs[request_id]
        this.enableDiffSending()
    }

    rejectDiff(request_id: number): void {
        delete this.unconfirmedDiffs[request_id]
        this.enableDiffSending()
    }

    applyDiffs(diffs: unknown[], cid: string | number | undefined): void {
        this.receiving = true
        const steps = diffs.map(j => Step.fromJSON(this.mod.editor.schema, j))
        const clientIds = diffs
            .map(_ => cid)
            .filter((id): id is string | number => id !== undefined)
        const tr = receiveTransaction(
            this.mod.editor.view.state,
            steps,
            clientIds
        )
        tr.setMeta("remote", true)
        this.mod.editor.view.dispatch(tr)
        this.setConfirmedDoc(tr, steps.length)
        this.receiving = false
        this.sendToCollaborators()
    }
}