"use strict";

const crypto = require("node:crypto");
const { EventEmitter } = require("node:events");
const { pool } = require("./db");
const { refundTaskBillingOnFailure } = require("./billing");
const { putObject, isOssConfigured } = require("./ossClient");

const taskEvents = new EventEmitter();
taskEvents.setMaxListeners(200);

const activePollers = new Map();
const POLL_INTERVAL_MS = 3000;
const MAX_POLL_ATTEMPTS = 120;
const GRS_IMAGE_MAX_POLL_ATTEMPTS = Number(process.env.GRSAI_IMAGE_MAX_POLL_ATTEMPTS || 60);

function normalizeTaskProgress(value) {
  const numeric = Number(value);
  if (!Number.isFinite(numeric)) return undefined;
  return Math.max(0, Math.min(100, Math.round(numeric)));
}

function formatTaskEvent(row) {
  return {
    taskId: row.id,
    status: row.status,
    progress: row.progress,
    resultUrl: row.result_url || null,
    error: row.error || null,
  };
}

async function createTaskLifecycleNotification(task) {
  if (!task || !task.user_id || !task.id) return;

  const isCompleted = task.status === "completed";
  const isFailed = task.status === "failed";
  if (!isCompleted && !isFailed) return;

  const typeLabel = task.type === "video" ? "视频" : "图像";
  const noticeType = isCompleted ? "task_completed" : "task_failed";
  const title = isCompleted ? `${typeLabel}生成已完成` : `${typeLabel}生成失败`;
  const description = isCompleted
    ? "生成结果已同步到任务历史，可继续编辑或保存到资产库。"
    : String(task.error || "任务处理失败，请稍后重试。").slice(0, 500);

  await pool.query(
    `
    INSERT INTO web_notifications (
      user_id, type, title, description, target_type, target_id, metadata_json
    )
    SELECT $1::integer, $2::varchar, $3::varchar, $4::text, 'generation_task', $5::varchar, $6::text
    WHERE NOT EXISTS (
      SELECT 1
      FROM web_notifications
      WHERE user_id = $1
        AND type = $2::varchar
        AND target_type = 'generation_task'
        AND target_id = $5::varchar
    )
  `,
    [
      task.user_id,
      noticeType,
      title,
      description,
      String(task.id),
      JSON.stringify({ taskType: task.type, resultUrl: task.result_url || null }),
    ],
  );
}

async function updateTaskInDb(taskId, updates) {
  const nextUpdates = { ...updates };

  const fields = [];
  const values = [];
  let idx = 1;

  if (nextUpdates.status !== undefined) { fields.push(`status = $${idx++}`); values.push(nextUpdates.status); }
  if (nextUpdates.progress !== undefined) {
    const progress = normalizeTaskProgress(nextUpdates.progress);
    if (progress !== undefined) { fields.push(`progress = $${idx++}`); values.push(progress); }
  }
  if (nextUpdates.resultUrl !== undefined) { fields.push(`result_url = $${idx++}`); values.push(nextUpdates.resultUrl); }
  if (nextUpdates.error !== undefined) { fields.push(`error = $${idx++}`); values.push(nextUpdates.error); }
  if (nextUpdates.providerTaskId !== undefined) { fields.push(`provider_task_id = $${idx++}`); values.push(nextUpdates.providerTaskId); }
  if (nextUpdates.status === "completed" || nextUpdates.status === "failed") {
    fields.push("completed_at = NOW()");
  }
  fields.push("updated_at = NOW()");

  if (fields.length === 0) return;
  values.push(taskId);
  const { rows } = await pool.query(
    `UPDATE generation_tasks SET ${fields.join(", ")} WHERE id = $${idx} RETURNING *`,
    values,
  );
  let updatedTask = rows[0];

  if (updatedTask) {
    taskEvents.emit(`task:${taskId}`, formatTaskEvent(updatedTask));
  }

  if (nextUpdates.status === "completed" && updatedTask?.result_url) {
    persistTaskResultUrlToOssInBackground(updatedTask);
  }

  if (nextUpdates.status === "completed" || nextUpdates.status === "failed") {
    await createTaskLifecycleNotification(updatedTask).catch((err) => {
      console.error(`[aiTaskWorker] notification error for task ${taskId}:`, err.message);
    });
  }

  if (nextUpdates.status === "failed") {
    await refundTaskBillingOnFailure(taskId).catch((err) => {
      console.error(`[aiTaskWorker] refund error for task ${taskId}:`, err.message);
    });
  }
}

