/**
 * Background render worker. Runs as a standalone Node.js process (dist/worker.js), detached from
 * LM Studio's process tree, so a render started with `runInBackground: true` survives the plugin
 * host being shut down after the tool call ends.
 *
 * Lifetime: the worker is started on demand by launcher.ts, renders the queued jobs one after
 * another and exits once everything is delivered and a short grace period has passed.
 */
import fs from "fs";
import path from "path";
import {
  CLAIMED_INFIX,
  HEARTBEAT_INTERVAL_MS,
  JOBS_DIR,
  JOB_FILE_SUFFIX,
  PID_FILE,
  WORKER_LOG_FILE,
  WORKER_STATE_DIR,
  readLiveWorkerPid,
  type WorkerJob,
} from "./workerShared.js";

const TAG = "[Worker]";
/** After the last job is delivered, keep scanning for new jobs this long before exiting. */
const IDLE_GRACE_MS = 10_000;
const POLL_INTERVAL_MS = 500;
/** Slightly above the 30 min absolute timer of the job lifecycle, which has marked the record failed by then. */
const JOB_ABANDON_AFTER_MS = 31 * 60 * 1000;
const SHUTDOWN_FORCE_EXIT_MS = 5_000;
const EMPTY_PID_FILE_GRACE_MS = 1_000;

function localTimestamp(d = new Date()): string {
  const pad = (n: number, len = 2) => String(n).padStart(len, "0");
  const offsetMinutes = -d.getTimezoneOffset();
  const sign = offsetMinutes >= 0 ? "+" : "-";
  const abs = Math.abs(offsetMinutes);
  return (
    `${d.getFullYear()}-${pad(d.getMonth() + 1)}-${pad(d.getDate())}T` +
    `${pad(d.getHours())}:${pad(d.getMinutes())}:${pad(d.getSeconds())}.${pad(d.getMilliseconds(), 3)}` +
    `${sign}${pad(Math.floor(abs / 60))}:${pad(abs % 60)}`
  );
}

function log(message: string): void {
  const line = `${localTimestamp()} ${TAG} ${message}\n`;
  try {
    fs.mkdirSync(WORKER_STATE_DIR, { recursive: true });
    fs.appendFileSync(WORKER_LOG_FILE, line, "utf8");
  } catch (e) {
    process.stderr.write(`${line}${TAG} log file write failed: ${String((e as any)?.message || e)}\n`);
  }
}

