backend/src/routes/publish.ts

import { Router } from "express";
import { writeFileSync, existsSync } from "fs";
import { randomUUID } from "crypto";
import { EventEmitter } from "events";
import * as Y from "yjs";
import type { Hocuspocus } from "@hocuspocus/server";
import { resolveUser, extractToken } from "../auth.js";
import { isHfStorageEnabled, setUserToken } from "../hf-storage.js";
import {
  publishDocument,
  previewDocument,
  type PublishResult,
  type PublishStage,
} from "../publisher/index.js";
import { docPath } from "../utils.js";

const DEFAULT_DOC_NAME = "default";

export interface PublishContext {
  oauthEnabled: boolean;
  hocuspocus: Hocuspocus;
}

/** Lock descriptor kept while a publish is running for a given docName. */
interface PublishLock {
  jobId: string;
  startedAt: number;
  userName?: string;
}

interface PublishJob {
  emitter: EventEmitter;
  /** Latest stage so a late SSE subscriber can immediately show the current label. */
  stage?: { stage: string; detail?: Record<string, unknown> };
  finished?: { ok: true; result: PublishResult } | { ok: false; error: string };
}

/** Per-docName mutex: a doc can only be published by one request at a time. */
const publishLocks = new Map<string, PublishLock>();

/** Active/finished jobs indexed by jobId so `/api/publish/stream` can attach. */
const publishJobs = new Map<string, PublishJob>();

/** How long we keep a finished job in memory so a late subscriber can still see the final event. */
const JOB_RETENTION_MS = 30_000;

/**
 * Broadcast publish status into the Y.Doc via a dedicated Y.Map("publish-status").
 * Every collaborative editor observes this map and disables its Publish button
 * while a publish is active. Called at the start and end of a publish.
 */
async function broadcastPublishStatus(
  hocuspocus: Hocuspocus,
  docName: string,
  status: { active: boolean; userName?: string; startedAt?: number; jobId?: string }
): Promise<void> {
  try {
    const conn = await hocuspocus.openDirectConnection(docName);
    if (conn.document) {
      const statusMap = conn.document.getMap("publish-status");
      conn.document.transact(() => {
        statusMap.set("active", status.active);
        statusMap.set("userName", status.userName ?? null);
        statusMap.set("startedAt", status.startedAt ?? null);
        statusMap.set("jobId", status.jobId ?? null);
      });
    }
    await conn.disconnect();
  } catch (err) {
    console.warn("[publish] failed to broadcast publish status:", (err as Error).message);
  }
}