function persistTaskResultUrlToOssInBackground(task) {
  if (!task?.id || !task?.result_url) return;

  Promise.resolve()
    .then(async () => {
      const durableUrl = await persistResultUrlToOss(task);
      if (!durableUrl || durableUrl === task.result_url) return;

      await pool.query(
        "UPDATE generation_tasks SET result_url = $1, updated_at = NOW() WHERE id = $2 AND result_url = $3",
        [durableUrl, task.id, task.result_url],
      );
      console.info(`[aiTaskWorker] task ${task.id} result persisted to OSS after completion`);
    })
    .catch((error) => {
      console.warn(`[aiTaskWorker] background result persistence failed for task ${task.id}:`, error.message);
    });
}

function asObject(value) {
  return value && typeof value === "object" && !Array.isArray(value) ? value : undefined;
}

function asString(value) {
  return typeof value === "string" && value.trim() ? value.trim() : undefined;
}

function firstString(...values) {
  for (const value of values) {
    const stringValue = asString(value);
    if (stringValue) return stringValue;
  }
  return undefined;
}

function mediaExtensionFromContentType(contentType, fallbackUrl, taskType) {
  const mime = String(contentType || "").split(";")[0].trim().toLowerCase();
  if (mime === "image/jpeg") return "jpg";
  if (mime === "image/png") return "png";
  if (mime === "image/webp") return "webp";
  if (mime === "image/gif") return "gif";
  if (mime === "video/webm") return "webm";
  if (mime === "video/quicktime") return "mov";
  if (mime === "video/mp4") return "mp4";

  try {
    const matched = new URL(fallbackUrl).pathname.match(/\.([a-z0-9]{2,5})$/i);
    if (matched?.[1]) return matched[1].toLowerCase();
  } catch {
    // Keep the type fallback below.
  }

  return taskType === "video" ? "mp4" : "png";
}

function isErrorDocumentContentType(contentType) {
  return /(?:application|text)\/(?:json|xml|html|plain)|\+xml/i.test(String(contentType || ""));
}

function isOwnPersistedResultUrl(value) {
  return /\/users\/[^/]+\/generation-results\//i.test(String(value || ""));
}

async function persistIncomingResultUrl(taskId, resultUrl) {
  const normalizedUrl = String(resultUrl || "").trim();
  if (!/^https?:\/\//i.test(normalizedUrl) || isOwnPersistedResultUrl(normalizedUrl) || !isOssConfigured()) {
    return null;
  }

  try {
    const { rows } = await pool.query(
      "SELECT id, user_id, type, result_url FROM generation_tasks WHERE id = $1",
      [taskId],
    );
    const task = rows[0];
    if (!task?.user_id || !task?.type) return null;
    return persistResultUrlToOss({ ...task, result_url: normalizedUrl });
  } catch (error) {
    console.warn(`[aiTaskWorker] result pre-persistence skipped for task ${taskId}:`, error.message);
    return null;
  }
}

async function persistResultUrlToOss(task) {
  const resultUrl = String(task?.result_url || "").trim();
  if (!/^https?:\/\//i.test(resultUrl) || isOwnPersistedResultUrl(resultUrl) || !isOssConfigured()) {
    return null;
  }

  try {
    const response = await fetch(resultUrl, { method: "GET" });
    if (!response.ok) {
      throw new Error(`result fetch returned ${response.status}`);
    }

    const contentType = response.headers.get("content-type") || (task.type === "video" ? "video/mp4" : "image/png");
    if (isErrorDocumentContentType(contentType)) {
      const text = await response.text().catch(() => "");
      throw new Error(`result fetch returned error document: ${text.slice(0, 120)}`);
    }

    const buffer = Buffer.from(await response.arrayBuffer());
    if (!buffer.length) {
      throw new Error("result fetch returned empty content");
    }

    const safeUserId = String(task.user_id).replace(/[^a-zA-Z0-9_-]/g, "");
    const extension = mediaExtensionFromContentType(contentType, resultUrl, task.type);
    const objectKey = `users/${safeUserId}/generation-results/${task.id}-${Date.now()}-${crypto.randomUUID()}.${extension}`;
    const uploaded = await putObject(objectKey, buffer, contentType, { "x-oss-object-acl": "public-read" });
    return uploaded.url;
  } catch (error) {
    console.warn(`[aiTaskWorker] result persistence skipped for task ${task?.id}:`, error.message);
    return null;
  }
}