function errMessage(e: unknown): string {
  return String((e as any)?.message || e);
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function pidAlive(pid: number): boolean {
  try {
    process.kill(pid, 0);
    return true;
  } catch (e: any) {
    return e?.code === "EPERM";
  }
}

/* ------------------------------ PID file -------------------------------- */

let heartbeat: NodeJS.Timeout | null = null;

function stopHeartbeat(): void {
  if (heartbeat) clearInterval(heartbeat);
  heartbeat = null;
}

function startHeartbeat(): void {
  stopHeartbeat();
  heartbeat = setInterval(() => {
    try {
      const now = new Date();
      fs.utimesSync(PID_FILE, now, now);
    } catch (e) {
      log(`heartbeat failed: ${errMessage(e)}`);
    }
  }, HEARTBEAT_INTERVAL_MS);
  heartbeat.unref();
}

/** Registers this process as the single active worker. Returns false when another live worker exists. */
function acquirePidFile(): boolean {
  fs.mkdirSync(WORKER_STATE_DIR, { recursive: true });
  for (let attempt = 0; attempt < 2; attempt++) {
    try {
      fs.writeFileSync(PID_FILE, String(process.pid), { flag: "wx" });
      startHeartbeat();
      return true;
    } catch (e: any) {
      if (e?.code !== "EEXIST") throw e;
      if (readLiveWorkerPid() !== null) return false;
      // A PID file that another worker has just created may not contain its PID yet.
      const content = fs.readFileSync(PID_FILE, "utf8").trim();
      if (content === "" && Date.now() - fs.statSync(PID_FILE).mtimeMs < EMPTY_PID_FILE_GRACE_MS) return false;
      fs.rmSync(PID_FILE, { force: true });
    }
  }
  return false;
}

function releasePidFile(): void {
  stopHeartbeat();
  try {
    if (fs.readFileSync(PID_FILE, "utf8").trim() === String(process.pid)) fs.rmSync(PID_FILE, { force: true });
  } catch (e: any) {
    if (e?.code !== "ENOENT") log(`PID file cleanup failed: ${errMessage(e)}`);
  }
}

/* ------------------------------- Job files ------------------------------ */

function listJobFiles(): string[] {
  return fs
    .readdirSync(JOBS_DIR)
    .filter((n) => n.endsWith(JOB_FILE_SUFFIX) && !n.startsWith("."))
    .sort();
}

/** Takes ownership of the oldest job by renaming it; the rename is atomic, so exactly one worker wins. */
function claimNextJob(): string | null {
  for (const name of listJobFiles()) {
    const src = path.join(JOBS_DIR, name);
    const dst = `${src}${CLAIMED_INFIX}${process.pid}`;
    try {
      fs.renameSync(src, dst);
      return dst;
    } catch (e: any) {
      if (e?.code === "ENOENT") continue;
      throw e;
    }
  }
  return null;
}

function readJob(file: string): WorkerJob | null {
  try {
    const job = JSON.parse(fs.readFileSync(file, "utf8"));
    if (
      typeof job?.requestId === "string" &&
      typeof job?.chatId === "string" &&
      typeof job?.workingDir === "string" &&
      job.args &&
      typeof job.args === "object"
    ) {
      return job as WorkerJob;
    }
    log(`job file has an unexpected shape: ${file}`);
  } catch (e) {
    log(`job file unreadable (${file}): ${errMessage(e)}`);
  }
  return null;
}

/** Fails the requests of jobs whose worker died mid-render (their claimed file is still there). */
async function recoverOrphanedClaims(): Promise<void> {
  const core = await import("../core-bundle.mjs");
  for (const name of fs.readdirSync(JOBS_DIR)) {
    const at = name.indexOf(CLAIMED_INFIX);
    if (at < 0) continue;
    const ownerPid = parseInt(name.slice(at + CLAIMED_INFIX.length), 10);
    if (ownerPid === process.pid || (Number.isFinite(ownerPid) && pidAlive(ownerPid))) continue;
    const file = path.join(JOBS_DIR, name);
    const job = readJob(file);
    if (job) {
      await core.markBackgroundRequestFailed(
        job.workingDir,
        job.requestId,
        "The background worker stopped while this job was running."
      );
      log(`recovered orphaned job ${job.requestId} (owner pid ${ownerPid})`);
    }
    fs.rmSync(file, { force: true });
  }
}

/* ------------------------------ Job execution --------------------------- */

let currentJob: WorkerJob | null = null;

async function configureCustomConfigs(customConfigsPath: string | null): Promise<void> {
  const loader = await import("../services/customConfigsLoader.js");
  const status = await loader.checkCustomConfigs(customConfigsPath ?? "");
  if (customConfigsPath && !status.available) {
    throw new Error(`custom_configs.json is not readable at ${customConfigsPath}: ${status.error ?? "unknown error"}`);
  }
  await loader.loadCustomConfigs();
}

async function executeJob(job: WorkerJob): Promise<void> {
  const core = await import("../core-bundle.mjs");
  try {
    for (const [key, value] of Object.entries(job.env)) process.env[key] = value;
    const tools: any = await import("../core/tools.js");
    if (typeof tools.handleGenerateImage !== "function") {
      throw new Error("generate_image core handler not available.");
    }
    await core.runBackgroundJobAwaited({
      chatWd: job.workingDir,
      requestId: job.requestId,
      run: async (onProgress) => {
        await configureCustomConfigs(job.customConfigsPath);
        // The render pipeline resolves the chat from this process-wide context; jobs run one at a time.
        core.setActiveChatContext({ chatId: job.chatId, workingDir: job.workingDir, requestId: job.requestId });
        return tools.handleGenerateImage(job.args, onProgress, { requestIdOverride: job.requestId });
      },
    });
  } catch (e) {
    log(`job ${job.requestId} failed before the render started: ${errMessage(e)}`);
    await core.markBackgroundRequestFailed(job.workingDir, job.requestId, errMessage(e));
  }
}

/** Resolves when `p` settles or after `ms`, so one hung render cannot block the queue forever. */
function abandonAfter(p: Promise<void>, ms: number, requestId: string): Promise<void> {
  return new Promise((resolve) => {
    const timer = setTimeout(() => {
      log(`abandoning job ${requestId} after ${Math.round(ms / 60000)} min without a result`);
      resolve();
    }, ms);
    p.finally(() => {
      clearTimeout(timer);
      resolve();
    });
  });
}

async function runJobFile(file: string): Promise<void> {
  const job = readJob(file);
  if (!job) {
    fs.renameSync(file, `${file}.invalid`);
    return;
  }
  currentJob = job;
  log(`start ${job.requestId} (chat ${job.chatId})`);
  try {
    await abandonAfter(executeJob(job), JOB_ABANDON_AFTER_MS, job.requestId);
  } finally {
    currentJob = null;
    fs.rmSync(file, { force: true });
    log(`finished ${job.requestId}`);
  }
}

/* -------------------------------- Lifecycle ----------------------------- */

async function failCurrentJob(reason: string): Promise<void> {
  const job = currentJob;
  if (!job) return;
  try {
    const core = await import("../core-bundle.mjs");
    await core.markBackgroundRequestFailed(job.workingDir, job.requestId, reason);
  } catch (e) {
    log(`could not record failure for ${job.requestId}: ${errMessage(e)}`);
  }
}

function exitAfter(code: number, work: Promise<void>): void {
  const force = setTimeout(() => process.exit(code), SHUTDOWN_FORCE_EXIT_MS);
  force.unref();
  work.finally(() => {
    releasePidFile();
    process.exit(code);
  });
}

process.on("SIGPIPE", () => {
  /* stdio is /dev/null; nothing to do */
});
process.on("SIGHUP", () => log("received SIGHUP, ignoring"));
process.on("SIGTERM", () => {
  log("received SIGTERM, shutting down");
  exitAfter(0, failCurrentJob("The background worker was stopped while this job was running."));
});
process.on("SIGINT", () => {
  log("received SIGINT, shutting down");
  exitAfter(0, failCurrentJob("The background worker was stopped while this job was running."));
});
process.on("uncaughtException", (err) => {
  log(`UNCAUGHT EXCEPTION: ${(err as any)?.stack ?? err}`);
  exitAfter(99, failCurrentJob(`The background worker crashed: ${errMessage(err)}`));
});
process.on("unhandledRejection", (reason) => {
  log(`UNHANDLED REJECTION: ${(reason as any)?.stack ?? reason}`);
});

async function main(): Promise<void> {
  fs.mkdirSync(JOBS_DIR, { recursive: true });
  if (!acquirePidFile()) {
    log("another worker is already running, exiting");
    process.exit(0);
  }
  log(`started (pid ${process.pid})`);
  await recoverOrphanedClaims();

  let idleSince: number | null = null;
  for (;;) {
    const file = claimNextJob();
    if (file) {
      idleSince = null;
      await runJobFile(file);
      continue;
    }
    if (idleSince === null) idleSince = Date.now();
    if (Date.now() - idleSince >= IDLE_GRACE_MS) {
      // Unregister first, then scan once more: a launcher that wrote its job after this scan also
      // sees the missing PID file and starts a new worker.
      releasePidFile();
      if (listJobFiles().length > 0) {
        if (acquirePidFile()) {
          idleSince = null;
          continue;
        }
        log("another worker took over the pending jobs, exiting");
        process.exit(0);
      }
      log("all jobs delivered, exiting");
      process.exit(0);
    }
    await sleep(POLL_INTERVAL_MS);
  }
}

main().catch((e) => {
  log(`fatal: ${(e as any)?.stack ?? e}`);
  releasePidFile();
  process.exit(1);
});

/**
 * Background render worker. Runs as a standalone Node.js process (dist/worker.js), detached from
 * LM Studio's process tree, so a render started with `runInBackground: true` survives the plugin
 * host being shut down after the tool call ends.
 *
 * Lifetime: the worker is started on demand by launcher.ts, renders the queued jobs one after
 * another and exits once everything is delivered and a short grace period has passed.
 */
import fs from "fs";
import path from "path";
import {
  CLAIMED_INFIX,
  HEARTBEAT_INTERVAL_MS,
  JOBS_DIR,
  JOB_FILE_SUFFIX,
  PID_FILE,
  WORKER_LOG_FILE,
  WORKER_STATE_DIR,
  readLiveWorkerPid,
  type WorkerJob,
} from "./workerShared.js";

const TAG = "[Worker]";
/** After the last job is delivered, keep scanning for new jobs this long before exiting. */
const IDLE_GRACE_MS = 10_000;
const POLL_INTERVAL_MS = 500;
/** Slightly above the 30 min absolute timer of the job lifecycle, which has marked the record failed by then. */
const JOB_ABANDON_AFTER_MS = 31 * 60 * 1000;
const SHUTDOWN_FORCE_EXIT_MS = 5_000;
const EMPTY_PID_FILE_GRACE_MS = 1_000;

function localTimestamp(d = new Date()): string {
  const pad = (n: number, len = 2) => String(n).padStart(len, "0");
  const offsetMinutes = -d.getTimezoneOffset();
  const sign = offsetMinutes >= 0 ? "+" : "-";
  const abs = Math.abs(offsetMinutes);
  return (
    `${d.getFullYear()}-${pad(d.getMonth() + 1)}-${pad(d.getDate())}T` +
    `${pad(d.getHours())}:${pad(d.getMinutes())}:${pad(d.getSeconds())}.${pad(d.getMilliseconds(), 3)}` +
    `${sign}${pad(Math.floor(abs / 60))}:${pad(abs % 60)}`
  );
}

function log(message: string): void {
  const line = `${localTimestamp()} ${TAG} ${message}\n`;
  try {
    fs.mkdirSync(WORKER_STATE_DIR, { recursive: true });
    fs.appendFileSync(WORKER_LOG_FILE, line, "utf8");
  } catch (e) {
    process.stderr.write(`${line}${TAG} log file write failed: ${String((e as any)?.message || e)}\n`);
  }
}

function errMessage(e: unknown): string {
  return String((e as any)?.message || e);
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function pidAlive(pid: number): boolean {
  try {
    process.kill(pid, 0);
    return true;
  } catch (e: any) {
    return e?.code === "EPERM";
  }
}

/* ------------------------------ PID file -------------------------------- */

let heartbeat: NodeJS.Timeout | null = null;

function stopHeartbeat(): void {
  if (heartbeat) clearInterval(heartbeat);
  heartbeat = null;
}

function startHeartbeat(): void {
  stopHeartbeat();
  heartbeat = setInterval(() => {
    try {
      const now = new Date();
      fs.utimesSync(PID_FILE, now, now);
    } catch (e) {
      log(`heartbeat failed: ${errMessage(e)}`);
    }
  }, HEARTBEAT_INTERVAL_MS);
  heartbeat.unref();
}

/** Registers this process as the single active worker. Returns false when another live worker exists. */
function acquirePidFile(): boolean {
  fs.mkdirSync(WORKER_STATE_DIR, { recursive: true });
  for (let attempt = 0; attempt < 2; attempt++) {
    try {
      fs.writeFileSync(PID_FILE, String(process.pid), { flag: "wx" });
      startHeartbeat();
      return true;
    } catch (e: any) {
      if (e?.code !== "EEXIST") throw e;
      if (readLiveWorkerPid() !== null) return false;
      // A PID file that another worker has just created may not contain its PID yet.
      const content = fs.readFileSync(PID_FILE, "utf8").trim();
      if (content === "" && Date.now() - fs.statSync(PID_FILE).mtimeMs < EMPTY_PID_FILE_GRACE_MS) return false;
      fs.rmSync(PID_FILE, { force: true });
    }
  }
  return false;
}

function releasePidFile(): void {
  stopHeartbeat();
  try {
    if (fs.readFileSync(PID_FILE, "utf8").trim() === String(process.pid)) fs.rmSync(PID_FILE, { force: true });
  } catch (e: any) {
    if (e?.code !== "ENOENT") log(`PID file cleanup failed: ${errMessage(e)}`);
  }
}

/* ------------------------------- Job files ------------------------------ */

function listJobFiles(): string[] {
  return fs
    .readdirSync(JOBS_DIR)
    .filter((n) => n.endsWith(JOB_FILE_SUFFIX) && !n.startsWith("."))
    .sort();
}

/** Takes ownership of the oldest job by renaming it; the rename is atomic, so exactly one worker wins. */
function claimNextJob(): string | null {
  for (const name of listJobFiles()) {
    const src = path.join(JOBS_DIR, name);
    const dst = `${src}${CLAIMED_INFIX}${process.pid}`;
    try {
      fs.renameSync(src, dst);
      return dst;
    } catch (e: any) {
      if (e?.code === "ENOENT") continue;
      throw e;
    }
  }
  return null;
}

function readJob(file: string): WorkerJob | null {
  try {
    const job = JSON.parse(fs.readFileSync(file, "utf8"));
    if (
      typeof job?.requestId === "string" &&
      typeof job?.chatId === "string" &&
      typeof job?.workingDir === "string" &&
      job.args &&
      typeof job.args === "object"
    ) {
      return job as WorkerJob;
    }
    log(`job file has an unexpected shape: ${file}`);
  } catch (e) {
    log(`job file unreadable (${file}): ${errMessage(e)}`);
  }
  return null;
}

/** Fails the requests of jobs whose worker died mid-render (their claimed file is still there). */
async function recoverOrphanedClaims(): Promise<void> {
  const core = await import("../core-bundle.mjs");
  for (const name of fs.readdirSync(JOBS_DIR)) {
    const at = name.indexOf(CLAIMED_INFIX);
    if (at < 0) continue;
    const ownerPid = parseInt(name.slice(at + CLAIMED_INFIX.length), 10);
    if (ownerPid === process.pid || (Number.isFinite(ownerPid) && pidAlive(ownerPid))) continue;
    const file = path.join(JOBS_DIR, name);
    const job = readJob(file);
    if (job) {
      await core.markBackgroundRequestFailed(
        job.workingDir,
        job.requestId,
        "The background worker stopped while this job was running."
      );
      log(`recovered orphaned job ${job.requestId} (owner pid ${ownerPid})`);
    }
    fs.rmSync(file, { force: true });
  }
}

/* ------------------------------ Job execution --------------------------- */

let currentJob: WorkerJob | null = null;

async function configureCustomConfigs(customConfigsPath: string | null): Promise<void> {
  const loader = await import("../services/customConfigsLoader.js");
  const status = await loader.checkCustomConfigs(customConfigsPath ?? "");
  if (customConfigsPath && !status.available) {
    throw new Error(`custom_configs.json is not readable at ${customConfigsPath}: ${status.error ?? "unknown error"}`);
  }
  await loader.loadCustomConfigs();
}

async function executeJob(job: WorkerJob): Promise<void> {
  const core = await import("../core-bundle.mjs");
  try {
    for (const [key, value] of Object.entries(job.env)) process.env[key] = value;
    const tools: any = await import("../core/tools.js");
    if (typeof tools.handleGenerateImage !== "function") {
      throw new Error("generate_image core handler not available.");
    }
    await core.runBackgroundJobAwaited({
      chatWd: job.workingDir,
      requestId: job.requestId,
      run: async (onProgress) => {
        await configureCustomConfigs(job.customConfigsPath);
        // The render pipeline resolves the chat from this process-wide context; jobs run one at a time.
        core.setActiveChatContext({ chatId: job.chatId, workingDir: job.workingDir, requestId: job.requestId });
        return tools.handleGenerateImage(job.args, onProgress, { requestIdOverride: job.requestId });
      },
    });
  } catch (e) {
    log(`job ${job.requestId} failed before the render started: ${errMessage(e)}`);
    await core.markBackgroundRequestFailed(job.workingDir, job.requestId, errMessage(e));
  }
}

/** Resolves when `p` settles or after `ms`, so one hung render cannot block the queue forever. */
function abandonAfter(p: Promise<void>, ms: number, requestId: string): Promise<void> {
  return new Promise((resolve) => {
    const timer = setTimeout(() => {
      log(`abandoning job ${requestId} after ${Math.round(ms / 60000)} min without a result`);
      resolve();
    }, ms);
    p.finally(() => {
      clearTimeout(timer);
      resolve();
    });
  });
}

async function runJobFile(file: string): Promise<void> {
  const job = readJob(file);
  if (!job) {
    fs.renameSync(file, `${file}.invalid`);
    return;
  }
  currentJob = job;
  log(`start ${job.requestId} (chat ${job.chatId})`);
  try {
    await abandonAfter(executeJob(job), JOB_ABANDON_AFTER_MS, job.requestId);
  } finally {
    currentJob = null;
    fs.rmSync(file, { force: true });
    log(`finished ${job.requestId}`);
  }
}

/* -------------------------------- Lifecycle ----------------------------- */

async function failCurrentJob(reason: string): Promise<void> {
  const job = currentJob;
  if (!job) return;
  try {
    const core = await import("../core-bundle.mjs");
    await core.markBackgroundRequestFailed(job.workingDir, job.requestId, reason);
  } catch (e) {
    log(`could not record failure for ${job.requestId}: ${errMessage(e)}`);
  }
}

function exitAfter(code: number, work: Promise<void>): void {
  const force = setTimeout(() => process.exit(code), SHUTDOWN_FORCE_EXIT_MS);
  force.unref();
  work.finally(() => {
    releasePidFile();
    process.exit(code);
  });
}

process.on("SIGPIPE", () => {
  /* stdio is /dev/null; nothing to do */
});
process.on("SIGHUP", () => log("received SIGHUP, ignoring"));
process.on("SIGTERM", () => {
  log("received SIGTERM, shutting down");
  exitAfter(0, failCurrentJob("The background worker was stopped while this job was running."));
});
process.on("SIGINT", () => {
  log("received SIGINT, shutting down");
  exitAfter(0, failCurrentJob("The background worker was stopped while this job was running."));
});
process.on("uncaughtException", (err) => {
  log(`UNCAUGHT EXCEPTION: ${(err as any)?.stack ?? err}`);
  exitAfter(99, failCurrentJob(`The background worker crashed: ${errMessage(err)}`));
});
process.on("unhandledRejection", (reason) => {
  log(`UNHANDLED REJECTION: ${(reason as any)?.stack ?? reason}`);
});

async function main(): Promise<void> {
  fs.mkdirSync(JOBS_DIR, { recursive: true });
  if (!acquirePidFile()) {
    log("another worker is already running, exiting");
    process.exit(0);
  }
  log(`started (pid ${process.pid})`);
  await recoverOrphanedClaims();

  let idleSince: number | null = null;
  for (;;) {
    const file = claimNextJob();
    if (file) {
      idleSince = null;
      await runJobFile(file);
      continue;
    }
    if (idleSince === null) idleSince = Date.now();
    if (Date.now() - idleSince >= IDLE_GRACE_MS) {
      // Unregister first, then scan once more: a launcher that wrote its job after this scan also
      // sees the missing PID file and starts a new worker.
      releasePidFile();
      if (listJobFiles().length > 0) {
        if (acquirePidFile()) {
          idleSince = null;
          continue;
        }
        log("another worker took over the pending jobs, exiting");
        process.exit(0);
      }
      log("all jobs delivered, exiting");
      process.exit(0);
    }
    await sleep(POLL_INTERVAL_MS);
  }
}

main().catch((e) => {
  log(`fatal: ${(e as any)?.stack ?? e}`);
  releasePidFile();
  process.exit(1);
});