export function createPublishRouter(ctx: PublishContext): Router {
  const router = Router();

  // ---------- Preview ----------

  router.get("/api/preview/:docName", async (req, res) => {
    const docName = req.params.docName || DEFAULT_DOC_NAME;

    if (ctx.oauthEnabled) {
      const token = extractToken(req.headers.cookie);
      const user = await resolveUser(token);
      if (!user || !user.canEdit) {
        res.status(403).json({ error: "Unauthorized" });
        return;
      }
    }

    try {
      const conn = await ctx.hocuspocus.openDirectConnection(docName);
      if (conn.document) {
        const update = Y.encodeStateAsUpdate(conn.document);
        writeFileSync(docPath(docName), Buffer.from(update));
      }
      await conn.disconnect();

      const result = await previewDocument(docName);
      if ("error" in result) {
        res.status(404).json({ error: result.error });
        return;
      }
      res.setHeader("Content-Type", "text/html; charset=utf-8");
      res.send(result.html);
    } catch (err: any) {
      console.error("[preview] error:", err);
      res.status(500).json({ error: err.message || "Preview failed" });
    }
  });

  // ---------- Publish ----------

  router.post("/api/publish", async (req, res) => {
    let userToken: string | undefined;
    let userName: string | undefined;

    if (ctx.oauthEnabled) {
      const token = extractToken(req.headers.cookie);
      const user = await resolveUser(token);
      if (!user) {
        console.warn("[publish] no valid user from token, cookie present:", !!token);
        res.status(403).json({ error: "Unauthorized: please log in first" });
        return;
      }
      if (!user.canEdit) {
        console.warn("[publish] user lacks write access:", user.name);
        res.status(403).json({ error: "Unauthorized: write access required" });
        return;
      }
      userToken = token ?? undefined;
      userName = user.name;
      if (userToken) setUserToken(userToken);
    }

    const docName = DEFAULT_DOC_NAME;

    // Mutex: reject with 409 if another publish is already in progress.
    const existing = publishLocks.get(docName);
    if (existing) {
      res.status(409).json({
        error: "Publish already in progress",
        jobId: existing.jobId,
        startedAt: existing.startedAt,
        userName: existing.userName ?? null,
      });
      return;
    }

    const jobId = randomUUID();
    const startedAt = Date.now();
    const emitter = new EventEmitter();
    emitter.setMaxListeners(32);
    const job: PublishJob = { emitter };
    publishJobs.set(jobId, job);
    publishLocks.set(docName, { jobId, startedAt, userName });

    // Broadcast "publishing" status to all collaborators immediately so their
    // Publish buttons become disabled.
    broadcastPublishStatus(ctx.hocuspocus, docName, {
      active: true,
      userName,
      startedAt,
      jobId,
    }).catch(() => {});

    // Return the jobId right away so the client can open the SSE stream.
    res.json({ jobId });

    // Run the publish pipeline asynchronously; SSE subscribers receive events.
    try {
      // Snapshot the Y.Doc state to disk (as before) so publishDocument can
      // reload it from a stable .yjs file.
      try {
        const conn = await ctx.hocuspocus.openDirectConnection(docName);
        if (conn.document) {
          const update = Y.encodeStateAsUpdate(conn.document);
          writeFileSync(docPath(docName), Buffer.from(update));
        }
        await conn.disconnect();
      } catch (err) {
        console.warn("[publish] snapshot failed:", (err as Error).message);
      }

      const emitStage = (stage: PublishStage | string, detail?: Record<string, unknown>) => {
        const event = { stage, detail };
        job.stage = event;
        emitter.emit("stage", event);
      };

      const result = await publishDocument(docName, userToken, { onStage: emitStage });

      if (!result.success) {
        job.finished = { ok: false, error: result.error || "Publish failed" };
        emitter.emit("done", { success: false, error: result.error || "Publish failed" });
      } else {
        job.finished = { ok: true, result };
        emitter.emit("done", { success: true, result });
      }
    } catch (err: any) {
      console.error("[publish] error:", err);
      const message = err?.message || "Publish failed";
      job.finished = { ok: false, error: message };
      emitter.emit("done", { success: false, error: message });
    } finally {
      publishLocks.delete(docName);
      broadcastPublishStatus(ctx.hocuspocus, docName, { active: false }).catch(() => {});

      // Keep the finished job around briefly so late SSE subscribers still receive the outcome.
      setTimeout(() => {
        publishJobs.delete(jobId);
      }, JOB_RETENTION_MS);
    }
  });

  // ---------- SSE stream ----------

  router.get("/api/publish/stream", async (req, res) => {
    const jobId = typeof req.query.jobId === "string" ? req.query.jobId : "";
    const job = publishJobs.get(jobId);

    if (ctx.oauthEnabled) {
      const token = extractToken(req.headers.cookie);
      const user = await resolveUser(token);
      if (!user || !user.canEdit) {
        res.status(403).end();
        return;
      }
    }

    if (!job) {
      res.status(404).end();
      return;
    }

    res.setHeader("Content-Type", "text/event-stream");
    res.setHeader("Cache-Control", "no-cache, no-transform");
    res.setHeader("Connection", "keep-alive");
    res.setHeader("X-Accel-Buffering", "no");
    res.flushHeaders?.();

    const write = (event: string, payload: unknown) => {
      res.write(`event: ${event}\n`);
      res.write(`data: ${JSON.stringify(payload)}\n\n`);
    };

    // Immediately replay the current stage (if any) so the subscriber sees the latest label.
    if (job.stage) write("stage", job.stage);

    const onStage = (payload: { stage: string; detail?: Record<string, unknown> }) => {
      write("stage", payload);
    };
    const onDone = (payload: { success: boolean; error?: string; result?: PublishResult }) => {
      write("done", payload);
      cleanup();
      res.end();
    };

    // If the job already finished (rare: client subscribed after completion),
    // emit the final event and close.
    if (job.finished) {
      if (job.finished.ok) {
        write("done", { success: true, result: job.finished.result });
      } else {
        write("done", { success: false, error: job.finished.error });
      }
      res.end();
      return;
    }

    job.emitter.on("stage", onStage);
    job.emitter.on("done", onDone);

    // Keep-alive comments every 15s to defeat proxy idle timeouts.
    const keepAlive = setInterval(() => {
      res.write(`: keepalive ${Date.now()}\n\n`);
    }, 15_000);

    const cleanup = () => {
      clearInterval(keepAlive);
      job.emitter.off("stage", onStage);
      job.emitter.off("done", onDone);
    };

    req.on("close", () => {
      cleanup();
      res.end();
    });
  });

  // ---------- Status (used by non-collaborative fallback) ----------

  router.get("/api/publish/status", (_req, res) => {
    const lock = publishLocks.get(DEFAULT_DOC_NAME);
    if (!lock) {
      res.json({ active: false });
      return;
    }
    res.json({
      active: true,
      jobId: lock.jobId,
      startedAt: lock.startedAt,
      userName: lock.userName ?? null,
    });
  });

  // ---------- Reset Document ----------

  router.post("/api/admin/reset-document", async (req, res) => {
    if (ctx.oauthEnabled) {
      const token = extractToken(req.headers.cookie);
      const user = await resolveUser(token);
      if (!user || !user.canEdit) {
        res.status(403).json({ error: "Unauthorized" });
        return;
      }
      if (token) setUserToken(token);
    }

    try {
      const p = docPath(DEFAULT_DOC_NAME);
      if (existsSync(p)) {
        const { unlinkSync } = await import("fs");
        unlinkSync(p);
        console.log("[admin] deleted local Y.Doc file:", p);
      }

      await ctx.hocuspocus.closeConnections(DEFAULT_DOC_NAME);
      console.log("[admin] closed connections for", DEFAULT_DOC_NAME);

      res.json({ success: true, message: "Document reset. Refresh the editor to start fresh." });
    } catch (err: any) {
      console.error("[admin] reset error:", err);
      res.status(500).json({ error: err.message || "Reset failed" });
    }
  });

  return router;
}