function normalizeImageResultValue(value) {
  const stringValue = firstString(value);
  if (!stringValue) return undefined;
  if (/^(https?:)?\/\//i.test(stringValue) || /^data:image\//i.test(stringValue)) return stringValue;
  const markdownImageMatch = stringValue.match(/!\[[^\]]*]\((https?:\/\/[^)\s]+)\)/i);
  if (markdownImageMatch?.[1]) return markdownImageMatch[1];
  const urlMatch = stringValue.match(/https?:\/\/[^\s"'<>）)]+/i);
  if (urlMatch?.[0]) return urlMatch[0];
  try {
    const parsed = JSON.parse(stringValue);
    const parsedResult = firstImageResult(parsed);
    if (parsedResult) return parsedResult;
  } catch {}
  if (/^[A-Za-z0-9+/]+={0,2}$/.test(stringValue) && stringValue.length > 128) {
    return `data:image/png;base64,${stringValue}`;
  }
  return undefined;
}

function firstImageResult(...values) {
  for (const value of values) {
    if (Array.isArray(value)) {
      for (const item of value) {
        const result = firstImageResult(item);
        if (result) return result;
      }
      continue;
    }
    if (value && typeof value === "object") {
      const result = firstImageResult(
        value.url,
        value.image_url,
        value.imageUrl,
        value.result_url,
        value.resultUrl,
        value.output_url,
        value.outputUrl,
        value.b64_image,
        value.b64_json,
        value.base64,
        value.image,
        value.content,
        value.text,
        value.message,
      );
      if (result) return result;
      continue;
    }
    const result = normalizeImageResultValue(value);
    if (result) return result;
  }
  return undefined;
}

function pickDeep(value, keys) {
  if (!value || typeof value !== "object") return undefined;
  const obj = value;
  for (const key of keys) {
    const direct = firstString(obj[key]);
    if (direct) return direct;
  }
  for (const child of Object.values(obj)) {
    if (child && typeof child === "object") {
      const nested = pickDeep(child, keys);
      if (nested) return nested;
    }
  }
  return undefined;
}

function normalizeStatus(value) {
  return String(value || "").trim().toLowerCase();
}

function isCompletedStatus(status) {
  return ["completed", "complete", "succeeded", "success", "succeed", "done", "finished", "successed"].includes(status);
}

function isFailedStatus(status) {
  return ["failed", "failure", "fail", "canceled", "cancelled", "expired", "error", "violation"].includes(status);
}

function extractProviderTaskId(json) {
  const data = asObject(json?.data) || json;
  const output = asObject(json?.output) || asObject(data?.output);
  return firstString(
    output?.task_id,
    output?.taskId,
    data?.task_id,
    data?.taskId,
    data?.id,
    json?.task_id,
    json?.taskId,
    json?.id,
  );
}

function extractImageUrl(json) {
  const rawData = json?.data;
  const data = asObject(rawData) || json;
  const rawOutput = json?.output ?? data?.output;
  const output = asObject(rawOutput) || asObject(data?.output);
  const rawResult = data?.result ?? json?.result ?? output?.result;
  const result = asObject(rawResult);
  const choices =
    (Array.isArray(output?.choices) && output.choices) ||
    (Array.isArray(data?.choices) && data.choices) ||
    (Array.isArray(json?.choices) && json.choices) ||
    [];
  const firstChoice = asObject(choices[0]);
  const message = asObject(firstChoice?.message);
  const content = Array.isArray(message?.content) ? message.content : [];
  const firstContent = asObject(content[0]);
  const outputResults = Array.isArray(output?.results) ? output.results : [];
  const topLevelResults = Array.isArray(json?.results) ? json.results : [];
  const dataResults = Array.isArray(data?.results) ? data.results : [];
  const resultResults = Array.isArray(result?.results) ? result.results : [];
  const dataImages = Array.isArray(data?.images) ? data.images : [];
  const dataImageUrls = Array.isArray(data?.image_urls) ? data.image_urls : [];
  const dataUrls = Array.isArray(data?.urls) ? data.urls : [];
  const outputImages = Array.isArray(output?.images) ? output.images : [];
  const outputImageUrls = Array.isArray(output?.image_urls) ? output.image_urls : [];
  const resultImages = Array.isArray(result?.images) ? result.images : [];
  const resultUrls = Array.isArray(result?.urls) ? result.urls : [];
  const candidates = [
    ...topLevelResults,
    ...dataResults,
    ...outputResults,
    ...resultResults,
    ...dataImages,
    ...dataImageUrls,
    ...dataUrls,
    ...outputImages,
    ...outputImageUrls,
    ...resultImages,
    ...resultUrls,
  ];

  return firstImageResult(
    rawData,
    rawOutput,
    rawResult,
    firstContent?.image,
    firstContent?.image_url,
    firstContent?.image_url?.url,
    message?.content,
    firstChoice?.delta?.content,
    candidates,
    data?.image_url,
    data?.imageUrl,
    data?.result_url,
    data?.resultUrl,
    data?.output_url,
    data?.outputUrl,
    output?.image_url,
    output?.imageUrl,
    output?.result_url,
    output?.resultUrl,
    output?.output_url,
    output?.outputUrl,
    result?.image_url,
    result?.imageUrl,
    result?.result_url,
    result?.resultUrl,
    result?.output_url,
    result?.outputUrl,
    pickDeep(json, ["image", "image_url", "imageUrl", "result_url", "resultUrl", "output_url", "outputUrl", "url", "b64_image", "b64_json", "base64"]),
  );
}

function extractGeminiImageUrl(json) {
  // Gemini response: candidates[].content.parts[].inlineData (base64) or text (URL)
  const candidates = Array.isArray(json?.candidates) ? json.candidates : [];
  for (const candidate of candidates) {
    const parts = Array.isArray(candidate?.content?.parts) ? candidate.content.parts : [];
    for (const part of parts) {
      const inlineData = part?.inlineData;
      if (inlineData?.data) {
        const mimeType = inlineData.mimeType || "image/png";
        return `data:${mimeType};base64,${inlineData.data}`;
      }
    }
  }
  // Also check for direct URL in candidate text
  for (const candidate of candidates) {
    const parts = Array.isArray(candidate?.content?.parts) ? candidate.content.parts : [];
    for (const part of parts) {
      if (part?.text && /^https?:\/\/.+\.(png|jpg|jpeg|webp|gif)/i.test(part.text)) {
        return part.text;
      }
    }
  }
  return null;
}

function extractVideoUrl(json) {
  const data = asObject(json?.data) || json;
  const output = asObject(json?.output) || asObject(data?.output);
  const result = asObject(data?.task_result) || asObject(json?.task_result) || asObject(output?.task_result);
  const videos = Array.isArray(result?.videos) ? result.videos : [];
  const firstVideo = asObject(videos[0]);

  return firstString(
    output?.video_url,
    output?.output_video_url,
    output?.outputVideoUrl,
    output?.watermark_video_url,
    data?.video_url,
    data?.videoUrl,
    data?.output_video_url,
    data?.outputVideoUrl,
    data?.url,
    data?.result?.url,
    firstVideo?.url,
    firstVideo?.video_url,
    pickDeep(json, ["video_url", "output_video_url", "outputVideoUrl", "watermark_video_url", "videoUrl", "download_url", "downloadUrl", "content_url", "contentUrl", "url"]),
  );
}

function extractErrorMessage(json, fallback) {
  const data = asObject(json?.data) || {};
  const output = asObject(json?.output) || {};
  const error = asObject(json?.error) || {};
  return firstString(
    output.message,
    output.code,
    data.task_status_msg,
    data.failure_reason,
    data.message,
    error.message,
    error.error,
    json?.message,
    json?.error,
  ) || fallback;
}

async function fetchJson(url, headers) {
  const res = await fetch(url, { method: "GET", headers });
  if (!res.ok) return { ok: false, json: null };
  return { ok: true, json: await res.json() };
}

async function pollGrsaiImage(_taskId, providerTaskId, apiKey, baseUrl, resultEndpoint) {
  const url = `${baseUrl}${resultEndpoint}?id=${encodeURIComponent(providerTaskId)}`;
  const { ok, json } = await fetchJson(url, {
    Authorization: `Bearer ${apiKey}`,
    Accept: "application/json",
  });
  if (!ok) {
    console.warn(`[grsai-poll] task ${_taskId} fetch not ok, url=${url}`);
    return { status: "running", progress: 50 };
  }

  const data = asObject(json?.data) || json;
  const output = asObject(data?.output) || asObject(json?.output) || {};
  const status = normalizeStatus(
    output.task_status ||
      output.status ||
      data.task_status ||
      data.status ||
      data.state ||
      json?.task_status ||
      json?.status ||
      json?.state,
  );
  const resultUrl = extractImageUrl(json);
  console.info(`[grsai-poll] task ${_taskId} status=${status} resultUrl=${resultUrl ? "yes" : "no"} raw=${JSON.stringify(json).slice(0, 300)}`);
  if (resultUrl) {
    return { status: "completed", progress: 100, resultUrl };
  }
  if (isCompletedStatus(status)) {
    const completedUrl = extractImageUrl(json);
    if (!completedUrl) return { status: "failed", error: "Image generation completed without a result url" };
    return { status: "completed", progress: 100, resultUrl: completedUrl };
  }
  if (isFailedStatus(status)) {
    return { status: "failed", error: extractErrorMessage(json, "Image generation failed") };
  }
  return { status: "running", progress: Math.min(90, 30 + Math.random() * 40) };
}

async function pollDashscopeImage(_taskId, providerTaskId, apiKey) {
  const url = `https://dashscope.aliyuncs.com/api/v1/tasks/${encodeURIComponent(providerTaskId)}`;
  const { ok, json } = await fetchJson(url, {
    Authorization: `Bearer ${apiKey}`,
    Accept: "application/json",
  });
  if (!ok) return { status: "running", progress: 50 };

  const output = asObject(json?.output) || {};
  const status = normalizeStatus(output.task_status || json?.task_status);
  const resultUrl = extractImageUrl(json);
  if (isCompletedStatus(status)) {
    if (!resultUrl) return { status: "failed", error: "DashScope image generation completed without a result url" };
    return { status: "completed", progress: 100, resultUrl };
  }
  if (isFailedStatus(status)) {
    return { status: "failed", error: extractErrorMessage(json, "DashScope image generation failed") };
  }
  return { status: "running", progress: Math.min(90, 30 + Math.random() * 40) };
}

function base64Url(input) {
  return Buffer.from(input)
    .toString("base64")
    .replace(/\+/g, "-")
    .replace(/\//g, "_")
    .replace(/=+$/, "");
}

function parseKlingCredential(apiKey) {
  const raw = String(apiKey || "");
  const colonIndex = raw.indexOf(":");
  if (colonIndex <= 0) return null;
  const accessKey = raw.slice(0, colonIndex).trim();
  const secretKey = raw.slice(colonIndex + 1).trim();
  return accessKey && secretKey ? { accessKey, secretKey } : null;
}

function createKlingJwt(accessKey, secretKey) {
  const header = { alg: "HS256", typ: "JWT" };
  const now = Math.floor(Date.now() / 1000);
  const payload = { iss: accessKey, exp: now + 1800, nbf: now - 5 };
  const unsigned = `${base64Url(JSON.stringify(header))}.${base64Url(JSON.stringify(payload))}`;
  const signature = crypto.createHmac("sha256", secretKey).update(unsigned).digest("base64")
    .replace(/\+/g, "-")
    .replace(/\//g, "_")
    .replace(/=+$/, "");
  return `${unsigned}.${signature}`;
}

function getPollRequest(providerTaskId, apiKey, providerConfig) {
  const protocol = providerConfig.protocol;
  const baseUrl = providerConfig.baseUrl || "";

  if (
    protocol === "wan-i2v" ||
    protocol === "wan-s2v" ||
    protocol === "kling-dashscope" ||
    String(protocol || "").startsWith("happyhorse-")
  ) {
    return {
      url: `${baseUrl}/api/v1/tasks/${encodeURIComponent(providerTaskId)}`,
      headers: { Authorization: `Bearer ${apiKey}`, Accept: "application/json" },
    };
  }

  if (protocol === "kling-omni") {
    const credential = parseKlingCredential(apiKey);
    return {
      url: `${baseUrl}/v1/videos/omni-video/${encodeURIComponent(providerTaskId)}`,
      headers: {
        Authorization: `Bearer ${credential ? createKlingJwt(credential.accessKey, credential.secretKey) : apiKey}`,
        Accept: "application/json",
      },
    };
  }

  return {
    url: `${baseUrl}${providerConfig.endpoint}/${encodeURIComponent(providerTaskId)}`,
    headers: { Authorization: `Bearer ${apiKey}`, Accept: "application/json" },
  };
}

async function pollVideoTask(_taskId, providerTaskId, apiKey, providerConfig) {
  const { url, headers } = getPollRequest(providerTaskId, apiKey, providerConfig);
  const { ok, json } = await fetchJson(url, headers);
  if (!ok) return { status: "running", progress: 50 };

  const data = asObject(json?.data) || json;
  const output = asObject(json?.output) || {};
  const status = normalizeStatus(
    output.task_status ||
    data.task_status ||
    data.status ||
    json?.task_status ||
    json?.status,
  );

  const resultUrl = extractVideoUrl(json);
  if (isCompletedStatus(status) || resultUrl) {
    return { status: "completed", progress: 100, resultUrl: resultUrl || null };
  }
  if (isFailedStatus(status)) {
    return { status: "failed", error: extractErrorMessage(json, "Video generation failed") };
  }
  const progress = Number(data.progress || output.progress);
  return { status: "running", progress: Number.isFinite(progress) ? Math.min(95, progress) : Math.min(90, 30 + Math.random() * 30) };
}

function getMaxPollAttempts(type, providerConfig) {
  if (type === "image" && providerConfig?.transport === "grsai-image") {
    return Number.isFinite(GRS_IMAGE_MAX_POLL_ATTEMPTS) && GRS_IMAGE_MAX_POLL_ATTEMPTS > 0
      ? Math.trunc(GRS_IMAGE_MAX_POLL_ATTEMPTS)
      : 40;
  }
  if (type === "video") return 400;
  return MAX_POLL_ATTEMPTS;
}

function startPolling(taskDbId, { providerTaskId, apiKey, type, providerConfig, leaseToken, keyManager, onTaskFailed }) {
  if (activePollers.has(taskDbId)) return;

  let attempts = 0;
  const maxPollAttempts = getMaxPollAttempts(type, providerConfig);
  const interval = setInterval(async () => {
    attempts++;
    if (attempts > maxPollAttempts) {
      clearInterval(interval);
      activePollers.delete(taskDbId);
      if (leaseToken && keyManager) await keyManager.releaseKey(leaseToken).catch(() => {});
      if (typeof onTaskFailed === "function") {
        const handled = await onTaskFailed("Task timed out").catch((fallbackErr) => {
          console.error(`[aiTaskWorker] fallback error for task ${taskDbId}:`, fallbackErr.message);
          return false;
        });
        if (handled) return;
      }
      await updateTaskInDb(taskDbId, { status: "failed", error: "Task timed out" });
      return;
    }

    try {
      // Check if task was cancelled by user
      const { rows: [taskRow] } = await pool.query("SELECT status FROM generation_tasks WHERE id = $1", [taskDbId]);
      if (!taskRow || taskRow.status === "cancelled") {
        clearInterval(interval);
        activePollers.delete(taskDbId);
        if (leaseToken && keyManager) await keyManager.releaseKey(leaseToken).catch(() => {});
        return;
      }

      let result;
      if (type === "image") {
        if (providerConfig.transport === "dashscope-image") {
          result = await pollDashscopeImage(taskDbId, providerTaskId, apiKey);
        } else {
          result = await pollGrsaiImage(taskDbId, providerTaskId, apiKey, providerConfig.baseUrl, providerConfig.resultEndpoint || "/result");
        }
      } else {
        result = await pollVideoTask(taskDbId, providerTaskId, apiKey, providerConfig);
      }

      if (result.status === "completed" || result.status === "failed") {
        clearInterval(interval);
        activePollers.delete(taskDbId);
        if (leaseToken && keyManager) await keyManager.releaseKey(leaseToken).catch(() => {});
        if (result.status === "failed" && typeof onTaskFailed === "function") {
          const handled = await onTaskFailed(result.error || "Task failed").catch((fallbackErr) => {
            console.error(`[aiTaskWorker] fallback error for task ${taskDbId}:`, fallbackErr.message);
            return false;
          });
          if (handled) return;
        }
      }

      await updateTaskInDb(taskDbId, result);
    } catch (err) {
      console.error(`[aiTaskWorker] poll error for task ${taskDbId}:`, err.message);
    }
  }, POLL_INTERVAL_MS);

  activePollers.set(taskDbId, { interval, leaseToken });
}

function stopPolling(taskDbId) {
  const poller = activePollers.get(taskDbId);
  if (poller) {
    clearInterval(poller.interval);
    activePollers.delete(taskDbId);
  }
}

function getActiveCount() {
  return activePollers.size;
}

// --- Periodic stale task cleanup ---
// Runs every 5 minutes, marks tasks stuck in 'pending'/'running' for too long as 'failed'.
// This catches cases where the worker crashed, the provider API never responded,
// or the cancel request failed silently on the client side.
const STALE_TASK_CLEANUP_INTERVAL_MS = 5 * 60 * 1000;
let staleTaskCleanupTimer = null;

async function runStaleTaskCleanup() {
  try {
    const { rows } = await pool.query(
      `UPDATE generation_tasks
       SET status = 'failed', error = '任务超时自动释放', updated_at = NOW()
       WHERE status IN ('pending', 'running')
         AND GREATEST(updated_at, COALESCE(last_poll_at, created_at)) < NOW() - INTERVAL '10 minutes'
       RETURNING id`,
    );
    for (const row of rows) {
      taskEvents.emit(`task:${row.id}`, {
        taskId: row.id,
        status: "failed",
        progress: null,
        resultUrl: null,
        error: "任务超时自动释放",
      });
      // Also stop any active poller for this task
      const poller = activePollers.get(row.id);
      if (poller) {
        clearInterval(poller.timer);
        activePollers.delete(row.id);
      }
    }
    if (rows.length > 0) {
      console.log(`[aiTaskWorker] Cleaned up ${rows.length} stale task(s)`);
    }
  } catch (err) {
    console.error("[aiTaskWorker] Stale task cleanup failed:", err.message);
  }
}

function startStaleTaskCleanup() {
  if (staleTaskCleanupTimer) return;
  staleTaskCleanupTimer = setInterval(runStaleTaskCleanup, STALE_TASK_CLEANUP_INTERVAL_MS);
  // Run once shortly after startup
  setTimeout(runStaleTaskCleanup, 10_000);
}

function stopStaleTaskCleanup() {
  if (staleTaskCleanupTimer) {
    clearInterval(staleTaskCleanupTimer);
    staleTaskCleanupTimer = null;
  }
}

module.exports = {
  startPolling,
  stopPolling,
  updateTaskInDb,
  getActiveCount,
  extractProviderTaskId,
  extractImageUrl,
  extractGeminiImageUrl,
  extractVideoUrl,
  parseKlingCredential,
  createKlingJwt,
  taskEvents,
  startStaleTaskCleanup,
  stopStaleTaskCleanup,
};